P3_DANN_USPS_training.py

import math
from itertools import chain
from pathlib import Path

import numpy as np
import torch
from torch import nn
from torch.utils.data import DataLoader
from torch.utils.tensorboard.writer import SummaryWriter
from torchvision import transforms
from tqdm import tqdm

from digit_dataloader import digit_dataset
from P3_USPS_model import DomainClassifier, FeatureExtractor, LabelPredictor


# https://github.com/NaJaeMin92/pytorch_DANN
def rm_tree(pth: Path):
    if pth.is_dir():
        for child in pth.iterdir():
            if child.is_file():
                child.unlink()
            else:
                rm_tree(child)
        pth.rmdir()


def cycle(iterable):
    while True:
        for x in iterable:
            yield x


source_train_set = digit_dataset(
    # [0.4631, 0.4666, 0.4195], [0.1979, 0.1845, 0.2083]
    root='hw2_data/digits/mnistm/data',
    transform=transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize([0.4631, 0.4666, 0.4195],
                             [0.1979, 0.1845, 0.2083])
    ]),
    label_csv='hw2_data/digits/mnistm/train.csv'
)
target_train_set = digit_dataset(
    # [0.2570, 0.2570, 0.2570], [0.3372, 0.3372, 0.3372]
    root='hw2_data/digits/usps/data',
    transform=transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize([0.2573, 0.2573, 0.2573],
                             [0.3373, 0.3373, 0.3373])
    ]),
    label_csv='hw2_data/digits/usps/train.csv'
)
target_val_set = digit_dataset(
    # [0.2570, 0.2570, 0.2570], [0.3372, 0.3372, 0.3372]
    root='hw2_data/digits/usps/data',
    transform=transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize([0.2573, 0.2573, 0.2573],
                             [0.3373, 0.3373, 0.3373])
    ]),
    label_csv='hw2_data/digits/usps/val.csv'
)

batch_size = 256
source_train_loader = DataLoader(
    source_train_set, batch_size, shuffle=True, num_workers=6)
target_train_loader = DataLoader(
    target_train_set, batch_size, shuffle=True, num_workers=6)
target_val_loader = DataLoader(
    target_val_set, 2 * batch_size, shuffle=False, num_workers=6)

target_train_loader = iter(cycle(target_train_loader))

num_epochs = 300
lr = 0.0003
gamma = 10
device = 'cuda' if torch.cuda.is_available() else 'cpu'
ckpt_path = Path('./P3_USPS_ckpt')
tb_path = Path('./P3_USPS_tb')

rm_tree(ckpt_path)
rm_tree(tb_path)

ckpt_path.mkdir(exist_ok=True)
tb_path.mkdir(exist_ok=True)

writer = SummaryWriter(tb_path)


F = FeatureExtractor().to(device)
L = LabelPredictor().to(device)
D = DomainClassifier().to(device)

label_loss_fn = nn.CrossEntropyLoss()
domain_loss_fn = nn.BCEWithLogitsLoss()
optim = torch.optim.Adam(
    chain(F.parameters(), L.parameters(), D.parameters()), lr=lr)

current_step = 0
total_steps = num_epochs * len(source_train_loader)
best_target_acc = 0.3
for epoch in range(num_epochs):
    for (src_x, src_y), (tgt_x, _) in tqdm(zip(source_train_loader, target_train_loader), total=len(source_train_loader)):
        src_x = src_x.to(device, non_blocking=True)
        src_y = src_y.to(device, non_blocking=True)
        tgt_x = tgt_x.to(device, non_blocking=True)

        # scheduling
        p = current_step / total_steps
        lambda_ = 2.0 / (1.0 + math.exp(-gamma * p)) - 1
        optim.param_groups[0]['lr'] = lr / (1.0 + gamma * p) ** 0.75

        # label classification loss
        source_feature = F(src_x)
        source_logits = L(source_feature)
        label_loss = label_loss_fn(source_logits, src_y)

        # domain discriminator loss
        all_x = torch.cat([src_x, tgt_x], dim=0)
        all_feature = F(all_x)
        domain_logits = D(all_feature, lambda_)
        domain_logits = domain_logits.squeeze()
        domain_labels = torch.cat([
            # source=0
            torch.zeros(src_x.shape[0], dtype=torch.float, device=device),
            # target=1
            torch.ones(tgt_x.shape[0], dtype=torch.float, device=device)
        ], dim=0)
        domain_loss = domain_loss_fn(domain_logits, domain_labels)

        writer.add_scalars('training', {
                           'label_loss': label_loss, 'domain_loss': domain_loss}, global_step=current_step)

        loss = label_loss + domain_loss
        loss.backward()
        optim.step()
        optim.zero_grad()

        current_step += 1

    # validation
    for model in [F, L, D]:
        model.eval()
    va_acc = 0
    for tgt_x, tgt_y in tqdm(target_val_loader):
        tgt_x = tgt_x.to(device)
        tgt_y = tgt_y.cpu().numpy()

        with torch.no_grad():
            logits = L(F(tgt_x))
        pred = logits.argmax(-1).cpu().numpy()
        va_acc += np.mean((pred == tgt_y).astype(int))
    va_acc /= len(target_val_loader)
    for model in [F, L, D]:
        model.train()
    writer.add_scalar('accuracy/validation', va_acc, global_step=current_step)

    print(f"epoch: {epoch}, va_acc: {va_acc}")
    if va_acc >= best_target_acc:
        best_target_acc = va_acc
        torch.save(F.state_dict(), ckpt_path / f'best_F.pth')
        torch.save(L.state_dict(), ckpt_path / f'best_L.pth')
        torch.save(D.state_dict(), ckpt_path / f'best_D.pth')
        print(f"[new model saved]")