train.py

# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License").
# You may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import argparse
from tqdm import tqdm

from algorithms import *
from datasets import *
from models import *
from samplers import *
from utils import *

TRAIN_ITERATIONS = 20000
CHECKPOINT_ITERATIONS = 1000
TEST_NUM_TASKS = 1000

args = argparse.ArgumentParser(
    description='few-shot-benchmark',
    formatter_class=argparse.ArgumentDefaultsHelpFormatter)
args.add_argument('--dataset', type=str, required=True)
args.add_argument('--model', type=str, required=True)
args.add_argument('--algorithm', type=str, required=True)
args.add_argument('--sampler', type=str, required=True)
args.add_argument('--ways', type=int, default=5)
args.add_argument('--support_shots', type=int, required=True)
args.add_argument('--query_shots', type=int, default=15)
args.add_argument('--params_path', type=str, default=None)
args = vars(args.parse_args())

if args['dataset'] == 'mini-imagenet':
    TRAIN_BATCH_SIZE = 16
elif args['dataset'] == 'tiered-imagenet':
    TRAIN_BATCH_SIZE = 32


def main():

    dataset = args['dataset']
    model = args['model']
    algorithm = args['algorithm']
    sampler = args['sampler']
    ways = args['ways']
    support_shots = args['support_shots']
    query_shots = args['query_shots']

    setup()

    train_tasks, val_tasks, test_tasks = DATASETS[dataset](
        ways, support_shots + query_shots, ways, support_shots + query_shots,
        'cpu')
    model = MODELS[model](ways)
    algo_model, adapt = ALGORITHMS[algorithm]
    model = algo_model(model)
    model.to('cuda')
    Sampler = SAMPLERS[sampler](TRAIN_BATCH_SIZE, TRAIN_ITERATIONS,
                                train_tasks, args)
    optimizer = torch.optim.Adam(model.parameters())

    best_meta_val_acc = 0.
    for iteration in tqdm(range(1, TRAIN_ITERATIONS + 1)):

        meta_train_accs = []
        sum_weights = 0.
        sum_weights_square = 0.
        optimizer.zero_grad()
        for _ in range(TRAIN_BATCH_SIZE):

            task = train_tasks.sample()
            loss, acc, weight = evaluate_task(
                task, ways, support_shots, query_shots, model, adapt, Sampler,
                True)
            meta_train_accs.append(acc.item())
            sum_weights += weight
            sum_weights_square += (weight ** 2.)

            (weight * loss).backward()

        if sum_weights == 0.:
            inv_effective_batch_size = 1.
        else:
            inv_effective_batch_size = sum_weights_square / (sum_weights ** 2.)
        for p in model.parameters():
            if hasattr(p, 'grad') and p.grad is not None:
                p.grad.data.mul_(inv_effective_batch_size)
        optimizer.step()

        if iteration % CHECKPOINT_ITERATIONS == 0:

            meta_val_accs = []
            meta_test_accs = []
            for idx in range(TEST_NUM_TASKS):

                task = val_tasks[idx]
                _, acc = evaluate_task(task, ways, support_shots,
                                       query_shots, model, adapt, None, False)
                meta_val_accs.append(acc.item())

                task = test_tasks[idx]
                _, acc = evaluate_task(task, ways, support_shots,
                                       query_shots, model, adapt, None, False)
                meta_test_accs.append(acc.item())

            meta_val_acc = np.mean(meta_val_accs)
            if meta_val_acc > best_meta_val_acc:
                best_meta_val_acc = meta_val_acc
                final_meta_test_acc = np.mean(meta_test_accs)
                final_meta_test_ci = 1.96 * np.std(meta_test_accs, ddof=1) \
                    / (len(meta_test_accs) ** .5)

            tqdm.write(
                '[%d]\t meta_train_accuracy:\t %.4f' %
                (iteration, np.mean(meta_train_accs)))
            tqdm.write(
                '[%d]\t meta_val_accuracy:\t %.4f' %
                (iteration, np.mean(meta_val_accs)))
            tqdm.write(
                '[%d]\t meta_test_accuracy:\t %.4f' %
                (iteration, np.mean(meta_test_accs)))
            tqdm.write('[%d]\t final_meta_test_accuracy:\t %.4f +/- %.4f' %
                       (iteration, final_meta_test_acc, final_meta_test_ci))


if __name__ == '__main__':

    main()