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

Create DataParallel model if several GPUs
DRL36 · Nov 3, 2018 · f6ed6ac · f6ed6ac
2 parents 6165f84 + 20294bd
commit f6ed6ac
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 0 deletions.
diff --git a/extract_features_pytorch.py b/extract_features_pytorch.py
@@ -249,6 +249,9 @@ def main():
     if args.init_checkpoint is not None:
         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)

diff --git a/run_classifier_pytorch.py b/run_classifier_pytorch.py
@@ -482,6 +482,9 @@ def main():
     if args.init_checkpoint is not None:
         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.}

diff --git a/run_squad_pytorch.py b/run_squad_pytorch.py
@@ -795,6 +795,9 @@ def main():
     if args.init_checkpoint is not None:
         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.}