| 128 | |
| 129 | |
| 130 | def train_one_epoch(config, model, data_loader, optimizer, epoch, lr_scheduler): |
| 131 | model.train() |
| 132 | optimizer.zero_grad() |
| 133 | |
| 134 | num_steps = len(data_loader) |
| 135 | batch_time = AverageMeter() |
| 136 | loss_meter = AverageMeter() |
| 137 | norm_meter = AverageMeter() |
| 138 | |
| 139 | start = time.time() |
| 140 | end = time.time() |
| 141 | for idx, (samples_1, samples_2, targets) in enumerate(data_loader): |
| 142 | samples_1 = samples_1.cuda(non_blocking=True) |
| 143 | samples_2 = samples_2.cuda(non_blocking=True) |
| 144 | targets = targets.cuda(non_blocking=True) |
| 145 | |
| 146 | loss = model(samples_1, samples_2) |
| 147 | |
| 148 | optimizer.zero_grad() |
| 149 | if config.AMP_OPT_LEVEL != "O0": |
| 150 | with amp.scale_loss(loss, optimizer) as scaled_loss: |
| 151 | scaled_loss.backward() |
| 152 | if config.TRAIN.CLIP_GRAD: |
| 153 | grad_norm = torch.nn.utils.clip_grad_norm_(amp.master_params(optimizer), config.TRAIN.CLIP_GRAD) |
| 154 | else: |
| 155 | grad_norm = get_grad_norm(amp.master_params(optimizer)) |
| 156 | else: |
| 157 | loss.backward() |
| 158 | if config.TRAIN.CLIP_GRAD: |
| 159 | grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), config.TRAIN.CLIP_GRAD) |
| 160 | else: |
| 161 | grad_norm = get_grad_norm(model.parameters()) |
| 162 | optimizer.step() |
| 163 | lr_scheduler.step_update(epoch * num_steps + idx) |
| 164 | |
| 165 | torch.cuda.synchronize() |
| 166 | |
| 167 | loss_meter.update(loss.item(), targets.size(0)) |
| 168 | norm_meter.update(grad_norm) |
| 169 | batch_time.update(time.time() - end) |
| 170 | end = time.time() |
| 171 | |
| 172 | if idx % config.PRINT_FREQ == 0: |
| 173 | lr = optimizer.param_groups[0]['lr'] |
| 174 | memory_used = torch.cuda.max_memory_allocated() / (1024.0 * 1024.0) |
| 175 | etas = batch_time.avg * (num_steps - idx) |
| 176 | logger.info( |
| 177 | f'Train: [{epoch}/{config.TRAIN.EPOCHS}][{idx}/{num_steps}]\t' |
| 178 | f'eta {datetime.timedelta(seconds=int(etas))} lr {lr:.6f}\t' |
| 179 | f'time {batch_time.val:.4f} ({batch_time.avg:.4f})\t' |
| 180 | f'loss {loss_meter.val:.4f} ({loss_meter.avg:.4f})\t' |
| 181 | f'grad_norm {norm_meter.val:.4f} ({norm_meter.avg:.4f})\t' |
| 182 | f'mem {memory_used:.0f}MB') |
| 183 | epoch_time = time.time() - start |
| 184 | logger.info(f"EPOCH {epoch} training takes {datetime.timedelta(seconds=int(epoch_time))}") |
| 185 | |
| 186 | |
| 187 | if __name__ == '__main__': |