Function train

pytorch-model/classification_train.py:112–156

Trains and evaluates model.

(args, model, train_dataloader, val_dataloader, criterion, opt)

def train(args, model, train_dataloader, val_dataloader, criterion, opt):
    """
    Trains and evaluates model.
    """
    print("\n--- Begin training! ---\n")
    train_losses = dict()
    train_accs = dict()
    val_losses = dict()
    val_accs = dict()
    model_save_dir = constants.get_model_dir(args.dataset, args.model_type, args.model_name)
    viz_path = constants.get_viz_dir(args.dataset, args.model_type, args.model_name)

    for epoch in range(args.num_epochs):

        if epoch % args.eval_every == 0:
            # --- Time to evaluate! ---
            val_avg_loss, val_avg_acc, _ = eval_model(model, val_dataloader, criterion, args)
            val_losses[epoch] = val_avg_loss
            val_accs[epoch] = val_avg_acc
            
            # --- Report and plot losses/ious ---
            print(f" Val avg loss: {val_avg_loss} | Val avg acc: {val_avg_acc}\n")
            # viz_utils.plot_losses_ious(val_losses, val_ious, viz_path, prefix="val")
            
        # --- Train ---
        train_avg_loss, train_avg_acc, _ = train_one_epoch(model, train_dataloader, criterion, opt, args)
        train_losses[epoch] = train_avg_acc
        train_accs[epoch] = train_avg_loss

        # --- Print results ---
        if epoch % args.print_every == 0:
            print(f"Epoch: {epoch} | Train avg loss: {train_avg_loss} | Train avg acc: {train_avg_acc}")

        # --- Plot loss/metrics so far (TODO: do this every epoch?) ---
        # --- Then save all train stats ---
        # viz_utils.plot_losses_ious(train_losses, train_accs, viz_path, prefix="train")
        save_train_stats(train_losses, train_accs, val_losses, val_accs, args)
        
        # --- Only save actual model files every so often ---
        if epoch % args.save_every == 0:
            model_save_path = os.path.join(model_save_dir, f"model_epoch_{epoch}.pth")
            print(f"Saving model to {model_save_path}...")
            torch.save(model.state_dict(), model_save_path)

    return train_losses, train_accs, val_losses, val_accs


def save_train_stats(train_losses, train_accs, val_losses, val_accs, args):