Merge pull request #1 from huggingface/multi-gpu-support

Create DataParallel model if several GPUs

extract_features_pytorch.py

@@ -250,6 +250,9 @@ def main():
         model.load_state_dict(torch.load(args.init_checkpoint, map_location='cpu'))
     model.to(device)
     
+    if n_gpu > 1:
+        model = nn.DataParallel(model)
+
     all_input_ids = torch.tensor([f.input_ids for f in features], dtype=torch.long)
     all_input_mask = torch.tensor([f.input_mask for f in features], dtype=torch.long)
     all_example_index = torch.arange(all_input_ids.size(0), dtype=torch.long)

run_classifier_pytorch.py

@@ -483,6 +483,9 @@ def main():
         model.bert.load_state_dict(torch.load(args.init_checkpoint, map_location='cpu'))
     model.to(device)
     
+    if n_gpu > 1:
+        model = torch.nn.DataParallel(model)
+
     optimizer = BERTAdam([{'params': [p for n, p in model.named_parameters() if n != 'bias'], 'l2': 0.01},
                           {'params': [p for n, p in model.named_parameters() if n == 'bias'], 'l2': 0.}
                          ],

run_squad_pytorch.py

@@ -796,6 +796,9 @@ def main():
         model.bert.load_state_dict(torch.load(args.init_checkpoint, map_location='cpu'))
     model.to(device)
     
+    if n_gpu > 1:
+        model = torch.nn.DataParallel(model)
+
     optimizer = BERTAdam([{'params': [p for n, p in model.named_parameters() if n != 'bias'], 'l2': 0.01},
                           {'params': [p for n, p in model.named_parameters() if n == 'bias'], 'l2': 0.}
                          ],