Patch speech_transformer data loading to load from absolute path (#699)

Summary:
Lazy Tensor users would like to run the benchmark from a different directory. Yet the data loading script in speech transformer model always assume the code runs under the `benchmark/` directory.

This PR patches the dataloading code of speech transformer model such that it finds the absolute path of input data, and load data from there.

It also fixes a bug in `run.py` related to `extra_args`.

Pull Request resolved: #699

Reviewed By: wconstab

Differential Revision: D33625399

Pulled By: xuzhao9

fbshipit-source-id: 2e9fdc9184bb397601dbb287dd872c281707d5a9

Author: xuzhao9

run.py

@@ -163,7 +163,8 @@ def _validate_devices(devices: str):
         print(f"Unable to find model matching {args.model}.")
         exit(-1)
     model_args = inspect.signature(Model)
-    if extra_args and not 'extra_args' in model_args.parameters:
+    support_extra_args = 'extra_args' in model_args.parameters
+    if extra_args and not support_extra_args:
         print(f"The model {args.model} doesn't accept extra args: {extra_args}")
         exit(-1)
     print(f"Running {args.test} method from {Model.name} on {args.device} in {args.mode} mode.")
@@ -174,14 +175,23 @@ def _validate_devices(devices: str):
     if args.bs:
         try:
             if args.test == "eval":
-                m = Model(device=args.device, jit=(args.mode == "jit"), eval_bs=args.bs, extra_args=extra_args)
+                if support_extra_args:
+                    m = Model(device=args.device, jit=(args.mode == "jit"), eval_bs=args.bs, extra_args=extra_args)
+                else:
+                    m = Model(device=args.device, jit=(args.mode == "jit"), eval_bs=args.bs)
             elif args.test == "train":
-                m = Model(device=args.device, jit=(args.mode == "jit"), train_bs=args.bs, extra_args=extra_args)
+                if support_extra_args:
+                    m = Model(device=args.device, jit=(args.mode == "jit"), train_bs=args.bs, extra_args=extra_args)
+                else:
+                    m = Model(device=args.device, jit=(args.mode == "jit"), eval_bs=args.bs)
         except:
             print(f"The model {args.model} doesn't support specifying batch size, please remove --bs argument in the commandline.")
             exit(1)
     else:
-        m = Model(device=args.device, jit=(args.mode == "jit"), extra_args=extra_args)
+        if support_extra_args:
+            m = Model(device=args.device, jit=(args.mode == "jit"), extra_args=extra_args)
+        else:
+            m = Model(device=args.device, jit=(args.mode == "jit"))
 
     test = getattr(m, args.test)
     model_flops = None

torchbenchmark/models/speech_transformer/__init__.py

@@ -30,8 +30,8 @@ def __init__(self, device=None, jit=False, train_bs=32):
             return
         self.traincfg = SpeechTransformerTrainConfig(prefetch=True, train_bs=train_bs, num_train_batch=NUM_TRAIN_BATCH)
         self.evalcfg = SpeechTransformerEvalConfig(self.traincfg, num_eval_batch=NUM_EVAL_BATCH)
-        self.traincfg.model.cuda()
-        self.evalcfg.model.cuda()
+        self.traincfg.model.to(self.device)
+        self.evalcfg.model.to(self.device)
 
     def get_module(self):
         if self.device == "cpu":
@@ -40,7 +40,7 @@ def get_module(self):
             raise NotImplementedError("JIT is not supported by this model")
         for data in self.traincfg.tr_loader:
             padded_input, input_lengths, padded_target = data
-            return self.traincfg.model, (padded_input.cuda(), input_lengths.cuda(), padded_target.cuda())
+            return self.traincfg.model, (padded_input.to(self.device), input_lengths.to(self.device), padded_target.to(self.device))
 
     def train(self, niter=1):
         if self.device == "cpu":

torchbenchmark/models/speech_transformer/config.py

@@ -37,7 +37,11 @@ class SpeechTransformerTrainConfig:
     batch_frames = 15000
     maxlen_in = 800
     maxlen_out = 150
-    num_workers = 4
+    # don't use subprocess in dataloader
+    # because TorchBench is only running 1 batch
+    num_workers = 0
+    # original value
+    # num_workers = 4
     # optimizer
     k = 0.2
     warmup_steps = 1

torchbenchmark/models/speech_transformer/speech_transformer/data/data.py

@@ -10,6 +10,7 @@
 """
 import json
 
+from pathlib import Path
 import numpy as np
 import torch
 import torch.utils.data as data
@@ -144,7 +145,10 @@ def load_inputs_and_targets(batch, LFR_m=1, LFR_n=1):
     # load acoustic features and target sequence of token ids
     # for b in batch:
     #     print(b[1]['input'][0]['feat'])
-    xs = [kaldi_io.read_mat(b[1]['input'][0]['feat']) for b in batch]
+    # TorchBench: Patch the input data with current file directory
+    # Current file path: TORCHBENCH_ROOT/torchbenchmark/models/speech_transformer/speech_transformer/data/data.py
+    TORCHBENCH_ROOT = Path(__file__).parents[5]
+    xs = [kaldi_io.read_mat(str(TORCHBENCH_ROOT.joinpath(b[1]['input'][0]['feat']).resolve())) for b in batch]
     ys = [b[1]['output'][0]['tokenid'].split() for b in batch]
 
     if LFR_m != 1 or LFR_n != 1: