| 24 | from torch.amp import autocast |
| 25 | |
| 26 | class Executor: |
| 27 | def __init__(self): |
| 28 | self.step = 0 |
| 29 | self.epoch = 0 |
| 30 | self.rank = int(os.environ.get('RANK', 0)) |
| 31 | if torch.cuda.is_available(): |
| 32 | if torch.cuda.is_available(): |
| 33 | self.device = torch.device('cuda:{}'.format(self.rank)) |
| 34 | elif torch.backends.mps.is_available(): |
| 35 | self.device = torch.device('mps') |
| 36 | elif torch.xpu.is_available(): |
| 37 | self.device = torch.device('xpu') |
| 38 | else: |
| 39 | self.device = torch.device('cpu') |
| 40 | def train_one_epoch(self, model, optimizer, scheduler, train_data_loader, cv_data_loader, writer, info_dict, group_join, scaler=None): |
| 41 | ''' Train one epoch |
| 42 | ''' |
| 43 | |
| 44 | lr = optimizer.param_groups[0]['lr'] |
| 45 | logging.info('Epoch {} TRAIN info lr {} rank {}'.format(self.epoch, lr, self.rank)) |
| 46 | logging.info('using accumulate grad, new batch size is {} times' |
| 47 | ' larger than before'.format(info_dict['accum_grad'])) |
| 48 | # A context manager to be used in conjunction with an instance of |
| 49 | # torch.nn.parallel.DistributedDataParallel to be able to train |
| 50 | # with uneven inputs across participating processes. |
| 51 | model.train() |
| 52 | model_context = model.join if info_dict['train_engine'] == 'torch_ddp' else nullcontext |
| 53 | with model_context(): |
| 54 | for batch_idx, batch_dict in enumerate(train_data_loader): |
| 55 | info_dict["tag"] = "TRAIN" |
| 56 | info_dict["step"] = self.step |
| 57 | info_dict["epoch"] = self.epoch |
| 58 | info_dict["batch_idx"] = batch_idx |
| 59 | if inspiremusic_join(group_join, info_dict): |
| 60 | break |
| 61 | |
| 62 | # Disable gradient synchronizations across DDP processes. |
| 63 | # Within this context, gradients will be accumulated on module |
| 64 | # variables, which will later be synchronized. |
| 65 | if info_dict['train_engine'] == 'torch_ddp' and (batch_idx + 1) % info_dict["accum_grad"] != 0: |
| 66 | context = model.no_sync |
| 67 | # Used for single gpu training and DDP gradient synchronization |
| 68 | # processes. |
| 69 | else: |
| 70 | context = nullcontext |
| 71 | |
| 72 | with context(): |
| 73 | with autocast(device_type='cuda', enabled=scaler is not None): |
| 74 | info_dict = batch_forward(model, batch_dict, info_dict, scaler) |
| 75 | info_dict = batch_backward(model, info_dict, scaler) |
| 76 | |
| 77 | info_dict = update_parameter_and_lr(model, optimizer, scheduler, info_dict, scaler) |
| 78 | log_per_step(writer, info_dict) |
| 79 | # NOTE specify save_per_step in inspiremusic.yaml if you want to enable step save |
| 80 | if info_dict['save_per_step'] > 0 and (self.step + 1) % info_dict['save_per_step'] == 0 and \ |
| 81 | (batch_idx + 1) % info_dict["accum_grad"] == 0: |
| 82 | dist.barrier() |
| 83 | self.cv(model, cv_data_loader, writer, info_dict, on_batch_end=False, scaler=scaler) |