論文題目：Spatial Transformer Networks

STN.png

• 定位網(wǎng)絡(luò)：一般是由一個(gè)全連接網(wǎng)絡(luò)實(shí)現(xiàn)，該網(wǎng)絡(luò)最終輸出6個(gè)數(shù)值，將其轉(zhuǎn)換為2x3矩陣，其中每?jī)蓚€(gè)為一組，分別表示仿射變換中的平移、旋轉(zhuǎn)、縮放所對(duì)應(yīng)的參數(shù)。
• 網(wǎng)格生成器：創(chuàng)建一個(gè)與輸出圖片大小相同的矩陣，把該矩陣與仿射參數(shù)矩陣相乘將得到的結(jié)果作為圖片中每個(gè)像素點(diǎn)對(duì)應(yīng)于原圖的坐標(biāo)。
• 采樣器：使用坐標(biāo)映射中每個(gè)像素點(diǎn)的坐標(biāo)值，在原始圖片上取相應(yīng)的像素，并將其填充到目標(biāo)圖片中，得到整幅目標(biāo)圖片。
  通俗一點(diǎn)講就是通過訓(xùn)練學(xué)習(xí)一個(gè)仿射變換矩陣，而不是人為定義這個(gè)矩陣，這個(gè)仿射變換矩陣可以將形變的圖像大致恢復(fù)回來(lái)。
  注意：STN更適合用于分類或文字識(shí)別這樣輸出結(jié)果跟圖像位置不太相關(guān)的任務(wù)，想目標(biāo)檢測(cè)或語(yǔ)義分割任務(wù)的改動(dòng)量就比較大，由于STN對(duì)輸入圖像做了形變所以標(biāo)簽也要做相應(yīng)變換的。
  代碼如下所示：

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-network
        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)
        )

        # Regressor for the 3 * 2 affine matrix
        self.fc_loc = nn.Sequential(
            nn.Linear(10 * 3 * 3, 32),
            nn.ReLU(True),
            nn.Linear(32, 3 * 2)
        )

        # Initialize the weights/bias with identity transformation
        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))

    # Spatial transformer network forward function
    def stn(self, x):
        xs = self.localization(x)
        # xs = xs.view(-1, 10 * 3 * 3)
        xs = torch.flatten(xs, 1)
        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):
        # transform the input
        x = self.stn(x)

        # Perform the usual forward pass
        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)

這是一個(gè)分類網(wǎng)絡(luò)中在開頭部分增加了STN模塊，用于消除或降低圖像形變的影響，從而降低后續(xù)模型的分類壓力