Function main

pytorch-model/classification_train.py:178–246

()

def main():
    
    # --- Check for GPU ---
    device = "cuda:0" if torch.cuda.is_available() else "cpu"
    
    # --- Args ---
    args = opts.get_train_args()
    print("\n" + "-" * 30 + " Args " + "-" * 30)
    for k, v in vars(args).items():
        print(f"{k}: {v}")
    print()
    
    # --- Model and viz save dir ---
    model_save_dir = constants.get_model_dir(args.dataset, args.model_type, args.model_name)
    viz_save_dir = constants.get_viz_dir(args.dataset, args.model_type, args.model_name)
    if os.path.isdir(model_save_dir):
        raise RuntimeError(f"Error: {model_save_dir} already exists! Exiting...")
    elif os.path.isdir(viz_save_dir):
        raise RuntimeError(f"Error: {viz_save_dir} already exists! Exiting...")
    else:
        print(f"--> Creating directory {model_save_dir}...")
        os.makedirs(model_save_dir)
        print(f"--> Creating directory {viz_save_dir}...")
        os.makedirs(viz_save_dir)
    print("Done!\n")

    # --- Setup dataset ---
    print("--> Setting up dataset...")
    train_dataset = datasets.DATASETS[args.dataset](mode="train")
    val_dataset = datasets.DATASETS[args.dataset](mode="val")
    print("Done!\n")

    # --- Dataloaders ---
    print("--> Setting up dataloaders...")
    train_dataloader = DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True, num_workers=1)
    val_dataloader = DataLoader(val_dataset, batch_size=args.batch_size, shuffle=False, num_workers=1)
    print("Done!\n")

    # --- Setup model ---
    print("--> Setting up model...")
    model = models.MODEL_TYPES[args.model_type](train_dataset)
    # torch.cuda.set_device(device)
    # model = model.cuda(device)
    print("Done!\n")

    # --- Optimizer ---
    print("--> Setting up optimizer/criterion...")
    opt = torch.optim.Adam(model.parameters(), lr=args.lr)
    
    # --- Loss fn ---
    criterion = nn.CrossEntropyLoss(weight=train_dataset.get_weights())#.cuda(constants.GPU)
    # criterion = nn.CrossEntropyLoss()#.cuda(constants.GPU)
    print("Done!\n")

    # --- Train ---
    train_losses, train_accs, val_losses, val_accs =\
        train(args, model, train_dataloader, val_dataloader, criterion, opt)