distribute GPU

Author: Wonters

src/mixins.py

@@ -8,6 +8,7 @@
 import optuna
 import random
 from torch.utils.data import Subset, DataLoader
+import torch.distributed as dist
 
 logger = logging.getLogger(__name__)
 
@@ -76,7 +77,7 @@ def _train_batch(self, x, y):
         acc = correct / len(labels)
         loss.backward()
         self.optimizer.step()
-        logger.info(f"loss {loss.item()}")
+        logger.info(f" Rank {dist.get_rank()} loss {loss.item()}")
         mlflow.log_metric("loss", loss.item())
         mlflow.log_metric("acc", acc)
         mlflow.log_metric("time", time.time())

src/ml.py

@@ -11,8 +11,10 @@
 import matplotlib.pyplot as plt
 from multiprocessing import cpu_count
 import torch
+import torch.nn as nn
+import torch.distributed as dist
 from tqdm import tqdm
-from torch.utils.data import DataLoader
+from torch.utils.data import DataLoader, DistributedSampler
 import torch.nn.functional as F
 from sklearn.linear_model import LogisticRegression
 import lightgbm as lgm
@@ -473,16 +475,14 @@ def load_checkpoint(self):
             }
             self.model.load_state_dict(embedding_weights, strict=False)
             self.model.eval()
-        import torch
-        import torch.nn as nn
-        import torch.distributed as dist
 
         dist.init_process_group("nccl")
         local_rank = torch.distributed.get_rank()
         torch.cuda.set_device(local_rank)
 
         self.model = self.model.cuda(local_rank)
         self.model = nn.parallel.DistributedDataParallel(self.model, device_ids=[local_rank], output_device=local_rank,find_unused_parameters=True)
+        self.dataset = DistributedSampler(self.dataset)
         self.optimizer = torch.optim.Adam(self.model.parameters(), lr=self.lr)
         self.scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
             self.optimizer, mode="min", factor=0.5, patience=2