add multi-GPU dataparrallel

mace/tools/torch_tools.py

@@ -119,3 +119,17 @@ def init_wandb(project: str, entity: str, name: str, config: dict):
     import wandb
 
     wandb.init(project=project, entity=entity, name=name, config=config)
+
+class DataParallelModel(torch.nn.Module):
+    def __init__(self, model):
+        super(DataParallelModel, self).__init__()
+        self.model = torch.nn.DataParallel(model).cuda()
+
+    def forward(self, batch, training):
+        return self.model(batch, training=training)
+
+    def __getattr__(self, name):
+        try:
+            return super().__getattr__(name)
+        except AttributeError:
+            return getattr(self.model.module, name)

scripts/run_train.py

@@ -318,6 +318,10 @@ def main() -> None:
     else:
         raise RuntimeError(f"Unknown model: '{args.model}'")
 
+
+    if torch.cuda.device_count() > 1:
+        logging.info(f"Multi-GPUs training on {torch.cuda.device_count()} GPUs.")
+        model = tools.DataParallelModel(model)
     model.to(device)
 
     # Optimizer