Merge pull request #55 from undertherain/torch-parallel-issue-gh47

Close #47 (for now)
undertherain · Jun 18, 2020 · 21b87dc · 21b87dc
2 parents 1101988 + 4243771
commit 21b87dc
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 12 deletions.
diff --git a/benchmarker/benchmarker.py b/benchmarker/benchmarker.py
@@ -52,7 +52,8 @@ def run(args, unknown_args):
         params["batch_size_per_device"] = int(args.batch_size)
     # params["misc"] = args.misc
     if args.gpus:
-        params["gpus"] = list(map(int, args.gpus.split(',')))
+        os.environ["CUDA_VISIBLE_DEVICES"] = args.gpus
+        params["gpus"] = list(map(int, args.gpus.split(",")))
     else:
         params["gpus"] = []
 

diff --git a/benchmarker/modules/do_pytorch.py b/benchmarker/modules/do_pytorch.py
@@ -34,13 +34,13 @@ def train(self, model, device, optimizer, epoch):
             # TODO: criterion should be included in the model, deepening on params
             #criterion = nn.CrossEntropyLoss()
             #loss = criterion(output, target)
-            loss.backward()
+            loss.mean().backward()
             optimizer.step()
             log_interval = 10
             if batch_idx % log_interval == 0:
                 print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                     epoch, batch_idx * len(data), len(self.x_train),
-                    100. * batch_idx / len(self.x_train), loss.item()))
+                    100. * batch_idx / len(self.x_train), loss.mean().item()))
         if device.type == "cuda":
             torch.cuda.synchronize()
 
@@ -74,13 +74,7 @@ def run_internal(self):
                 torch.backends.mkldnn.enabled = False
             else:
                 raise RuntimeError("Unknown backend")
-        if self.params["nb_gpus"] > 1:
-            raise NotADirectoryError("multyple GPUs not supported yet")
-        if self.params["gpus"]:
-            torch.cuda.set_device(self.params["gpus"][0])
-            device = torch.device("cuda")
-        else:
-            device = torch.device("cpu")
+        device = torch.device("cuda" if self.params["gpus"] else "cpu")
 
         x_train, y_train = self.load_data()
         # TODO: make of/on-core optional
@@ -89,7 +83,10 @@ def run_internal(self):
         # train_dataset = torch.utils.data.TensorDataset(x_train, y_train)
         # train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=self.params["batch_size"], shuffle=False)
 
-        model = self.net.to(device)
+        model = self.net
+        if len(self.params["gpus"]) > 1:
+            model = nn.DataParallel(model)
+        model.to(device)
         # TODO: args for training hyperparameters
         optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.95)
         start = timer()

diff --git a/benchmarker/modules/do_tensorflow.py b/benchmarker/modules/do_tensorflow.py
@@ -15,7 +15,6 @@ class Benchmark(INeuralNet):
 
     def __init__(self, params, remaining_args=None):
         gpus = params["gpus"]
-        os.environ["CUDA_VISIBLE_DEVICES"] = ",".join(map(str, gpus)) if gpus else "-1"
         super().__init__(params, remaining_args)
         self.params["channels_first"] = False
         os.environ["KERAS_BACKEND"] = "tensorflow"