Function do_train

engine.py:32–206

(
    args,
    model,
    accelerator,
    optimizer,
    dataloaders,
    best_val_metrics,
    logger
)

def do_train(
    args,
    model,
    accelerator,
    optimizer,
    dataloaders,
    best_val_metrics,
    logger
):
    
    if accelerator.is_main_process:
        logger.log_messages(f"call with args: {args}")
        logger.log_messages(f"{model}")
    
    curr_iter = args.start_epoch * len(dataloaders['train'])
    max_iters = args.max_epoch * len(dataloaders['train'])
    
    time_delta = SmoothedValue(window_size=10)
    loss_avg = SmoothedValue(window_size=10)
    loss_break_down_avg = {}
    
    model.train()
    accelerator.wait_for_everyone()
    
    for curr_epoch in range(args.start_epoch, args.max_epoch):
        
        for batch_idx, batch_data_label in enumerate(dataloaders['train']):
            
            curr_time = time.time()
            
            ### core for model training
            
            curr_iter = curr_epoch * len(dataloaders['train']) + batch_idx
            curr_lr = adjust_learning_rate(args, optimizer, curr_iter, max_iters)
            
            with accelerator.accumulate(model):
                
                with accelerator.autocast():
                    outputs = model(batch_data_label)
                loss = outputs['loss']
                
                # sanity check, skip the infinite loss
                if not math.isfinite(loss.item()):
                    logger.log_messages("Loss in not finite. Skip this iteration.")
                    model.eval()
                    model.train()
                    torch.cuda.empty_cache()
                    continue
                
                accelerator.backward(loss)
                if args.clip_gradient > 0:
                    accelerator.clip_grad_norm_(model.parameters(), args.clip_gradient)
                
                optimizer.step()
                optimizer.zero_grad()
            
            ### logging training loss status