| |
| |
| |
| |
| |
| |
| |
| |
| |
|
|
|
|
| """ |
| Train and eval functions used in main.py |
| """ |
| import math |
| import os |
| import sys |
| from typing import Iterable |
|
|
| import torch |
| import util.misc as utils |
|
|
| from datasets.data_prefetcher import data_dict_to_cuda |
|
|
|
|
| def train_one_epoch_mot(model: torch.nn.Module, criterion: torch.nn.Module, |
| data_loader: Iterable, optimizer: torch.optim.Optimizer, |
| device: torch.device, epoch: int, max_norm: float = 0): |
| model.train() |
| criterion.train() |
| metric_logger = utils.MetricLogger(delimiter=" ") |
| metric_logger.add_meter('lr', utils.SmoothedValue(window_size=1, fmt='{value:.6f}')) |
| |
| metric_logger.add_meter('grad_norm', utils.SmoothedValue(window_size=1, fmt='{value:.2f}')) |
| header = 'Epoch: [{}]'.format(epoch) |
| print_freq = 10 |
|
|
| |
| for data_dict in metric_logger.log_every(data_loader, print_freq, header): |
| data_dict = data_dict_to_cuda(data_dict, device) |
| outputs = model(data_dict) |
|
|
| loss_dict = criterion(outputs, data_dict) |
| |
| weight_dict = criterion.weight_dict |
| losses = sum(loss_dict[k] * weight_dict[k] for k in loss_dict.keys() if k in weight_dict) |
|
|
| |
| loss_dict_reduced = utils.reduce_dict(loss_dict) |
| |
| |
| loss_dict_reduced_scaled = {k: v * weight_dict[k] |
| for k, v in loss_dict_reduced.items() if k in weight_dict} |
| losses_reduced_scaled = sum(loss_dict_reduced_scaled.values()) |
|
|
| loss_value = losses_reduced_scaled.item() |
|
|
| if not math.isfinite(loss_value): |
| print("Loss is {}, stopping training".format(loss_value)) |
| print(loss_dict_reduced) |
| sys.exit(1) |
|
|
| optimizer.zero_grad() |
| losses.backward() |
| if max_norm > 0: |
| grad_total_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm) |
| else: |
| grad_total_norm = utils.get_total_grad_norm(model.parameters(), max_norm) |
| optimizer.step() |
|
|
| |
| metric_logger.update(loss=loss_value, **loss_dict_reduced_scaled) |
| |
| metric_logger.update(lr=optimizer.param_groups[0]["lr"]) |
| metric_logger.update(grad_norm=grad_total_norm) |
| |
|
|
| metric_logger.synchronize_between_processes() |
| print("Averaged stats:", metric_logger) |
| return {k: meter.global_avg for k, meter in metric_logger.meters.items()} |
|
|
