PyTorch框架实战系列（3）——空间变换器网络STN

神经网络在对图像识别的实际应用过程中，经常会遇到这样的问题：需要识别的目标只是图片的一小部分；目标区域大小不一；目标的视角有差异或者是扭曲的。

这些情形如果不做任何处理，直接使用样本，对于CNN模型的效果就会造成一定消极影响。比如，对于手写文字的识别，我们都希望输入的样本是按文字切割好的、大小一致的、清晰而工整的。所以对于上述问题，我们常常会在模型训练之前先对样本做一定的预处理，但这些处理往往是复杂而困难的。那能不能让神经网络自己学习如何对样本做有效的预处理呢？

空间变换器网络（STN），是解决这一类问题的一种神经网络模型，可以自我学习如何更好的在输入图像上做有益的空间变换，比如裁剪、缩放和校正。更为方便的一点是，它可以嵌入现有的CNN模型中。

本例教程：https://pytorch.org/tutorials/intermediate/spatial_transformer_tutorial.html

中文教程：http://pytorch123.com/FourSection/SpatialTranNet/

STN论文地址：Spatial Transformer Networks

1、手写数字识别——MNIST数据集

MNIST数据集是很多人入门图像识别时，必然会下手的一个经典数据集。我们通过使用CNN模型和嵌入了STN的CNN对比，来实践空间变换网络的作用。

下载数据集：

from torchvision import datasets, transforms
from torch.utils.data import DataLoader

# 存储路径
root = './data'
# 下载MNIST 数据集
train_set = datasets.MNIST(root=root, train=True, download=True)
test_set = datasets.MNIST(root=root, train=False, download=True)

下载完后，我们看一下数据集里都有什么：

# 查看样本
print('train_set', len(train_set))
print('test_set', len(test_set))
sample = next(iter(train_set))
print(sample)
# 图片转化为tensor，查看最大最小值
totensor = transforms.ToTensor()
t = totensor(sample[0])
print(t.size())
print('min', t.min().item(), 'max', t.max().item())

可以看到，训练集有60000个样本，测试集有10000个，每个样本包含图片数据和它的标签。

图片是PIL.Image格式的，28*28的灰度图。

PIL.Image可以直接转换为张量，因为是灰度图，所以是1通道。

数值范围是[0, 1]之间。

这样，我们知道了如何通过transformers把样本数据转换为模型需要的，归一化到[-1, 1]之间的张量数据集：

# 数据集转换为归一化张量
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5,), (0.5,))
])
train_set = datasets.MNIST(root=root, train=True, download=False, transform=transform)
test_set = datasets.MNIST(root=root, train=False, download=False, transform=transform)
# 查看样本
sample = next(iter(train_set))
print('min', sample[0].min().item(), 'max', sample[0].max().item())

可以看到张量数值转为[-1, 1]了。

根据数据集就可以构造训练用的迭代器了：

# 构造迭代器
batch_size = 64
train_loader = DataLoader(dataset=train_set, batch_size=batch_size, shuffle=True)
test_loader = DataLoader(dataset=test_set, batch_size=batch_size, shuffle=True)
# 迭代器输出的张量
samples, labels = next(iter(train_loader))
print(samples.size(), labels.size())

得到了模型所需的标准四维张量[样本数, 通道数, 高, 宽]。

再利用OpenCV预览一下输入样本，直观感受一下：

# 预览图片
import cv2
import torchvision

imgs = torchvision.utils.make_grid(samples)
# 通道转置到最内维度
imgs = imgs.numpy().transpose(1, 2, 0)
# 逆归一化
imgs = imgs * 0.5 + 0.5
cv2.imshow('win',imgs)
cv2.waitKey(0)

2、卷积神经网络训练

设计一个简单的两层卷积的CNN，MNIST数据集识别难度很低，所以不需要太复杂。

CNN.py

import torch.nn as nn


class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        # 卷积层
        self.convs = nn.Sequential(
            nn.Conv2d(1, 64, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=4),
            nn.Conv2d(64, 128, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=3),
        )
        # 线性层
        self.linear = nn.Sequential(
            nn.Dropout(0.5),
            nn.Linear(in_features=512, out_features=10),
        )

    def forward(self, x):
        x = self.convs(x)
        x = x.view(x.size()[0], -1)
        x = self.linear(x)
        # print(x.size())
        return x


if __name__ == '__main__':
    import torch

    x = torch.rand(1, 1, 28, 28)
    model = Model()
    print(model)
    y = model(x)
    print(y)

训练模型：

train.py

# -*- coding: utf-8 -*-
from torchvision import datasets, transforms
from torch.utils.data import DataLoader
import torch.optim as optim
import torch.nn as nn
import torch
import numpy as np
from sklearn import metrics
import time
from datetime import timedelta


def get_time_dif(start_time):
    """获取已使用时间"""
    end_time = time.time()
    time_dif = end_time - start_time
    return timedelta(seconds=int(round(time_dif)))

def train(save_path, model, trainloader, testloader):
    start_time = time.time()
    # 训练模式
    model.train()
    # 指定损失函数和优化器
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=0.001)

    total_batch = 0  # 记录进行到多少batch
    dev_best_loss = float('inf')  # 记录验证集最佳损失率
    last_improve = 0  # 记录上次验证集loss下降的batch数
    flag = False  # 记录是否很久没有效果提升

    # 批次训练
    for epoch in range(20):
        print('Epoch [{}/{}]'.format(epoch + 1, 20))
        # 从迭代器中按mini-batch训练
        for trains, labels in trainloader:
            outputs = model(trains)
            # 模型梯度归零
            model.zero_grad()
            # 损失函数反馈传递
            loss = criterion(outputs, labels)
            loss.backward()
            # 执行优化
            optimizer.step()
            # 每多少轮在测试集上查看训练的效果
            if total_batch % 200 == 0:
                # 获得训练集准确率
                true = labels.data
                predic = torch.max(outputs.data, 1)[1]
                train_acc = metrics.accuracy_score(true, predic)
                # 如果验证集上继续收敛则保存模型参数
                dev_acc, dev_loss = evaluate(model, testloader)
                if dev_loss < dev_best_loss:
                    dev_best_loss = dev_loss
                    torch.save(model.state_dict(), save_path)
                    improve = '*'
                    last_improve = total_batch
                else:
                    improve = ''
                time_dif = get_time_dif(start_time)
                # 训练成果
                msg = 'Iter: {0:>6},  Train Loss: {1:>5.2},  Train Acc: {2:>6.2%},  Val Loss: {3:>5.2},  Val Acc: {4:>6.2%},  Time: {5}  {6}'
                print(msg.format(total_batch, loss.item(), train_acc, dev_loss, dev_acc, time_dif, improve))
                # 恢复训练
                model.train()
            total_batch += 1
            # 验证集loss超过多少batch没下降，结束训练
            if total_batch - last_improve > 1000:
                print("Training Finished ...")
                # torch.save(model.state_dict(), save_path)
                flag = True
                break
        if flag:
            break

    # 使用测试集测试评估模型
    model_test(save_path, model, testloader)

# 测试模型
def model_test(save_path, model, testloader):
    # 加载模型参数
    model.load_state_dict(torch.load(save_path))
    # 模型测试评估
    test_acc, test_loss, test_report, test_confusion = evaluate(model, testloader, test=True)
    msg = 'Test Loss: {0:>5.2},  Test Acc: {1:>6.2%}'
    print(msg.format(test_loss, test_acc))
    print("Confusion Matrix...")
    print(test_confusion)
    print("Classification Report...")
    print(test_report)

# 验证模型
def evaluate(model, dataloader, test=False):
    # 模型预测模式
    model.eval()
    loss_total = 0
    predict_all = np.array([], dtype=int)
    labels_all = np.array([], dtype=int)
    loss_func = nn.CrossEntropyLoss()
    # 不记录模型梯度
    with torch.no_grad():
        for texts, labels in dataloader:
            outputs = model(texts)
            loss = loss_func(outputs, labels)
            loss_total += loss
            labels = labels.data.numpy()
            predic = torch.max(outputs.data, 1)[1].numpy()
            labels_all = np.append(labels_all, labels)
            predict_all = np.append(predict_all, predic)
    # 验证集准确度
    acc = metrics.accuracy_score(labels_all, predict_all)
    # 给出模型测试结果，评估报告和混淆矩阵
    if test:
        report = metrics.classification_report(labels_all, predict_all, digits=4)
        confusion = metrics.confusion_matrix(labels_all, predict_all)
        return acc, loss_total / len(dataloader), report, confusion
    else:
        return acc, loss_total / len(dataloader)


if __name__ == '__main__':
    root = './data'    # 数据集
    save_path = './cnn_model.pth'   # 模型保存路径

    # 设置随机数种子
    np.random.seed(1)
    torch.manual_seed(1)
    torch.cuda.manual_seed_all(1)
    # 保证每次结果一样
    torch.backends.cudnn.deterministic = True

    # 数据集转换为归一化张量
    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.5,), (0.5,))
    ])
    trainset = datasets.MNIST(root=root, train=True, download=False, transform=transform)
    testset = datasets.MNIST(root=root, train=False, download=False, transform=transform)

    # 查看数据集大小
    print('trainset', len(trainset))
    print('testset', len(testset))

    batch_size = 50    # mini-batch
    # 构造迭代器
    trainloader = DataLoader(dataset=trainset, batch_size=batch_size, shuffle=True)
    testloader = DataLoader(dataset=testset, batch_size=batch_size, shuffle=True)

    # 迭代器输出的张量
    samples, labels = next(iter(trainloader))
    print(samples.size(), labels.size())

    # 训练测试CNN模型
    from CNN import Model

    # 模型实例化
    model = Model()
    train(save_path, model, trainloader, testloader)

 输出结果如下：

本例使用MNIST数据集验证STN的作用真是比较尴尬，主要这个数据集识别难度太低了，即使只用全连接层的神经网络，也有98%以上的准确率。这个简单的CNN达到几乎99%的准确率，实在没有什么提升空间了……

不过我们还是试验一下嵌入STN的CNN是否有积极的作用吧。

3、空间变换网络（STN）

论文中给出的网络结构如下：

STN分为三个部分：

1、定位网络（Localisation net）：也是一个卷积神经网络，用于学习变换参数矩阵θ，即裁剪区域和坐标转置的偏移量。定位网络并不是直接从数据集学习如何转换，而是嵌入在图像分类的CNN中，根据分类CNN的损失最小化原则自动学习空间转换的参数θ。也就是说，其驱动力不是样本本身，而是努力提高样本分类的准确性、减少损失值。

2、网格生成器（Grid genator）：根据变换参数θ，生成一个输入图像和转换图像之间的坐标对照网格。

3、采样器（Sampler）：将输入图像按网格的坐标对照关系，转换为一个新图像，转发给分类CNN。

论文中直观的展示了如何对图像进行网格转换：

4、STN-CNN训练

下面设计STN-CNN模型：

STNCNN.py

import torch.nn as nn
import torch


class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        # 定位网络-卷积层
        self.localization_convs = nn.Sequential(
            nn.Conv2d(1, 8, kernel_size=7),
            nn.ReLU(),
            nn.MaxPool2d(2, stride=2),
            nn.Conv2d(8, 10, kernel_size=5),
            nn.ReLU(),
            nn.MaxPool2d(2, stride=2),
        )
        # 定位网络-线性层
        self.localization_linear = nn.Sequential(
            nn.Linear(in_features=10 * 3 * 3, out_features=32),
            nn.ReLU(),
            nn.Linear(in_features=32, out_features=2 * 3)
        )
        # 初始化定位网络仿射矩阵的权重/偏置，即是初始化θ值。使得图片的空间转换从原始图像开始。
        self.localization_linear[2].weight.data.zero_()
        self.localization_linear[2].bias.data.copy_(torch.tensor([1, 0, 0,
                                                                  0, 1, 0], dtype=torch.float))

        # 图片分类-卷积层
        self.convs = nn.Sequential(
            nn.Conv2d(1, 64, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=4),
            nn.Conv2d(64, 128, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=3),
        )
        # 图片分类-线性层
        self.linear = nn.Sequential(
            nn.Dropout(0.5),
            nn.Linear(in_features=512, out_features=10),
        )

    # 空间变换器网络，转发图片张量
    def stn(self, x):
        # 使用CNN对图像结构定位，生成变换参数矩阵θ（2*3矩阵）
        x2 = self.localization_convs(x)
        x2 = x2.view(x2.size()[0], -1)
        x2 = self.localization_linear(x2)
        theta = x2.view(x2.size()[0], 2, 3)   # [1, 2, 3]
        # print(theta)
        '''
        2D空间变换初始θ参数应置为tensor([[[1., 0., 0.],
                                        [0., 1., 0.]]])
        '''
        # 网格生成器，根据θ建立原图片的坐标仿射矩阵
        grid = nn.functional.affine_grid(theta, x.size(), align_corners=True)   # [1, 28, 28, 2]
        # 采样器，根据网格对原图片进行转换，转发给CNN分类网络
        x = nn.functional.grid_sample(x, grid, align_corners=True)  # [1, 1, 28, 28]
        return x

    def forward(self, x):
        x = self.stn(x)
        # print(x.size())
        x = self.convs(x)
        x = x.view(x.size()[0], -1)
        x = self.linear(x)
        # print(x.size())
        return x


if __name__ == '__main__':
    x = torch.rand(1, 1, 28, 28)
    model = Model()
    print(model)
    y = model(x)
    print(y)

 对train.py做一下修改，跑一下嵌入STN的CNN训练效果是否有所提升：

    # # 训练测试CNN模型
    # from CNN import Model

    # 训练测试STN-CNN模型
    from STNCNN import Model
    save_path = './stn_cnn_model.pth'

    # 模型实例化
    model = Model()
    train(save_path, model, trainloader, testloader)

 训练结果如下：

可以看到对比：loss从0.035降低到了0.026，准确率从98.97%提升到了99.04%。

刚才说了，MNIST提升空间太小，不过从结果看STN对于图像识别的确是有一定的积极作用。

对比一下训练过程中模型在验证集上的损失值变化：

红色：CNN

蓝色：STN-CNN

5、STN转换图像可视化

STN网络输入的是图像，输出的是空间转换后的图像。这就比较有意思了，因为我们可以在数据集上直观的感受到图像转换的效果。

教程中使用matplotlib可视化，不好玩，个人比较喜欢OpenCV。随机在测试集上找64个图片看一下STN的结果。

visual.py

from torchvision import datasets, transforms
from torch.utils.data import DataLoader
import cv2
import torchvision


def visualize(samples):
    # 预览图片
    imgs = torchvision.utils.make_grid(samples)
    # 通道转置到最内维度
    imgs = imgs.numpy().transpose(1, 2, 0)
    # 逆归一化
    imgs = imgs * 0.5 + 0.5
    cv2.imshow('win', imgs)
    cv2.waitKey(0)


if __name__ == '__main__':
    # 存储路径
    root = './data'

    # 数据集转换为归一化张量
    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.5,), (0.5,))
    ])
    test_set = datasets.MNIST(root=root, train=False, download=False, transform=transform)

    # 构造迭代器
    batch_size = 64
    test_loader = DataLoader(dataset=test_set, batch_size=batch_size, shuffle=True)
    # 迭代器输出的张量
    samples, labels = next(iter(test_loader))
    print(samples.size(), labels.size())
    # 原图片
    visualize(samples)

    from STNCNN import Model
    import torch

    save_path = './stn_cnn_model.pth'
    model = Model()
    model.load_state_dict(torch.load(save_path))
    transform_samples = model.stn(samples).detach()
    # 空间转换后的图片
    visualize(transform_samples)

展示结果：

              

MNIST数据集是已经按数字裁剪好而且大小一致的图片，所以我们感受到的STN效果基本上只是将歪斜的数字扶正了。不过，这种效果已经不得不让人感叹数学的强大和神经网络编程模式的优势。

STN论文中还举例了鸟类图片识别的应用，可以推断出，STN在背景噪音明显的图片识别、视频类模式识别可以有很广泛的应用场景。