基于DANN的图像分类任务迁移学习

注:本博客的数据和任务来自NTU-ML2020作业，Kaggle网址为Kaggle.

数据预处理

我们要进行迁移学习的对象是10000张32x32x3的有标签正常照片，共有10类，和另外100000张人类画的手绘图，28x28x1黑白照片，类别也是10类但无标签。我们希望做到，让模型从有标签的原始分布数据中学到的知识能应用于无标签的，相似但与原始分布不相同的目标分布中，并提高黑白手绘图的正确率。
为此，训练前还要对数据做预处理。首先让原始分布的图像和目标分布的图像尽可能相似，我们要做有色图转灰度图，然后做边缘检测。为了模型的输入维度相同，要把28x28转为32x32.此外还可以增加一些平移旋转来让学习更鲁棒。

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.autograd import Function
import torch.optim as optim
import torchvision.transforms as transforms
from torchvision.datasets import ImageFolder
from torch.utils.data import DataLoader
import cv2
import matplotlib.pyplot as plt

# 在transform中使用转灰度-canny边缘提取-水平移动-小幅度旋转-转张量操作

source_transform = transforms.Compose([
    transforms.Grayscale(),
    transforms.Lambda(lambda x: cv2.Canny(np.array(x), 170, 300)),
    transforms.ToPILImage(),
    transforms.RandomHorizontalFlip(),
    transforms.RandomRotation(15, fill=(0,)),
    transforms.ToTensor(),
])
target_transform = transforms.Compose([
    transforms.Grayscale(),
    transforms.Resize((32, 32)),
    transforms.RandomHorizontalFlip(),
    transforms.RandomRotation(15, fill=(0,)),
    transforms.ToTensor(),
])

# 读取数据集，分为source和target两部分

source_dataset = ImageFolder('E:/real_or_drawing/train_data', transform=source_transform)
target_dataset = ImageFolder('E:/real_or_drawing/test_data', transform=target_transform)

source_dataloader = DataLoader(source_dataset, batch_size=32, shuffle=True)
target_dataloader = DataLoader(target_dataset, batch_size=32, shuffle=True)
test_dataloader = DataLoader(target_dataset, batch_size=128, shuffle=False)

DANN

Domain-Adversarial Training of NNs，值域对抗学习。这种算法是我们这里将要用的迁移学习方法，它被提出的起因是让CNN能够同时用于不同分布的数据，如果模型直接接收原值域的数据分布进行训练，即使原分布和目标分布有类似的地方，在接收目标值域的数据时，也会出现相当异常的特征提取和分类结果。我们可以理解为是模型在源数据分布上出现了过拟合(并不是对数据的过拟合)，在接收一些没有见到过的数据时自然会表现不佳。

解决这个问题最好的办法就是让模型在训练时也接收目标数据分布的数据。但是目标数据分布是无标签的，我们要用什么标准来训练模型呢？回忆CNN的架构，CNN使用卷积-池化的特征提取层来提取图片特征，后接全连接层进行预测。我们只需要让特征提取层既能提取原数据分布的特征，又能提取目标数据分布的特征，这样全连接层就能对两种值域但具有相同特征的数据进行同样的分类，从而目标数据分布的输入也很有可能被正确分类。

那么问题就变成了如何训练输入两个不同分布的数据，输出却是同种分布的特征提取层。回忆GAN的架构，我们让分布朝着源数据分布发展的方法是建立判别器，让判别器能分辨两种数据，而让生成器改变参数骗过判别器。这里也可以用同样的思想，我们建立能分辨原始分布和目标分布的二分类判别器，把特征提取层和二分类判别层接在一起。首先训练判别器，让判别器能分辨两类数据分布。然后训练特征提取层，逆梯度更新让特征提取层生成能骗过判别器的数据(目标输出0.5).如此训练多次直到特征提取层能把两种值域的输入变成同种分布的输出。

但是只是用GAN方法train特征提取层并不明智，因为我们的目标输出只有0-1的二分类，训练很有可能只是让特征提取层提取到一些没有用的特征。因此我们要一边训练正常的标签预测任务，一边训练判别器的判别任务和混淆两类输入的任务。这可能需要自己定义特殊的loss function

最后，我们就获得了能同时提取两个值域的特征的特征提取层，它后面的多分类层就可以对目标分布的数据做出还算称心如意的预测。

模型、训练、测试代码

这里使用类VGG(用多个3x3的卷积核代替大型卷积核以节约参数)的搭建方式，写一个高度卷积的特征提取层

class FeatureExtractor(nn.Module):
    def __init__(self):
        super(FeatureExtractor, self).__init__()

        self.conv = nn.Sequential(
            nn.Conv2d(1, 64, 3, 1, 1),
            nn.BatchNorm2d(64),
            nn.ReLU(),
            nn.MaxPool2d(2),

            nn.Conv2d(64, 128, 3, 1, 1),
            nn.BatchNorm2d(128),
            nn.ReLU(),
            nn.MaxPool2d(2),

            nn.Conv2d(128, 256, 3, 1, 1),
            nn.BatchNorm2d(256),
            nn.ReLU(),
            nn.MaxPool2d(2),

            nn.Conv2d(256, 256, 3, 1, 1),
            nn.BatchNorm2d(256),
            nn.ReLU(),
            nn.MaxPool2d(2),

            nn.Conv2d(256, 512, 3, 1, 1),
            nn.BatchNorm2d(512),
            nn.ReLU(),
            nn.MaxPool2d(2)
        )
        
    def forward(self, x):
        x = self.conv(x).squeeze()
        return x

#值域分类器，即GAN中的discriminator
class DomainClassifier(nn.Module):
    def __init__(self):
        super(DomainClassifier, self).__init__()

        self.layer = nn.Sequential(
            nn.Linear(512, 512),
            nn.BatchNorm1d(512),
            nn.ReLU(),

            nn.Linear(512, 512),
            nn.BatchNorm1d(512),
            nn.ReLU(),

            nn.Linear(512, 512),
            nn.BatchNorm1d(512),
            nn.ReLU(),

            nn.Linear(512, 512),
            nn.BatchNorm1d(512),
            nn.ReLU(),

            nn.Linear(512, 1),
        )

    def forward(self, h):
        y = self.layer(h)
        return y

#标签预测器，对特征作进一步分类
class LabelPredictor(nn.Module):

    def __init__(self):
        super(LabelPredictor, self).__init__()

        self.layer = nn.Sequential(
            nn.Linear(512, 512),
            nn.ReLU(),

            nn.Linear(512, 512),
            nn.ReLU(),

            nn.Linear(512, 10),
        )

    def forward(self, h):
        c = self.layer(h)
        return c

feature_extractor = FeatureExtractor().cuda()
label_predictor = LabelPredictor().cuda()
domain_classifier = DomainClassifier().cuda()

# 多分类使用交叉熵损失进行训练
class_criterion = nn.CrossEntropyLoss()
# domain_classifier的输出是1维，要先sigmoid转概率再计算交叉熵，使用BCEWithlogits
domain_criterion = nn.BCEWithLogitsLoss()

# 使用adam训练
optimizer_F = optim.Adam(feature_extractor.parameters())
optimizer_C = optim.Adam(label_predictor.parameters())
optimizer_D = optim.Adam(domain_classifier.parameters())

我们训练200个epoch，让数据尽量收敛

def train_epoch(source_dataloader, target_dataloader, lamb):
    '''
      Args:
        source_dataloader: source data的dataloader
        target_dataloader: target data的dataloader
        lamb: 对抗的lamb系数
    '''

    # D loss: Domain Classifier的loss
    # F loss: Feature Extrator & Label Predictor的loss
    # total_hit: 計算目前對了幾筆 total_num: 目前經過了幾筆
    running_D_loss, running_F_loss = 0.0, 0.0
    total_hit, total_num = 0.0, 0.0

    for i, ((source_data, source_label), (target_data, _)) in enumerate(zip(source_dataloader, target_dataloader)):

        source_data = source_data.cuda()
        source_label = source_label.cuda()
        target_data = target_data.cuda()
        
        # 把source data和target data混在一起，否则batch_norm会出错
        mixed_data = torch.cat([source_data, target_data], dim=0)
        # 设置判别器的目标标签
        domain_label = torch.zeros([source_data.shape[0] + target_data.shape[0], 1]).cuda()
        domain_label[:source_data.shape[0]] = 1

        # Step 1 : 训练Domain Classifier
        feature = feature_extractor(mixed_data)
        # 这里detach feature，因为不需要更新extractor的参数
        domain_logits = domain_classifier(feature.detach())
        loss = domain_criterion(domain_logits, domain_label)
        running_D_loss+= loss.item()
        loss.backward()
        optimizer_D.step()

        # Step 2 : 训练Feature Extractor和Domain Classifier
        class_logits = label_predictor(feature[:source_data.shape[0]])
        domain_logits = domain_classifier(feature)
        # 这里使用的loss是原值域数据的任务分类交叉熵损失减去，原值域数据和目标值域数据的判别损失
        # 因为我们想让extractor骗过判别器，判别损失加负号，而且为了调控训练使用lambda作为系数
        loss = class_criterion(class_logits, source_label) - lamb * domain_criterion(domain_logits, domain_label)
        running_F_loss+= loss.item()
        loss.backward()
        optimizer_F.step()
        optimizer_C.step()

        optimizer_D.zero_grad()
        optimizer_F.zero_grad()
        optimizer_C.zero_grad()

        total_hit += torch.sum(torch.argmax(class_logits, dim=1) == source_label).item()
        total_num += source_data.shape[0]
        print(i, end='\r')

    return running_D_loss / (i+1), running_F_loss / (i+1), total_hit / total_num


# 训练50 epochs
for epoch in range(50):
    train_D_loss, train_F_loss, train_acc = train_epoch(source_dataloader, target_dataloader, lamb=0.1)

    torch.save(feature_extractor.state_dict(), f'extractor_model.bin')
    torch.save(domain_classifier.state_dict(), f'domain_model.bin')
    torch.save(label_predictor.state_dict(), f'predictor_model.bin')

    print('epoch {:>3d}: train D loss: {:6.4f}, train F loss: {:6.4f}, acc {:6.4f}'.format(epoch, train_D_loss, train_F_loss, train_acc))

训练好之后可以看见原值域上的训练集正确率有98以上，想看手绘图片的正确率可以在Kaggle上提交一下。我们这里随便打印一些手绘图片和模型预测的标签。

feature_extractor.load_state_dict(torch.load('extractor_model.bin'))
domain_classifier.load_state_dict(torch.load('domain_model.bin'))
label_predictor.load_state_dict(torch.load('predictor_model.bin'))

for i, (data, _) in enumerate(test_dataloader):
    break
    
class_logits = label_predictor(feature_extractor(data.cuda()))

#我们看50张手绘图的预测

def no_axis_show(img, title='', cmap=None):
    fig = plt.imshow(img, interpolation='nearest', cmap=cmap)
    fig.axes.get_xaxis().set_visible(False)
    fig.axes.get_yaxis().set_visible(False)
    plt.title(title)

titles = ['horse', 'bed', 'clock', 'apple', 'cat', 'plane', 'television', 'dog', 'dolphin', 'spider']
plt.figure(figsize=(18, 18))
data = data.cuda()
for i in range(50):
    plt.subplot(5, 10, i+1)
    label = torch.argmax(class_logits[i]).cpu().detach().numpy()
    img = data[i].cpu().detach().numpy().reshape(32,32)
    fig = no_axis_show(img, title=titles[label])

正确率不能说有多高，但是模型似乎学会了分辨一些特征比较明显的图片。

值域

把特征提取层得到的特征用PCA降维可以在2D平面上看到值域的分布。

在不使用DANN时，原值域和目标值域是分开的，这样的特征投入全连接层必然不work。但是当我们强制让模型把两种数据的特征混在一起，就变成右图，这时目标值域的特征有机会被正确分类。