Function train_on_epoch

script/dm/direct_pose_model.py:279–298

Perform 1 epoch of training with batch

(args, data_loaders, model, hwf, optimizer, half_res, device, **render_kwargs_test)

def train_on_epoch(args, data_loaders, model, hwf, optimizer, half_res, device, **render_kwargs_test):
    ''' Perform 1 epoch of training with batch '''
    model.train()
    model = freeze_bn_layer_train(model)

    # Prepare dataloaders for PoseNet, each batch contains (image, pose)
    train_dl, val_dl, test_dl = data_loaders
    total_loss = []
    total_psnr = []
    
    ####  Core optimization loop  #####
    for data, pose, img_idx in train_dl:
        # print("img_idx: {}, pose: {}".format(img_idx, pose) )
        # training one step with batch_size = args.batch_size
        loss, psnr = train_on_batch(args, data, model, pose, img_idx, hwf, optimizer, half_res, device, **render_kwargs_test)
        total_loss.append(loss.item())
        total_psnr.append(psnr.item())
    total_loss_mean = np.mean(total_loss)
    total_psnr_mean = np.mean(total_psnr)
    return total_loss_mean, total_psnr_mean

def save_val_result_7Scenes(args, epoch, val_dl, model, hwf, half_res, device, num_samples=1, **render_kwargs_test):
    ''' Perform inference on a random val image and save the result '''