Function train_feature

script/run_feature.py:232–422

(args, train_dl, val_dl, test_dl, hwf, i_split, near, far)

    return train_loss

def train_feature(args, train_dl, val_dl, test_dl, hwf, i_split, near, far):

    # # load pretrained PoseNet model
    if args.DFNet_s:
        feat_model = DFNet_s()
    else:
        feat_model = DFNet()
    
    if args.pretrain_model_path != '':
        print("load posenet from ", args.pretrain_model_path)
        feat_model.load_state_dict(torch.load(args.pretrain_model_path))
    
    # # Freeze BN to not updating gamma and beta
    if args.freezeBN:
        feat_model = freeze_bn_layer(feat_model)

    feat_model.to(device)
    # summary(feat_model, (3, 240, 427))

    # set optimizer
    optimizer = optim.Adam(feat_model.parameters(), lr=args.learning_rate)
    scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, factor=0.95, patience=args.patience[1], verbose=True)

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

    # loss function
    loss_func = nn.MSELoss(reduction='mean')

    i_train, i_val, i_test = i_split
    # Cast intrinsics to right types
    H, W, focal = hwf
    H, W = int(H), int(W)
    hwf = [H, W, focal]

    # Create log dir and copy the config file
    basedir = args.basedir
    expname = args.expname
    os.makedirs(os.path.join(basedir, expname), exist_ok=True)
    f = os.path.join(basedir, expname, 'args.txt')
    with open(f, 'w') as file:
        for arg in sorted(vars(args)):
            attr = getattr(args, arg)
            file.write('{} = {}\n'.format(arg, attr))
    if args.config is not None:
        f = os.path.join(basedir, expname, 'config.txt')
        with open(f, 'w') as file:
            file.write(open(args.config, 'r').read())

    # load NeRF
    _, render_kwargs_test, start, _, _ = create_nerf(args)
    global_step = start

    bds_dict = {
        'near' : near,
        'far' : far,
    }
    # render_kwargs_train.update(bds_dict)