torch学习 (三十九)：空间Transformer网络_mnist transformer

1 引入

  基于MNIST数据集的空间Transformer网络示意。
  参考文献：pytorch中文版。

2 代码

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 numpy as np
import matplotlib.pyplot as plt

# 设置设备
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# 载入训练集，可以把num_workers调大一点，反正我是出错了，无奈之下智能0
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=0)
# 载入测试集
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=0)


# 空间transformer网络，包含定位网格、网格生成器、采样器
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)
        )
        # 3 * 2仿射矩阵回归器
        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 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)
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)
            # sum up batch loss
            test_loss += F.nll_loss(output, target, size_average=False).item()
            # get the index of the max log-probability
            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)))


def convert_image_np(inp):
    """转换tensor为numpy图像"""
    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

# 可视化输入图像和相应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()
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154

3 结果

3.1 分类

Train Epoch: 1 [0/60000 (0%)]	Loss: 2.340652
Train Epoch: 1 [32000/60000 (53%)]	Loss: 1.166227

Test set: Average loss: 0.2037, Accuracy: 9438/10000 (94%)

Train Epoch: 2 [0/60000 (0%)]	Loss: 0.349879
Train Epoch: 2 [32000/60000 (53%)]	Loss: 0.405450

Test set: Average loss: 1.3985, Accuracy: 6490/10000 (65%)

Train Epoch: 3 [0/60000 (0%)]	Loss: 2.455684
Train Epoch: 3 [32000/60000 (53%)]	Loss: 0.238986

Test set: Average loss: 0.0862, Accuracy: 9749/10000 (97%)

Train Epoch: 4 [0/60000 (0%)]	Loss: 0.117145
Train Epoch: 4 [32000/60000 (53%)]	Loss: 0.158333

Test set: Average loss: 0.0740, Accuracy: 9767/10000 (98%)

Train Epoch: 5 [0/60000 (0%)]	Loss: 0.146124
Train Epoch: 5 [32000/60000 (53%)]	Loss: 0.277766

Test set: Average loss: 0.0631, Accuracy: 9804/10000 (98%)

Train Epoch: 6 [0/60000 (0%)]	Loss: 0.095257
Train Epoch: 6 [32000/60000 (53%)]	Loss: 0.145955

Test set: Average loss: 0.0741, Accuracy: 9785/10000 (98%)

Train Epoch: 7 [0/60000 (0%)]	Loss: 0.130261
Train Epoch: 7 [32000/60000 (53%)]	Loss: 0.151724

Test set: Average loss: 0.0743, Accuracy: 9781/10000 (98%)

Train Epoch: 8 [0/60000 (0%)]	Loss: 0.183105
Train Epoch: 8 [32000/60000 (53%)]	Loss: 0.089648

Test set: Average loss: 0.0695, Accuracy: 9779/10000 (98%)

Train Epoch: 9 [0/60000 (0%)]	Loss: 0.079548
Train Epoch: 9 [32000/60000 (53%)]	Loss: 0.108383

Test set: Average loss: 0.0479, Accuracy: 9854/10000 (99%)

Train Epoch: 10 [0/60000 (0%)]	Loss: 0.242064
Train Epoch: 10 [32000/60000 (53%)]	Loss: 0.055248

Test set: Average loss: 0.0575, Accuracy: 9827/10000 (98%)

Train Epoch: 11 [0/60000 (0%)]	Loss: 0.439416
Train Epoch: 11 [32000/60000 (53%)]	Loss: 0.312931

Test set: Average loss: 0.0443, Accuracy: 9855/10000 (99%)

Train Epoch: 12 [0/60000 (0%)]	Loss: 0.245579
Train Epoch: 12 [32000/60000 (53%)]	Loss: 0.081791

Test set: Average loss: 0.0692, Accuracy: 9801/10000 (98%)

Train Epoch: 13 [0/60000 (0%)]	Loss: 0.215368
Train Epoch: 13 [32000/60000 (53%)]	Loss: 0.481491

Test set: Average loss: 0.0438, Accuracy: 9865/10000 (99%)

Train Epoch: 14 [0/60000 (0%)]	Loss: 0.067364
Train Epoch: 14 [32000/60000 (53%)]	Loss: 0.153307

Test set: Average loss: 0.0829, Accuracy: 9749/10000 (97%)

Train Epoch: 15 [0/60000 (0%)]	Loss: 0.131934
Train Epoch: 15 [32000/60000 (53%)]	Loss: 0.050842

Test set: Average loss: 0.0469, Accuracy: 9859/10000 (99%)

Train Epoch: 16 [0/60000 (0%)]	Loss: 0.327519
Train Epoch: 16 [32000/60000 (53%)]	Loss: 0.077101

Test set: Average loss: 0.0396, Accuracy: 9870/10000 (99%)

Train Epoch: 17 [0/60000 (0%)]	Loss: 0.134313
Train Epoch: 17 [32000/60000 (53%)]	Loss: 0.076286

Test set: Average loss: 0.0423, Accuracy: 9875/10000 (99%)

Train Epoch: 18 [0/60000 (0%)]	Loss: 0.078192
Train Epoch: 18 [32000/60000 (53%)]	Loss: 0.207639

Test set: Average loss: 0.0410, Accuracy: 9876/10000 (99%)

Train Epoch: 19 [0/60000 (0%)]	Loss: 0.108308
Train Epoch: 19 [32000/60000 (53%)]	Loss: 0.045296

Test set: Average loss: 0.0493, Accuracy: 9852/10000 (99%)

Train Epoch: 20 [0/60000 (0%)]	Loss: 0.106333
Train Epoch: 20 [32000/60000 (53%)]	Loss: 0.022837

Test set: Average loss: 0.0377, Accuracy: 9884/10000 (99%)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99

3.2 可视化