Beispiel #1
0
def loadWarmUpState(filename):
    print('  load pre-trained model file: {}'.format(filename))
    warmup_resume_file = get_resume_file(filename)
    tmp = torch.load(warmup_resume_file)
    if tmp is not None:
        state = tmp['state']
        state_keys = list(state.keys())
        for i, key in enumerate(state_keys):
            if "feature." in key:  # and '.7.' not in key:
                newkey = key.replace("feature.", "")
                state[newkey] = state.pop(key)
            else:
                state.pop(key)
    else:
        raise ValueError('No warm_up file')
    return state
Beispiel #2
0
    test_few_shot_params = dict(n_way=params.test_n_way,
                                n_support=params.n_shot)
    val_datamgr = SetDataManager(image_size,
                                 n_query=n_query,
                                 **test_few_shot_params)
    val_loader = val_datamgr.get_data_loader(val_file, aug=False)

    model = LFTNet(params, tf_path=params.tf_dir)
    model.cuda()

    # resume training
    start_epoch = params.start_epoch
    stop_epoch = params.stop_epoch
    if params.resume != '':
        resume_file = get_resume_file(
            '%s/checkpoints/%s' % (params.save_dir, params.resume),
            params.resume_epoch)
        if resume_file is not None:
            start_epoch = model.resume(resume_file)
            print('  resume the training with at {} epoch (model file {})'.
                  format(start_epoch, params.resume))
        else:
            raise ValueError('No resume file')
    # load pre-trained feature encoder
    else:
        if params.warmup == 'gg3b0':
            raise Exception(
                'Must provide pre-trained feature-encoder file using --warmup option!'
            )
        model.model.feature.load_state_dict(load_warmup_state(
            '%s/checkpoints/%s' % (params.save_dir, params.warmup),
Beispiel #3
0
                     tf_path=params.tf_dir,
                     **train_few_shot_params)
    else:
        raise ValueError('Unknown method')
    model = model.cuda()

    max_acc = 0
    total_it = 0
    start_epoch = params.start_epoch
    stop_epoch = params.stop_epoch
    if params.method not in ['baseline', 'baseline++']:
        stop_epoch = params.stop_epoch * model.batch_size

    # resume/warmup or not
    if params.resume != '':
        resume_file = get_resume_file('%s/checkpoints/%s' %
                                      (params.save_dir, params.resume))
        if resume_file is not None:
            print('  load resume file: {}'.format(params.resume))
            tmp = torch.load(resume_file)
            max_acc = tmp['max_acc']
            total_it = tmp['total_it']
            start_epoch = tmp['epoch'] + 1
            model.load_state_dict(tmp['state'])
            model.optimizer.load_state_dict(tmp['optimizer'])
        else:
            raise ValueError('No resume file')
    elif params.warmup != 'gg3b0':
        warmup_resume_file = get_resume_file('%s/checkpoints/%s' %
                                             (params.save_dir, params.warmup))
        if warmup_resume_file is not None:
            print('  load pretrained file: {}'.format(params.warmup))