Prepare workflow for multiple gpus

src/data/dataset.py

@@ -39,7 +39,7 @@ def cl_fn(batch, tokenizer):
 
  return img_emb, input_ids, attention_mask
 
-def get_loader(dataset, bs_exp=5, shuffle=True, num_workers=0, pin_memory=False):
+def get_loader(dataset, bs_exp=5, shuffle=True, num_workers=0, pin_memory=False, sampler=None):
  tokenizer = GPT2Tokenizer.from_pretrained('gpt2-xl')
  tokenizer.pad_token = tokenizer.eos_token
 
@@ -49,5 +49,6 @@ def get_loader(dataset, bs_exp=5, shuffle=True, num_workers=0, pin_memory=False)
  collate_fn=lambda b: cl_fn(b, tokenizer),
  shuffle=shuffle,
  num_workers=num_workers,
- pin_memory=pin_memory
+ pin_memory=pin_memory,
+ sampler=sampler
  )

src/model/loops.py

@@ -13,7 +13,7 @@
 from tqdm import tqdm
 
 class Trainer:
- def __init__(self, model, optimizer, scaler, scheduler, train_loader, valid_loader, test_dataset, test_path, ckp_path, device):
+ def __init__(self, model, optimizer, scaler, scheduler, train_loader, valid_loader, test_dataset, test_path, ckp_path, device, multi_gpu=False):
  self.model = model
  self.optimizer = optimizer
  self.scaler = scaler
@@ -25,6 +25,8 @@ def __init__(self, model, optimizer, scaler, scheduler, train_loader, valid_load
  self.ckp_path = ckp_path
  self.device = device
 
+ self.multi_gpu = multi_gpu
+
  # load checkpoint
  if os.path.isfile(ckp_path):
  self._load_ckp(ckp_path, optimizer, scheduler, scaler, device=device)
@@ -141,7 +143,7 @@ def save_ckp(self, ckp_path):
  torch.save(
  {
  'epoch': self.epoch,
- 'model_state_dict': self.model.state_dict(),
+ 'model_state_dict': self.model.module.state_dict() if self.multi_gpu else self.model.state_dict(),
  'optimizer_state_dict': self.optimizer.state_dict(),
  'scheduler_state_dict': self.scheduler.state_dict(),
  'scaler_state_dict': self.scaler.state_dict(),
@@ -153,6 +155,12 @@ def save_ckp(self, ckp_path):
 
  return True
 
+ def set_samplers_epoch(self, epoch):
+ self.train_loader.sampler.set_epoch(epoch)
+ self.valid_loader.sampler.set_epoch(epoch)
+
+ return True
+
  def _load_ckp(
  self, 
  checkpoint_fpath,

src/model/model.py

@@ -1,17 +1,29 @@
 '''
  Module contains final Model and all pieces of it.
 '''
+import os
 
 import torch
 import torch.nn as nn
-# from torch.distributed import init_process_group, destroy_process_group
+from torch.distributed import init_process_group, destroy_process_group
 from transformers import CLIPModel, CLIPProcessor, GPT2LMHeadModel, GPT2Tokenizer
 
-# def ddp_setup(rank, world_size):
-# init_process_group('nccl', rank=rank, world_size=world_size)
+def ddp_setup(rank, world_size):
+ '''
+ Setup distributed training.
+ '''
+
+ os.environ['MASTER_ADDR'] = 'localhost'
+ os.environ['MASTER_PORT'] = '12355'
+
+ init_process_group('nccl', rank=rank, world_size=world_size)
+
+def ddp_cleanup():
+ '''
+ Cleanup distributed training.
+ '''
 
-# def ddp_cleanup():
-# destroy_process_group()
+ destroy_process_group()
 
 class ImageEncoder(nn.Module):
  '''

src/training.py

@@ -7,19 +7,16 @@
 import random
 
 import numpy as np
-
-import wandb
 import torch
+import torch.multiprocessing as mp
+from torch.nn.parallel import DistributedDataParallel as DDP
 import torch.optim as optim
 from torch.utils.data import random_split
-
-import torch.multiprocessing as mp
 from torch.utils.data.distributed import DistributedSampler
-from torch.nn.parallel import DistributedDataParallel as DDP
 
+import wandb
 from data import MiniFlickrDataset, get_loader
-# from model import ddp_cleanup, ddp_setup, Net, Trainer
-from model import Net, Trainer
+from model import ddp_cleanup, ddp_setup, Net, Trainer
 from utils import Config, LRWarmup
 
 config = Config()
@@ -42,9 +39,18 @@
 torch.cuda.manual_seed(config.seed)
 torch.backends.cudnn.deterministic = True
 
-def main(config, ckp_name):
+def main(rank, world_size, config, ckp_name=''):
+ # more than 1 GPU
  is_cuda = torch.cuda.is_available()
- device = 'cuda' if is_cuda else 'cpu'
+ MULTIGPU = world_size > 1
+
+ if MULTIGPU:
+ ddp_setup(rank, world_size)
+ device = rank
+
+ else:
+ device = torch.device('cuda' if is_cuda else 'cpu')
+
  model = Net(
  ep_len=config.ep_len,
  num_layers=config.num_layers, 
@@ -55,6 +61,9 @@ def main(config, ckp_name):
  device=device
  )
 
+ if MULTIGPU:
+ model = DDP(model, device_ids=[device])
+
  dataset = MiniFlickrDataset(os.path.join('data', 'processed', 'dataset.pkl'))
 
  config.train_size = int(config.train_size * len(dataset))
@@ -66,17 +75,19 @@ def main(config, ckp_name):
  train_loader = get_loader(
  train_dataset, 
  bs_exp=config.batch_size_exp, 
- shuffle=True, 
+ shuffle=not MULTIGPU, 
  num_workers=config.num_workers if is_cuda else 0,
- pin_memory=is_cuda
+ pin_memory=is_cuda,
+ sampler=DistributedSampler(train_dataset) if MULTIGPU else None
  )
 
  valid_loader = get_loader(
  val_dataset, 
  bs_exp=config.batch_size_exp, 
  shuffle=False, 
  num_workers=config.num_workers if is_cuda else 0,
- pin_memory=is_cuda
+ pin_memory=is_cuda,
+ sampler=DistributedSampler(val_dataset) if MULTIGPU else None
  )
 
  optimizer = optim.Adam(model.parameters(), lr=config.lr)
@@ -105,6 +116,9 @@ def main(config, ckp_name):
  wandb.init(project='clipXgpt2 captioner', config=config.__dict__)
  wandb.watch(trainer.model, log='all')
  for epoch in range(trainer.epoch, config.epochs):
+ if MULTIGPU:
+ trainer.set_samplers_epoch(epoch)
+
  trainer.train_epoch()
  trainer.valid_epoch()
  trainer.test_result()
@@ -122,8 +136,15 @@ def main(config, ckp_name):
  if not os.path.exists(config.weights_dir):
  os.makedirs(config.weights_dir)
 
- if (epoch + 1) % 10 == 0: 
- trainer.save_ckp(os.path.join(config.weights_dir, f'epoch_{epoch}.pt'))
+ if (epoch + 1) % 10 == 0 and rank == 0:
+ trainer.save_ckp(os.path.join(config.weights_dir, f'epoch_{epoch + 1}.pt'))
+
+ ddp_cleanup()
+
 
 if __name__ == '__main__':
- main(config, args.checkpoint_name)
+ # check if there is no GPU - use CPU -> world_size = 1
+
+ world_size = torch.cuda.device_count() if torch.cuda.is_available() else 1
+
+ mp.spawn(main, args=(world_size, config, ''), nprocs=world_size)