(model: torch.nn.Module, criterion: torch.nn.Module,
data_loader: Iterable, optimizer: torch.optim.Optimizer,
device: torch.device, scaler: torch.cuda.amp.GradScaler,
epoch: int, max_norm: float = 0, fp16=False)
| 23 | |
| 24 | |
| 25 | def train_one_epoch(model: torch.nn.Module, criterion: torch.nn.Module, |
| 26 | data_loader: Iterable, optimizer: torch.optim.Optimizer, |
| 27 | device: torch.device, scaler: torch.cuda.amp.GradScaler, |
| 28 | epoch: int, max_norm: float = 0, fp16=False): |
| 29 | model.train() |
| 30 | criterion.train() |
| 31 | tensor_type = torch.cuda.HalfTensor if fp16 else torch.cuda.FloatTensor |
| 32 | metric_logger = utils.MetricLogger(delimiter=" ") |
| 33 | metric_logger.add_meter('lr', utils.SmoothedValue(window_size=1, fmt='{value:.6f}')) |
| 34 | metric_logger.add_meter('class_error', utils.SmoothedValue(window_size=1, fmt='{value:.2f}')) |
| 35 | metric_logger.add_meter('grad_norm', utils.SmoothedValue(window_size=1, fmt='{value:.2f}')) |
| 36 | header = 'Epoch: [{}]'.format(epoch) |
| 37 | print_freq = 10 |
| 38 | |
| 39 | prefetcher = data_prefetcher(data_loader, device, prefetch=True) |
| 40 | samples, targets = prefetcher.next() |
| 41 | |
| 42 | # for samples, targets in metric_logger.log_every(data_loader, print_freq, header): |
| 43 | for _ in metric_logger.log_every(range(len(data_loader)), print_freq, header): |
| 44 | samples.tensors = samples.tensors.type(tensor_type) |
| 45 | samples.mask = samples.mask.type(tensor_type) |
| 46 | |
| 47 | with torch.cuda.amp.autocast(enabled=fp16): |
| 48 | outputs, pre_outputs, pre_targets = model([samples, targets]) |
| 49 | loss_dict = criterion(outputs, targets, pre_outputs, pre_targets) |
| 50 | weight_dict = criterion.weight_dict |
| 51 | losses = sum(loss_dict[k] * weight_dict[k] for k in loss_dict.keys() if k in weight_dict) |
| 52 | |
| 53 | # reduce losses over all GPUs for logging purposes |
| 54 | loss_dict_reduced = utils.reduce_dict(loss_dict) |
| 55 | loss_dict_reduced_unscaled = {f'{k}_unscaled': v |
| 56 | for k, v in loss_dict_reduced.items()} |
| 57 | loss_dict_reduced_scaled = {k: v * weight_dict[k] |
| 58 | for k, v in loss_dict_reduced.items() if k in weight_dict} |
| 59 | losses_reduced_scaled = sum(loss_dict_reduced_scaled.values()) |
| 60 | |
| 61 | loss_value = losses_reduced_scaled.item() |
| 62 | |
| 63 | if not math.isfinite(loss_value): |
| 64 | print("Loss is {}, stopping training".format(loss_value)) |
| 65 | print(loss_dict_reduced) |
| 66 | sys.exit(1) |
| 67 | |
| 68 | optimizer.zero_grad() |
| 69 | scaler.scale(losses).backward() |
| 70 | scaler.unscale_(optimizer) |
| 71 | if max_norm > 0: |
| 72 | grad_total_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm) |
| 73 | else: |
| 74 | grad_total_norm = utils.get_total_grad_norm(model.parameters(), max_norm) |
| 75 | scaler.step(optimizer) |
| 76 | scaler.update() |
| 77 | |
| 78 | metric_logger.update(loss=loss_value, **loss_dict_reduced_scaled, **loss_dict_reduced_unscaled) |
| 79 | metric_logger.update(class_error=loss_dict_reduced['class_error']) |
| 80 | metric_logger.update(lr=optimizer.param_groups[0]["lr"]) |
| 81 | metric_logger.update(grad_norm=grad_total_norm) |
| 82 |
no test coverage detected