文章目錄

• 數據集的加載
• 空間變換網絡的介紹
• 定義網絡
• 訓練和測試模型
• 可視化 STN 結果

官方文檔地址： https://pytorch.org/tutorials/intermediate/spatial_transformer_tutorial.html

在本教程中，您將學會如何使用 空間變換網絡 的視覺注意力機制來擴充網絡。如果需要了解更多 空間變換網絡 可以在 DeepMind 論文中詳細閱讀。

空間變換網絡允許神經網絡學習如何對輸入圖像執行空間變換，以增強模型的幾何不變性。例如，他可以裁剪感興趣的區域，縮放并校正圖像的方向。這可能是一個有用的機制，因為CNN不會對旋轉和縮放以及更一般的仿射變換保持不變。

對于 STN 的最大優點之一就是：能夠將其簡單地插入到任何現有的 CNN 中，而無需進行任何修改。

# License: BSD # Author: Ghassen Hamrounifrom __future__ import print_function 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() # interactive mode

數據集的加載

本文以經典的 MNIST 數據集為例，使用標準卷積網絡和空間變換網絡。

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")# Training dataset 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 dataset 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)

空間變換網絡的介紹

空間變換網絡主要由三個主要部分組成：

• 本地網絡（Localisation Network）： 本地網絡為常規的 CNN，是一個用來回歸變換參數θ的網絡。
• 網格生成器（Grid Genator）： 網格生成器在輸入圖像中生成與輸出圖像的每個像素相對應的坐標網絡。
• 采樣器（Sampler）： 采樣器利用采樣網絡和輸入的特征圖同時作為輸入，然后輸入，得到了特征圖經過變換之后的結果。

定義網絡

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)# Spatial transformer localization-networkself.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))# Regressor for the 3 * 2 affine matrixself.fc_loc = nn.Sequential(nn.Linear(10 * 3 * 3, 32),nn.ReLU(True),nn.Linear(32, 3 * 2))# Initialize the weights/bias with identity transformationself.fc_loc[2].weight.data.zero_()self.fc_loc[2].bias.data.copy_(torch.tensor([1, 0, 0, 0, 1, 0], dtype=torch.float))# Spatial transformer network forward functiondef 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):# transform the inputx = self.stn(x)# Perform the usual forward passx = 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 算法訓練模型。網絡正在以監督學習的方式來學習分類任務。同時該模型以端到端的方式自動學習 STN。

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())) # # A simple test procedure to measure STN the performances on MNIST. #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)# sum up batch losstest_loss += F.nll_loss(output, target, size_average=False).item()# get the index of the max log-probabilitypred = 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)))

可視化 STN 結果

現在，我們將檢查學習到的視覺注意力機制的結果。

我們定義了一個小的輔助函數，以便訓練的時候進行可視化轉換。

def convert_image_np(inp):"""Convert a Tensor to numpy image."""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# We want to visualize the output of the spatial transformers layer # after the training, we visualize a batch of input images and # the corresponding transformed batch using STN.def visualize_stn():with torch.no_grad():# Get a batch of training datadata = 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))# Plot the results side-by-sidef, 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 the STN transformation on some input batch visualize_stn()plt.ioff() plt.show()

輸出：

Train Epoch: 18 [0/60000 (0%)] Loss: 0.197388 Train Epoch: 18 [32000/60000 (53%)] Loss: 0.070883Test set: Average loss: 0.0386, Accuracy: 9885/10000 (99%)Train Epoch: 19 [0/60000 (0%)] Loss: 0.127624 Train Epoch: 19 [32000/60000 (53%)] Loss: 0.078326Test set: Average loss: 0.0345, Accuracy: 9898/10000 (99%)Train Epoch: 20 [0/60000 (0%)] Loss: 0.035128 Train Epoch: 20 [32000/60000 (53%)] Loss: 0.028159Test set: Average loss: 0.0361, Accuracy: 9892/10000 (99%)

總結

以上是生活随笔為你收集整理的PyTorch 进阶学习（二）————STN:空间变换网络（Spatial Transformer Network）的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。