train_uda.py

"""Stage 1: train the feature extractor G_f with unsupervised domain adaptation
(paper Sec. 3.3-3.4, Algorithm 1).

Labelled source-domain samples train G_f + G_y with cross-entropy (Eq. 13);
unlabelled target-domain samples join the source ones in the domain
discriminator G_d behind a gradient reversal layer (Eq. 10-12, 14), and the
total loss is L = L_y + L_d (Eq. 15).

Set ``uda: false`` in the config to drop G_d and the target data (the
"w/o UDA" ablation); G_f is then trained on the source domain only.

    python train_uda.py --hypes_yaml configs/uda/cross_env1_adam.yaml
"""
import argparse
import os

import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from tqdm import tqdm

from metatrans.data import CSIHeatmapDataset, default_transform
from metatrans.models import GestureNet
from metatrans.utils import 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)
    return parser.parse_args()


def build_optimizer(hypes, params):
    cfg = hypes['optimizer']
    if cfg['core_method'] == 'Adam':
        return optim.Adam(params, lr=cfg['lr'])
    if cfg['core_method'] == 'SGD':
        return optim.SGD(params, lr=cfg['lr'], momentum=cfg.get('momentum', 0))
    raise ValueError('Optimizer not supported')


def build_scheduler(hypes, optimizer):
    cfg = hypes['lr_scheduler']
    if cfg['core_method'] == 'step':
        return optim.lr_scheduler.StepLR(optimizer, step_size=cfg['step_size'], gamma=cfg['gamma'])
    if cfg['core_method'] == 'multistep':
        return optim.lr_scheduler.MultiStepLR(optimizer, milestones=cfg['step_size'], gamma=cfg['gamma'])
    raise ValueError('lr_scheduler not supported')


def main():
    opt = train_parser()
    hypes = load_yaml(opt.hypes_yaml)
    use_uda = hypes.get('uda', True)
    num_classes = len(hypes['test_data']['ges'])
    batch_size = hypes['train_params']['batch_size']
    transform = default_transform()
    device = get_device()

    print('Dataset Building')
    train_loader = DataLoader(CSIHeatmapDataset.from_config(hypes['source_domain_data'], transform),
                              batch_size=batch_size, shuffle=True)
    if use_uda:
        target_loader = DataLoader(CSIHeatmapDataset.from_config(hypes['target_domain_data'], transform),
                                   batch_size=batch_size, shuffle=True)
    test_loader = DataLoader(CSIHeatmapDataset.from_config(hypes['test_data'], transform),
                             batch_size=batch_size, shuffle=True)

    print('Creating Model')
    model = GestureNet(num_classes=num_classes, use_uda=use_uda).to(device)
    class_criterion = nn.CrossEntropyLoss()
    domain_criterion = nn.BCEWithLogitsLoss()
    optimizer = build_optimizer(hypes, model.parameters())
    scheduler = build_scheduler(hypes, optimizer)

    saved_path = setup_train(hypes)
    log_loss1 = open(os.path.join(saved_path, 'loss1.txt'), 'w')  # gesture classification loss L_y
    log_loss2 = open(os.path.join(saved_path, 'loss2.txt'), 'w') if use_uda else None  # domain loss L_d
    log_acc = open(os.path.join(saved_path, 'acc.txt'), 'w')

    epochs = hypes['train_params']['epoches']
    global_test_acc = [0.0]
    global_confusion_matrix = np.zeros((num_classes, num_classes))
    best_confusion_matrix = None

    for epoch in range(epochs):
        model.train()
        sum_loss = 0.0
        correct = 0.0
        total = 0.0
        if use_uda:
            target_iter = iter(target_loader)

        pbar = tqdm(total=len(train_loader), desc=f"Epoch-Training {epoch + 1}/{epochs}", miniters=0.3)
        for i, (inputs, labels) in enumerate(train_loader):
            inputs, labels = inputs.to(device), labels.to(device)
            optimizer.zero_grad()

            if use_uda:
                # Restart the target iterator before it runs out, cycling over the target set.
                if (i + 1) % len(target_loader) >= len(target_loader) - 1:
                    target_iter = iter(target_loader)
                target_inputs, _ = next(target_iter)
                target_inputs = target_inputs.to(device)

                # GRL coefficient ramps from 0 to 1 over training (Ganin et al., JMLR 2016).
                p = float(i + epoch * len(train_loader)) / (epochs * len(train_loader))
                grl_lambda = 2. / (1. + np.exp(-10 * p)) - 1

                source_domain = torch.zeros((len(labels), 1), dtype=torch.float).to(device)
                class_pred, domain_pred = model(inputs, grl_lambda)
                loss_cls = class_criterion(class_pred, labels)
                loss_source = domain_criterion(domain_pred, source_domain)

                target_domain = torch.ones((len(target_inputs), 1), dtype=torch.float).to(device)
                _, domain_pred = model(target_inputs, grl_lambda)
                loss_target = domain_criterion(domain_pred, target_domain)

                loss_domain = loss_source + loss_target
                loss = loss_cls + loss_domain
                log_loss2.write("{:.6f}\n".format(loss_domain.item()))
            else:
                class_pred, _ = model(inputs)
                loss_cls = class_criterion(class_pred, labels)
                loss = loss_cls
            log_loss1.write("{:.6f}\n".format(loss_cls.item()))

            loss.backward()
            optimizer.step()

            sum_loss += loss.item()
            _, predicted = torch.max(class_pred.data, dim=1)
            total += labels.size(0)
            correct += predicted.eq(labels.data).cpu().sum().numpy()
            pbar.set_postfix({"Acc": correct / total, "Loss": sum_loss / (i + 1)})
            pbar.update(1)
        # Passing the epoch index is kept from the original runs: the decay lands one epoch
        # after each milestone (it emits a deprecation warning on recent PyTorch).
        scheduler.step(epoch)
        pbar.close()

        model.eval()
        with torch.no_grad():
            correct = 0.0
            total = 0.0
            tmp_confusion_matrix = np.zeros((num_classes, num_classes))

            pbar = tqdm(total=len(test_loader), desc=f"Epoch-Testing {epoch + 1}/{epochs}", miniters=0.3)
            for inputs, labels in test_loader:
                inputs, labels = inputs.to(device), labels.to(device)
                outputs, _ = model(inputs)
                _, predicted = torch.max(outputs.data, 1)
                total += labels.size(0)
                correct += predicted.eq(labels.data).cpu().sum().numpy()
                for x, y in zip(labels.cpu().numpy(), predicted.cpu().numpy()):
                    tmp_confusion_matrix[x][y] += 1
                pbar.set_postfix({"Acc": correct / total})
                pbar.update(1)
            pbar.close()

            acc = correct / total
            log_acc.write("{:.6f}\n".format(acc))
            if acc > max(global_test_acc):
                print('Saving model......')
                global_confusion_matrix = tmp_confusion_matrix
                torch.save(model.feature_extractor.state_dict(),
                           os.path.join(saved_path, f'epoch_{epoch}_featureExtractor.pth'))
                torch.save(model.state_dict(), os.path.join(saved_path, f'epoch_{epoch}_model.pth'))
                best_confusion_matrix = global_confusion_matrix / total * num_classes
            global_test_acc.append(acc)
            print(hypes['name'] + ' accuracy: {:.2%}, best accuracy: {:.2%}'.format(acc, max(global_test_acc)))
            print_matrix(global_confusion_matrix / total * num_classes)

    if best_confusion_matrix is not None:
        save_matrix(best_confusion_matrix, os.path.join(saved_path, 'best_confusion_matrix.txt'))

    for f in (log_loss1, log_loss2, log_acc):
        if f is not None:
            f.close()


if __name__ == '__main__':
    main()