PyTorch 空间变压器网络教程

空间变换器网络（STN）是一种视觉注意力机制，可以增强神经网络对输入图像的空间变换能力。通过学习如何对输入图像进行裁剪、缩放和方向校正等操作，STN 能够提升模型的几何不变性。

一、什么是空间变换器网络？

STN 是一种可插入到任何现有卷积神经网络（CNN）中的网络结构，它包含三个主要部分：

1. 本地化网络 ：常规的 CNN，用于预测变换参数。
2. 网格生成器 ：根据变换参数生成目标图像的采样网格。
3. 采样器 ：将采样网格应用于输入图像，生成变换后的图像。

这种机制允许网络学习如何对输入图像进行空间变换，从而增强模型对图像的几何变换适应能力。

二、加载数据并可视化

我们将使用经典的 MNIST 数据集来演示 STN 的应用。

import torch
import torchvision
from torchvision import datasets, transforms
import matplotlib.pyplot as plt


## 检测设备
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")


## 加载 MNIST 数据集
train_loader = torch.utils.data.DataLoader(
    datasets.MNIST(root='.', train=True, download=True,
                   transform=transforms.Compose([
                       transforms.ToTensor(),
                       transforms.Normalize((0.1307,), (0.3081,))
                   ])),
    batch_size=64, shuffle=True, num_workers=4
)


test_loader = torch.utils.data.DataLoader(
    datasets.MNIST(root='.', train=False, transform=transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.1307,), (0.3081,))
    ])),
    batch_size=64, shuffle=True, num_workers=4
)


## 可视化数据集中的图像
data, _ = next(iter(train_loader))
plt.figure(figsize=(10, 10))
for i in range(25):
    plt.subplot(5, 5, i + 1)
    plt.imshow(data[i][0].numpy(), cmap='gray')
    plt.axis('off')
plt.show()

三、定义 STN 模型

我们将定义一个包含 STN 的 CNN 模型。

import torch.nn as nn
import torch.nn.functional as F


class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 10, kernel_size=5)
        self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
        self.conv2_drop = nn.Dropout2d()
        self.fc1 = nn.Linear(320, 50)
        self.fc2 = nn.Linear(50, 10)


        # STN 的本地化网络
        self.localization = nn.Sequential(
            nn.Conv2d(1, 8, kernel_size=7),
            nn.MaxPool2d(2, stride=2),
            nn.ReLU(True),
            nn.Conv2d(8, 10, kernel_size=5),
            nn.MaxPool2d(2, stride=2),
            nn.ReLU(True)
        )


        # STN 的全连接层，用于预测变换参数
        self.fc_loc = nn.Sequential(
            nn.Linear(10 * 3 * 3, 32),
            nn.ReLU(True),
            nn.Linear(32, 3 * 2)
        )


        # 初始化变换参数为单位矩阵
        self.fc_loc[2].weight.data.zero_()
        self.fc_loc[2].bias.data.copy_(torch.tensor([1, 0, 0, 0, 1, 0], dtype=torch.float))


    # STN 的前向传播函数
    def stn(self, x):
        xs = self.localization(x)
        xs = xs.view(-1, 10 * 3 * 3)
        theta = self.fc_loc(xs)
        theta = theta.view(-1, 2, 3)
        grid = F.affine_grid(theta, x.size())
        x = F.grid_sample(x, grid)
        return x


    # 整个网络的前向传播函数
    def forward(self, x):
        x = self.stn(x)
        x = F.relu(F.max_pool2d(self.conv1(x), 2))
        x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))
        x = x.view(-1, 320)
        x = F.relu(self.fc1(x))
        x = F.dropout(x, training=self.training)
        x = self.fc2(x)
        return F.log_softmax(x, dim=1)


model = Net().to(device)

四、训练和测试模型

编写训练和测试函数，使用 SGD 优化器训练模型。

import torch.optim as optim


## 定义优化器
optimizer = optim.SGD(model.parameters(), lr=0.01)


def train(epoch):
    model.train()
    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad()
        output = model(data)
        loss = F.nll_loss(output, target)
        loss.backward()
        optimizer.step()
        if batch_idx % 500 == 0:
            print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                epoch, batch_idx * len(data), len(train_loader.dataset),
                100. * batch_idx / len(train_loader), loss.item()))


def test():
    with torch.no_grad():
        model.eval()
        test_loss = 0
        correct = 0
        for data, target in test_loader:
            data, target = data.to(device), target.to(device)
            output = model(data)
            test_loss += F.nll_loss(output, target, reduction='sum').item()
            pred = output.max(1, keepdim=True)[1]
            correct += pred.eq(target.view_as(pred)).sum().item()
        test_loss /= len(test_loader.dataset)
        print('\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(
            test_loss, correct, len(test_loader.dataset),
            100. * correct / len(test_loader.dataset)))


## 进行多轮训练
for epoch in range(1, 21):
    train(epoch)
    test()

五、可视化 STN 结果

训练完成后，我们可以可视化 STN 学习到的变换效果。

def convert_image_np(inp):
    inp = inp.numpy().transpose((1, 2, 0))
    mean = np.array([0.485, 0.456, 0.406])
    std = np.array([0.229, 0.224, 0.225])
    inp = std * inp + mean
    inp = np.clip(inp, 0, 1)
    return inp


def visualize_stn():
    with torch.no_grad():
        data = next(iter(test_loader))[0].to(device)
        input_tensor = data.cpu()
        transformed_input_tensor = model.stn(data).cpu()
        in_grid = convert_image_np(torchvision.utils.make_grid(input_tensor))
        out_grid = convert_image_np(torchvision.utils.make_grid(transformed_input_tensor))
        f, axarr = plt.subplots(1, 2)
        axarr[0].imshow(in_grid)
        axarr[0].set_title('原始图像')
        axarr[1].imshow(out_grid)
        axarr[1].set_title('STN 变换后的图像')
        plt.show()


visualize_stn()

六、总结

通过本教程，你学会了如何在 PyTorch 中实现空间变换器网络（STN），并在 MNIST 数据集上应用它。STN 的引入使得模型能够学习到对输入图像进行空间变换的能力，从而提高了模型的几何不变性和分类性能。在编程狮（W3Cschool）网站上，你可以找到更多关于 PyTorch 的详细教程和实战案例，帮助你进一步提升深度学习技能，成为人工智能领域的编程大神。