Function train

train_strandsVAE.py:147–254

(args, hyperparams, loader_train, loader_test, experiment_name, output_training_path)

def train(args, hyperparams, loader_train, loader_test, experiment_name, output_training_path):

    cb=create_callbacks(with_tensorboard=hyperparams.with_tensorboard,\
                        with_visualizer=hyperparams.with_visualizer,\
                        viewer_config_path=hyperparams.viewer_config_path,\
                        experiment_name=experiment_name)

    #create phases
    phases= [
        Phase('train', loader_train, grad=True),
        Phase('test', loader_test, grad=False),
    ]

    #model 
    model = StrandCodec(do_vae=hyperparams.enable_vae, 
                        scale_init=hyperparams.scale_init,
                        nr_verts_per_strand=hyperparams.nr_verts_per_strand, nr_values_to_decode=hyperparams.nr_values_to_decode,
                        dim_per_value_decoded=hyperparams.dim_per_value_decoded).to(args.device)
    model = torch.compile(model)

    #misc
    world2local=torch.compile(World2Local())
    loss_computer= torch.compile(StrandVAELoss())
    normalization_dict=loader_train.dataset.get_normalization_data()
    

    # #optimizer
    optimizer = torch.optim.AdamW (model.parameters(), amsgrad=False, lr=hyperparams.lr, weight_decay=0.0)
    lr_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=hyperparams.nr_iters_to_train)
    scheduler_warmup = UntunedLinearWarmup(optimizer)

    progress_bar = tqdm(range(0, hyperparams.nr_iters_to_train), desc="Training progress")


    is_in_training_loop=True
    while is_in_training_loop:

        for phase in phases:
            model.train(phase.grad)
            if hyperparams.enable_vae:
                model.encoder.do_vae=phase.grad #when testing we don't do any VAE stuff and rather just predict the mean

            cb.phase_started(phase=phase)
            cb.epoch_started(phase=phase)


            #run epoch 
            for batch in iter(phase.loader):
                cb.before_forward_pass(phase=phase)

                #progress
                if phase.grad and phase.iter_nr%100==0 :
                    progress_bar.update(100)

            
                #world_to_local
                with torch.no_grad():
                    gt_dict = prepare_gt_batch(batch, hyperparams, world2local, do_augmentation=phase.grad)