可视化 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 + mean

inp = 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 data

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

# Plot the results side-by-side

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 the STN transformation on some input batch

visualize_stn()

plt.ioff()

plt.show()

/usr/local/lib/python3.10/dist-packages/torch/nn/functional.py:5167: UserWarning:

Default grid_sample and affine_grid behavior has changed to align_corners=False since 1.3.0. Please specify align_corners=True if the old behavior is desired. See the documentation of grid_sample for details.

/usr/local/lib/python3.10/dist-packages/torch/nn/functional.py:5100: UserWarning:

Default grid_sample and affine_grid behavior has changed to align_corners=False since 1.3.0. Please specify align_corners=True if the old behavior is desired. See the documentation of grid_sample for details.

Train Epoch: 1 [0/60000 (0%)] Loss: 2.321845

Train Epoch: 1 [32000/60000 (53%)] Loss: 0.739860

/usr/local/lib/python3.10/dist-packages/torch/nn/_reduction.py:51: UserWarning:

size_average and reduce args will be deprecated, please use reduction='sum' instead.

Test set: Average loss: 0.2149, Accuracy: 9381/10000 (94%)

Train Epoch: 2 [0/60000 (0%)] Loss: 0.310411

Train Epoch: 2 [32000/60000 (53%)] Loss: 0.443614

Test set: Average loss: 0.1403, Accuracy: 9581/10000 (96%)

Train Epoch: 3 [0/60000 (0%)] Loss: 0.184066

Train Epoch: 3 [32000/60000 (53%)] Loss: 0.237139

Test set: Average loss: 0.0851, Accuracy: 9721/10000 (97%)

Train Epoch: 4 [0/60000 (0%)] Loss: 0.235744

Train Epoch: 4 [32000/60000 (53%)] Loss: 0.374082

Test set: Average loss: 0.0717, Accuracy: 9783/10000 (98%)

Train Epoch: 5 [0/60000 (0%)] Loss: 0.209939

Train Epoch: 5 [32000/60000 (53%)] Loss: 0.106230

Test set: Average loss: 0.0677, Accuracy: 9797/10000 (98%)

Train Epoch: 6 [0/60000 (0%)] Loss: 0.210056

Train Epoch: 6 [32000/60000 (53%)] Loss: 0.252697

Test set: Average loss: 0.0668, Accuracy: 9783/10000 (98%)

Train Epoch: 7 [0/60000 (0%)] Loss: 0.123260

Train Epoch: 7 [32000/60000 (53%)] Loss: 0.127563

Test set: Average loss: 0.0553, Accuracy: 9835/10000 (98%)

Train Epoch: 8 [0/60000 (0%)] Loss: 0.108394

Train Epoch: 8 [32000/60000 (53%)] Loss: 0.029457

Test set: Average loss: 0.0483, Accuracy: 9840/10000 (98%)

Train Epoch: 9 [0/60000 (0%)] Loss: 0.176748

Train Epoch: 9 [32000/60000 (53%)] Loss: 0.127248

Test set: Average loss: 0.1173, Accuracy: 9636/10000 (96%)

Train Epoch: 10 [0/60000 (0%)] Loss: 0.243940

Train Epoch: 10 [32000/60000 (53%)] Loss: 0.156468

Test set: Average loss: 0.2388, Accuracy: 9345/10000 (93%)

Train Epoch: 11 [0/60000 (0%)] Loss: 0.378354

Train Epoch: 11 [32000/60000 (53%)] Loss: 0.059640

Test set: Average loss: 0.0448, Accuracy: 9862/10000 (99%)

Train Epoch: 12 [0/60000 (0%)] Loss: 0.129616

Train Epoch: 12 [32000/60000 (53%)] Loss: 0.052506

Test set: Average loss: 0.0436, Accuracy: 9868/10000 (99%)

Train Epoch: 13 [0/60000 (0%)] Loss: 0.045932

Train Epoch: 13 [32000/60000 (53%)] Loss: 0.079384

Test set: Average loss: 0.0446, Accuracy: 9858/10000 (99%)

Train Epoch: 14 [0/60000 (0%)] Loss: 0.031097

Train Epoch: 14 [32000/60000 (53%)] Loss: 0.106284

Test set: Average loss: 0.0422, Accuracy: 9874/10000 (99%)

Train Epoch: 15 [0/60000 (0%)] Loss: 0.133345

Train Epoch: 15 [32000/60000 (53%)] Loss: 0.106248

Test set: Average loss: 0.0414, Accuracy: 9873/10000 (99%)

Train Epoch: 16 [0/60000 (0%)] Loss: 0.044279

Train Epoch: 16 [32000/60000 (53%)] Loss: 0.046603

Test set: Average loss: 0.0626, Accuracy: 9806/10000 (98%)

Train Epoch: 17 [0/60000 (0%)] Loss: 0.157027

Train Epoch: 17 [32000/60000 (53%)] Loss: 0.127502

Test set: Average loss: 0.0497, Accuracy: 9851/10000 (99%)

Train Epoch: 18 [0/60000 (0%)] Loss: 0.041967

Train Epoch: 18 [32000/60000 (53%)] Loss: 0.137598

Test set: Average loss: 0.0397, Accuracy: 9887/10000 (99%)

Train Epoch: 19 [0/60000 (0%)] Loss: 0.115559

Train Epoch: 19 [32000/60000 (53%)] Loss: 0.034766

Test set: Average loss: 0.0362, Accuracy: 9894/10000 (99%)

Train Epoch: 20 [0/60000 (0%)] Loss: 0.078510

Train Epoch: 20 [32000/60000 (53%)] Loss: 0.096980

Test set: Average loss: 0.0715, Accuracy: 9794/10000 (98%)

脚本总运行时间： (1 分钟 36.699 秒)

下载 Jupyter Notebook： spatial_transformer_tutorial.ipynb

下载 Python 源 代码： spatial_transformer_tutorial.py

下载 zip 文件： spatial_transformer_tutorial.zip