train_meta.py

"""Stage 2: meta-train the gesture classifier G_y with MAML on top of the frozen
Stage-1 feature extractor, then evaluate N-way K-shot adaptation in the target
domain (paper Sec. 3.5, Algorithm 2, Fig. 10).

Meta-training: tasks are sampled from the source domain; for each task the
classifier adapts on the support set (inner loop, Eq. 16) and the query losses
of ``tasks_per_update`` tasks are summed for one meta-update (outer loop, Eq. 17).

Meta-testing: for each target-domain task the meta-learned classifier adapts on
the K-shot support set and is scored on the query set.

    python train_meta.py --hypes_yaml configs/meta/cross_env1.yaml

Ablations:
    --reinit_before_test   adapt a freshly initialised classifier in every test
                           episode, i.e. discard meta-training ("w/o meta-training")
    --test_noise_mean M --test_noise_std S
                           add N(M, S^2) noise to the target-domain inputs; the
                           original noise experiment used M=0.4, S=0.1
"""
import argparse
import os

import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
from tqdm import tqdm

from metatrans.data import EpisodicBatchSampler, EpisodicCSIDataset, default_transform
from metatrans.models import MAML, load_feature_extractor
from metatrans.utils import add_gaussian_noise, get_device, load_yaml, print_matrix, save_matrix, setup_train


def train_parser():
    parser = argparse.ArgumentParser()
    parser.add_argument('--hypes_yaml', type=str, required=True)
    parser.add_argument('--feature_extractor_pth', type=str, default=None,
                        help='Stage-1 checkpoint; overrides `feature_extractor_pth` in the config')
    parser.add_argument('--reinit_before_test', action='store_true',
                        help='ablation: adapt a freshly initialised classifier in every test episode')
    parser.add_argument('--test_noise_mean', type=float, default=0.0,
                        help='ablation: mean of Gaussian noise added to target-domain inputs')
    parser.add_argument('--test_noise_std', type=float, default=0.0,
                        help='ablation: std of Gaussian noise added to target-domain inputs')
    return parser.parse_args()


def episode_features(feature_extractor, x, n_way, n_support, n_query, noise=None):
    """Turn an episode batch [n_way, n_support + n_query, C, H, W] into class-major
    support / query features. ``noise`` is an optional (mean, std) pair."""
    x = x.view(n_way * (n_support + n_query), *x.size()[2:])
    if noise is not None:
        x = add_gaussian_noise(x, mean=noise[0], std=noise[1])
    with torch.no_grad():
        z = feature_extractor(x)
    z = z.view(n_way, n_support + n_query, -1)
    z_support = z[:, :n_support].contiguous().view(n_way * n_support, -1)
    z_query = z[:, n_support:].contiguous().view(n_way * n_query, -1)
    return z_support, z_query


def main():
    opt = train_parser()
    hypes = load_yaml(opt.hypes_yaml)
    meta_cfg = hypes['meta_learning']
    device = get_device()
    transform = default_transform()

    source_data = EpisodicCSIDataset.from_config(hypes['source_domain_data'], transform)
    target_data = EpisodicCSIDataset.from_config(hypes['target_domain_data'], transform)

    n_way = meta_cfg.get('n_way', 6)
    n_support = meta_cfg['shot']
    n_query = source_data.samples_per_class - n_support
    assert target_data.samples_per_class == source_data.samples_per_class

    loader_source = torch.utils.data.DataLoader(
        source_data, batch_sampler=EpisodicBatchSampler(len(source_data), n_way, meta_cfg['training']['n_episodes']),
        num_workers=0, pin_memory=True)
    loader_target = torch.utils.data.DataLoader(
        target_data, batch_sampler=EpisodicBatchSampler(len(target_data), n_way, meta_cfg['testing']['n_episodes']),
        num_workers=0, pin_memory=True)

    fe_path = opt.feature_extractor_pth or hypes['feature_extractor_pth']
    feature_extractor = load_feature_extractor(fe_path, map_location=device).to(device)
    feature_extractor.eval()

    inner_steps = meta_cfg.get('inner_steps', 5)
    inner_lr = meta_cfg.get('inner_lr', 0.05)
    tasks_per_update = meta_cfg.get('tasks_per_update', 3)
    classifier = MAML(n_way=n_way, inner_steps=inner_steps, inner_lr=inner_lr, device=device).to(device)
    optimizer = optim.Adam(classifier.parameters(), lr=meta_cfg.get('meta_lr', 1e-4))
    loss_fn = nn.CrossEntropyLoss()

    n_train_episodes = meta_cfg['training']['n_episodes']
    if n_train_episodes % tasks_per_update:
        print(f'Warning: {n_train_episodes} training episodes is not a multiple of tasks_per_update='
              f'{tasks_per_update}; the last {n_train_episodes % tasks_per_update} episode(s) '
              f'do not contribute to any meta-update.')

    saved_path = setup_train(hypes)
    log_acc_training = open(os.path.join(saved_path, 'acc_training.txt'), 'w')
    log_acc_testing = open(os.path.join(saved_path, 'acc_testing.txt'), 'w')

    # ---------------- meta-training on the source domain ----------------
    classifier.train()
    correct = 0.0
    total = 0.0
    loss_all = []
    pbar = tqdm(total=len(loader_source), desc="Meta-Training", miniters=0.3)
    for x, _ in loader_source:
        z_support, z_query = episode_features(feature_extractor, x.to(device), n_way, n_support, n_query)
        y_query = torch.from_numpy(np.repeat(range(n_way), n_query)).long().to(device)

        optimizer.zero_grad()
        classifier.adapt(z_support)
        y = classifier(z_query)
        loss = loss_fn(y, y_query)
        loss_all.append(loss)
        if len(loss_all) == tasks_per_update:
            torch.stack(loss_all).sum(0).backward()
            optimizer.step()
            loss_all = []

        _, predicted = torch.max(y.data, dim=1)
        total += n_query * n_way
        correct += predicted.eq(y_query.data).cpu().sum().numpy()
        log_acc_training.write(f'{correct / total}\n')
        pbar.set_postfix({"Acc": correct / total, "Loss": loss.item()})
        pbar.update(1)
    pbar.close()

    # ---------------- meta-testing on the target domain ----------------
    correct = 0.0
    total = 0.0
    confusion_matrix = np.zeros((n_way, n_way))
    test_noise = None
    if opt.test_noise_mean or opt.test_noise_std:
        test_noise = (opt.test_noise_mean, opt.test_noise_std)
    pbar = tqdm(total=len(loader_target), desc="Meta-Testing", miniters=0.3)
    for x, _ in loader_target:
        if opt.reinit_before_test:
            classifier = MAML(n_way=n_way, inner_steps=meta_cfg.get('reinit_inner_steps', 6),
                              inner_lr=meta_cfg.get('reinit_inner_lr', 0.03), device=device).to(device)
            classifier.eval()
        z_support, z_query = episode_features(feature_extractor, x.to(device), n_way, n_support, n_query,
                                              noise=test_noise)
        y_query = torch.from_numpy(np.repeat(range(n_way), n_query)).long().to(device)

        classifier.adapt(z_support)
        with torch.no_grad():
            y = classifier(z_query)
        _, predicted = torch.max(y.data, 1)
        total += n_query * n_way
        correct += predicted.eq(y_query.data).cpu().sum().numpy()
        # NOTE: indices are episode-local class slots (the sampler shuffles the classes
        # of every episode), not gesture ids.
        for t, p in zip(y_query.cpu().numpy(), predicted.cpu().numpy()):
            confusion_matrix[t][p] += 1

        log_acc_testing.write(f'{correct / total}\n')
        pbar.set_postfix({"Acc": correct / total})
        pbar.update(1)
    pbar.close()

    acc = correct / total
    print(confusion_matrix)
    torch.save(classifier.state_dict(), os.path.join(saved_path, 'meta_classifier.pth'))
    print('Meta-testing accuracy: {:.2%}'.format(acc))
    matrix = confusion_matrix / total * n_way
    print_matrix(matrix)
    save_matrix(matrix, os.path.join(saved_path, 'best_confusion_matrix_testing.txt'))

    log_acc_training.close()
    log_acc_testing.close()


if __name__ == "__main__":
    main()