Function train_one_epoch

engine.py:34–179

(model: torch.nn.Module,
                    criterion: torch.nn.Module,
                    data_loader: Iterable,
                    optimizer: torch.optim.Optimizer,
                    device: torch.device,
                    epoch: int,
                    max_norm: float = 0,
                    wo_class_error=False,
                    lr_scheduler=None,
                    args=None,
                    logger=None,
                    ema_m=None,
                    tf_writer=None)

        return items

def train_one_epoch(model: torch.nn.Module,
                    criterion: torch.nn.Module,
                    data_loader: Iterable,
                    optimizer: torch.optim.Optimizer,
                    device: torch.device,
                    epoch: int,
                    max_norm: float = 0,
                    wo_class_error=False,
                    lr_scheduler=None,
                    args=None,
                    logger=None,
                    ema_m=None,
                    tf_writer=None):
    scaler = torch.cuda.amp.GradScaler(enabled=args.amp)

    try:
        need_tgt_for_training = args.use_dn
    except:
        need_tgt_for_training = False

    model.train()
    criterion.train()
    # criterion_smpl.to(device)
    metric_logger = utils.MetricLogger(delimiter='  ')
    metric_logger.add_meter(
        'lr', utils.SmoothedValue(window_size=1, fmt='{value:.6f}'))
    if not wo_class_error:
        metric_logger.add_meter(
            'class_error', utils.SmoothedValue(window_size=1,
                                               fmt='{value:.2f}'))
    header = 'Epoch: [{}]'.format(epoch)
    print_freq = 10

    _cnt = 0

    for step_i, data_batch in enumerate(metric_logger.log_every(data_loader,
                                              print_freq,
                                              header,
                                              logger=logger)):
        with torch.cuda.amp.autocast(enabled=args.amp):
            if need_tgt_for_training:
                outputs, targets, data_batch_nc = model(data_batch)
            else:
                outputs, targets, data_batch_nc = model(data_batch)
            
            ['hand_kp3d_4', 'face_kp3d_4', 'hand_kp2d_4',]
            loss_dict = criterion(outputs, targets, data_batch=data_batch_nc)
            weight_dict = criterion.weight_dict
            
            for k,v in weight_dict.items():
                for n in ['hand_kp3d_4', 'face_kp3d_4', 'hand_kp2d_4']:
                    if n in k:
                        weight_dict[k] = weight_dict[k]/10

            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)