这篇博客瞄准的是 pytorch 官方教程中 Image and Video 章节的 Spatial Transformer Networks Tutorial 部分。
- 官网链接：https://pytorch.org/tutorials/intermediate/spatial_transformer_tutorial.html
完整网盘链接: https://pan.baidu.com/s/1L9PVZ-KRDGVER-AJnXOvlQ?pwd=aa2m 提取码: aa2m
Spatial Transformer Networks Tutorial
在本教程将介绍如何使用 Spatial Transformer Network 的视觉注意机制来增强模型。可以在 DeepMind 论文 中信息。
Spatial Transformer Network, STN 是对任何可微空间注意力概括，允许神经网络学习如何对输入图像进行空间变换，以增强模型的几何不变性，例如裁剪ROI区域、缩放、校正图像的方向。
STN 的最大优点之一是只需进行很少的改动就可以将其插入任何现有的 CNN 中。
导入必要的包
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import torchvision
from torchvision import datasets, transforms
import matplotlib.pyplot as plt
import numpy as npplt.ion()
Loading the data
from six.moves import urllib
opener = urllib.request.build_opener()
opener.addheaders = [('User-agent', 'Mozilla/5.0')]
urllib.request.install_opener(opener)
计算加速设备
device = torch.accelerator.current_accelerator().type if torch.accelerator.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, num_workers=4
)test_loader = torch.utils.data.DataLoader(datasets.MNIST(root='./data',train=False,download=True,transform=transforms.Compose([transforms.ToTensor(),transforms.Normalize((0.1307,),(0.3081))])),batch_size=64, shuffle=True, num_workers=4
)
Depicting spatial transformer networks
STN 由三个主要组成部分：
- 定位网络是一个常规的 CNN，负责将转换的参数进行回归。这部分的模型参数（即模型所拥有的转换能力）不从当前数据集中训练得到，而是从整体模型中以提高全局准确性的训练目标上学习；
- 网格生成器将输入图像中每个像素对应到输出图像中的网格上；
- 采样器使用变换参数并将其作用在输入图像上；
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)# 定位模块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))# 回归模块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)grid = F.affine_grid(theta, x.size())x = F.grid_sample(x, grid)return xdef 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)
Training the model
模型整体以监督的方式学习分类任务。同时以端到端的方式自动学习 STN 部分。
定义训练函数
def train(epochs):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 = 0correct = 0for data, target in test_loader:data, target = data.to(device), target.to(device)output = model(data)test_loss += F.nll_loss(output, target, size_average=False).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)))
Visualizing the STN results
定义一个辅助函数来查看 STN 的训练结果。
将图像转为numpy数组：
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 + meaninp = np.clip(inp, 0, 1)return inp
可视化stn
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('Dataset Images')axarr[1].imshow(out_grid)axarr[1].set_title('Transformed Images')
训练
for epoch in range(1, 20+1):train(epoch)test()
可视化
visualize_stn()
plt.ioff()
plt.show()
