Function train

script/train.py:95–136

()

    return

def train():
    print(parser.format_values())
    # Load data
    if args.dataset_type == '7Scenes':
        train_dl, val_dl, test_dl, hwf, i_split, near, far = load_7Scenes_dataloader(args)
    elif args.dataset_type == 'Cambridge':
        train_dl, val_dl, test_dl, hwf, i_split, near, far = load_Cambridge_dataloader(args)
    else:
        print("please choose dataset_type: 7Scenes or Cambridge, exiting...")
        sys.exit()

    ### pose regression module, here requires a pretrained DFNet for Pose Estimator F
    assert(args.pretrain_model_path != '') # make sure to add a valid PATH using --pretrain_model_path
    # load pretrained DFNet model
    model = load_exisiting_model(args)

    if args.freezeBN:
        model = freeze_bn_layer(model)
    model.to(device)

    ### feature extraction module, here requires a pretrained DFNet for Feature Extractor G using --pretrain_featurenet_path
    if args.pretrain_featurenet_path == '':
        print('Use the same DFNet for Feature Extraction and Pose Regression')
        feat_model = load_exisiting_model(args)
    else: 
        # you can optionally load different pretrained DFNet for feature extractor and pose estimator
        feat_model = load_exisiting_model(args, isFeatureNet=True)

    feat_model.eval()
    feat_model.to(device)

    # set optimizer
    optimizer = optim.Adam(model.parameters(), lr=args.learning_rate) #weight_decay=weight_decay, **kwargs

    # set callbacks parameters
    early_stopping = EarlyStopping(args, patience=args.patience[0], verbose=False)

    # start training
    if args.dataset_type == '7Scenes':
        train_feature_matching(args, model, feat_model, optimizer, i_split, hwf, near, far, device, early_stopping, train_dl=train_dl, val_dl=val_dl, test_dl=test_dl)
    elif args.dataset_type == 'Cambridge':
        train_feature_matching(args, model, feat_model, optimizer, i_split, hwf, near, far, device, early_stopping, train_dl=train_dl, val_dl=val_dl, test_dl=test_dl)

def eval():
    print(parser.format_values())