(train_loader, model, optimizer, scheduler, scaler, epoch, args)
| 14 | |
| 15 | |
| 16 | def train(train_loader, model, optimizer, scheduler, scaler, epoch, args): |
| 17 | batch_time = AverageMeter('Batch', ':2.2f') |
| 18 | data_time = AverageMeter('Data', ':2.2f') |
| 19 | lr = AverageMeter('Lr', ':1.6f') |
| 20 | loss_meter = AverageMeter('Loss', ':2.4f') |
| 21 | iou_meter = AverageMeter('IoU', ':2.2f') |
| 22 | pr_meter = AverageMeter('Prec@50', ':2.2f') |
| 23 | progress = ProgressMeter( |
| 24 | len(train_loader), |
| 25 | [batch_time, data_time, lr, loss_meter, iou_meter, pr_meter], |
| 26 | prefix="Training: Epoch=[{}/{}] ".format(epoch, args.epochs)) |
| 27 | |
| 28 | model.train() |
| 29 | time.sleep(2) |
| 30 | end = time.time() |
| 31 | |
| 32 | # size_list = [320, 352, 384, 416, 448, 480, 512] |
| 33 | # idx = np.random.choice(len(size_list)) |
| 34 | # new_size = size_list[idx] |
| 35 | |
| 36 | for i, (image, text, target, l_mask) in enumerate(train_loader): |
| 37 | data_time.update(time.time() - end) |
| 38 | # data |
| 39 | image = torch.stack(image).cuda(non_blocking=True) |
| 40 | text = torch.stack(text).cuda(non_blocking=True) |
| 41 | target = torch.stack(target).cuda(non_blocking=True) |
| 42 | l_mask = torch.stack(l_mask).cuda(non_blocking=True) |
| 43 | # # multi-scale training |
| 44 | # image = F.interpolate(image, size=(new_size, new_size), mode='bilinear', align_corners=True) |
| 45 | text = text.squeeze(1) |
| 46 | l_mask = l_mask.squeeze(1) |
| 47 | # forward |
| 48 | with amp.autocast(): |
| 49 | pred, target, loss = model(image, text, l_mask, target) |
| 50 | # backward |
| 51 | optimizer.zero_grad() |
| 52 | scaler.scale(loss).backward() |
| 53 | # if args.max_norm: |
| 54 | # torch.nn.utils.clip_grad_norm_(model.parameters(), args.max_norm) |
| 55 | scaler.step(optimizer) |
| 56 | scaler.update() |
| 57 | scheduler.step() |
| 58 | |
| 59 | # metric |
| 60 | iou, pr5 = trainMetricGPU(pred, target, 0.35) |
| 61 | dist.all_reduce(loss.detach()) |
| 62 | dist.all_reduce(iou) |
| 63 | dist.all_reduce(pr5) |
| 64 | loss = loss / dist.get_world_size() |
| 65 | iou = iou / dist.get_world_size() |
| 66 | pr5 = pr5 / dist.get_world_size() |
| 67 | |
| 68 | loss_meter.update(loss.item(), image.size(0)) |
| 69 | iou_meter.update(iou.item(), image.size(0)) |
| 70 | pr_meter.update(pr5.item(), image.size(0)) |
| 71 | lr.update(optimizer.param_groups[0]["lr"]) |
| 72 | batch_time.update(time.time() - end) |
| 73 | end = time.time() |
no test coverage detected