Function train_posenet

script/dm/pose_model.py:359–391

(args, train_dl, val_dl, model, epochs, optimizer, loss_func, scheduler, device, early_stopping)

    return train_loss_epoch_mean

def train_posenet(args, train_dl, val_dl, model, epochs, optimizer, loss_func, scheduler, device, early_stopping):
    writer = SummaryWriter()
    model_log = tqdm(total=0, position=1, bar_format='{desc}')
    for epoch in tqdm(range(epochs), desc='epochs'):
        
        # train 1 epoch
        train_loss = train_on_epoch(args, train_dl, model, optimizer, loss_func, device)
        writer.add_scalar("Loss/train", train_loss, epoch)
        
        # validate every epoch
        val_loss = eval_on_epoch(args, val_dl, model, optimizer, loss_func, device)
        writer.add_scalar("Loss/val", val_loss, epoch)
        
        # reduce LR on plateau
        scheduler.step(val_loss)
        writer.add_scalar("lr", optimizer.param_groups[0]['lr'], epoch)

        # logging
        tqdm.write('At epoch {0:6d} : train loss: {1:.4f}, val loss: {2:.4f}'.format(epoch, train_loss, val_loss))
                
        # check wether to early stop
        early_stopping(val_loss, model, epoch=epoch, save_multiple=(not args.no_save_multiple), save_all=args.save_all_ckpt)
        if early_stopping.early_stop:
            print("Early stopping")
            break

        model_log.set_description_str(f'Best val loss: {early_stopping.val_loss_min:.4f}')

        if epoch % args.i_eval == 0:
            get_error_in_q(args, val_dl, model, len(val_dl.dataset), device, batch_size=1)


    writer.flush()