空间变换网络（Spatial Transformer）

CNN 靠卷积和池化获得了一定的平移不变性：图片里的物体挪个位置，网络多半还认得。但遇到旋转、缩放、倾斜，它就有点抓瞎。传统解法是数据增强：训练时把图随机转一转、缩一缩，让网络「见多识广」。这条路有效，但有个缺点：变换方式是外面预设的，网络只能被动适应。

空间变换网络（Spatial Transformer Networks，简称 STN）换个思路：把「摆正图片」做成网络的一部分，让网络自己学怎么变换输入对自己最有利。它出自 DeepMind 2015 年的论文，是视觉注意力机制的一个经典应用。

一个比喻：自动取景器

想象你在拍证件照：人太小就放大一点，歪了就转一下，不在中间就挪一挪。STN 干的正是这几件事——裁剪、缩放、旋转、平移，统称空间变换。它把输入图调整到「最适合分类」的样子，再交给后面的网络。

关键点是：这些变换参数不是人设定的，而是网络根据每张图自己算出来的。给一张歪七扭八的数字「9」，它先把图摆正再分类，准确率自然更高。

三个组件

STN 拆开是三个模块，串成一条流水线：

1. 定位网络（Localization Network）：一个小 CNN，看一眼输入图，输出一组变换参数。
2. 网格生成器（Grid Generator）：根据参数算出「输出图的每个像素，要从输入图的哪个位置取」。
3. 采样器（Sampler）：按算好的位置去输入图上取值，产出变换后的图。

在 PyTorch 里，第 2、3 步有现成函数：F.affine_grid 生成采样网格，F.grid_sample 完成采样。中间的可微性框架都帮你处理好了，梯度能一路传回定位网络。

仿射变换与 2×3 矩阵

STN 用的是仿射变换（Affine Transformation），覆盖旋转、缩放、平移、剪切，但不包括鱼眼镜头那种扭曲。平面上的仿射变换用 2×3 矩阵表示，共 6 个数：

[ a  b  c ]
[ d  e  f ]

a、e 管缩放，b、d 管旋转和剪切，c、f 管平移。定位网络就负责回归出这 6 个数。几个常见变换的矩阵长这样：平移 tx、ty 是 [1 0 tx; 0 1 ty]，等比缩放 s 倍是 [s 0 0; 0 s 0]，旋转角度 θ 是 [cosθ -sinθ 0; sinθ cosθ 0]。特别地，[1 0 0; 0 1 0] 表示「什么也不做」的恒等变换——这是初始化时的理想起点。

搭一个带 STN 的网络

把 STN 接在普通 CNN 前面：MNIST 数字先进 STN「摆正」，再走常规卷积。定位网络是两个卷积加两个全连接：

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

class Net(nn.Module):
    def __init__(self):
        super().__init__()
        # 分类用的常规网络
        self.conv1 = nn.Conv2d(1, 10, kernel_size=5)
        self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
        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),
        )
        # 回归 6 个仿射参数
        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))

    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)           # (batch, 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(x), 2))
        x = x.view(-1, 320)
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        return F.log_softmax(x, dim=1)

两个细节值得说。

stn 方法就三步：定位网络吐参数，affine_grid 生成网格，grid_sample 采样，全程可导。

采样器有个细节：变换后的像素位置不一定正好落在输入图的整数坐标上，比如要取的坐标是 (3.4, 5.7)。grid_sample 默认用双线性插值（Bilinear Interpolation），取周围四个像素加权平均，而且这个操作可导，梯度能顺畅流回定位网络。这正是 STN「可学习」的根基——整个变换链路处处可微。

恒等初始化很关键。fc_loc 最后一层的权重清零、bias 设成 [1,0,0,0,1,0]，保证训练刚开始时 STN 输出的是原图。要是随机初始化，网络一上来就把图乱扭，训练会很不稳定。这是论文里的经典技巧。

在 MNIST 上跑

数据还是老朋友 MNIST：单通道 28×28 的手写数字。

import torchvision
from torchvision import datasets, transforms

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

train_loader = torch.utils.data.DataLoader(
    datasets.MNIST(root='./data', train=True, download=True,
                   transform=transforms.Compose([
                       transforms.ToTensor(),
                       transforms.Normalize((0.1307,), (0.3081,)),
                   ])),
    batch_size=64, shuffle=True)

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

训练和普通分类没什么两样——这是 STN 最妙的地方：定位网络不用单独训练，跟着分类任务端到端一起学。模型最后输出的是 log_softmax，所以损失用 F.nll_loss：

import torch.optim as optim

model = Net().to(device)
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(f"Epoch {epoch} [{batch_idx * len(data)}"
                  f"/{len(train_loader.dataset)}] loss: {loss.item():.6f}")

def test():
    model.eval()
    correct = 0
    with torch.no_grad():
        for data, target in test_loader:
            data, target = data.to(device), target.to(device)
            output = model(data)
            pred = output.max(1, keepdim=True)[1]
            correct += pred.eq(target.view_as(pred)).sum().item()
    print(f"准确率: {correct}/{len(test_loader.dataset)} "
          f"({100. * correct / len(test_loader.dataset):.0f}%)")

for epoch in range(1, 6):
    train(epoch)
    test()

跑 20 个 epoch，测试准确率能到 99% 上下（官方教程的数据）。不过别激动：不装 STN 的普通 CNN 在 MNIST 上也能到 99% 附近。MNIST 的数字本来就摆得比较正，STN 的优势显不出来。它的用武之地是街景门牌号（SVHN）、自然场景文字这类歪来扭去的数据——物体位置、角度千变万化，数据增强管不过来，让网络自己学摆正更省心。

看看它学到了什么

STN 是可以可视化的：把输入图和 STN 变换后的图并排画出来，能直接看到网络学会了什么。

import torchvision

model.eval()
with torch.no_grad():
    data = next(iter(test_loader))[0].to(device)
    out = model.stn(data).cpu()

grid_in = torchvision.utils.make_grid(data.cpu())
grid_out = torchvision.utils.make_grid(out)
print("变换前像素均值:", round(grid_in.mean().item(), 3),
      " 变换后像素均值:", round(grid_out.mean().item(), 3))

训练初期，变换后的图还和原图差不多；训练后期，你会看到网络学会了把数字放大、居中、摆正。用 matplotlib 的 imshow 把两张网格图并列显示，效果更直观。这也验证了前面的说法：变换参数是网络自己从任务里学出来的，没人手把手教它。

小结

STN 把一个「可学习的空间变换」插进 CNN，让网络自己决定怎么摆正输入。三个组件各司其职：定位网络算参数、网格生成器算采样位置、采样器取值。实现上就靠 affine_grid 加 grid_sample 两个函数，外加恒等初始化的技巧。它几乎不增加多少参数，却能提升模型对几何变形的鲁棒性，而且能塞进任何现成 CNN 里，改动量很小。

到此为止，我们处理的对象都是一整张图。现实里的图片往往不止一个物体：马路上有行人、汽车、红绿灯，每个都要框出来。下一章，我们把视野从「整张图一个标签」扩大到「图上多个物体分别打框」——目标检测。