12.13 深度学习-空间注意力仿射变换STN
# 就是对图像做一个仿射 变换 仿射变换的参数要去学习 直接用nn.Linear()就行
# 仿射变换的原理就是 线性变换
# 这个仿射变换要加在模型最开始的地方
import torch
import torch.nn as nn
import torch.nn.functional as F
class STN(nn.Module):
def __init__(self, c, h, w):
super(STN, self).__init__()
# 定义线性变换 要输出 1,6的矩阵 然后变为2,3的 然后 这个2,3的矩阵 通过affine_grid变为一个 仿射变换矩阵 然后 通过grid_sample 把 数据 和仿射变换矩阵 完成仿射变换
# 两个线性层是根据图来的
self.fc = nn.Sequential(
nn.Linear(in_features=c * h * w, out_features=32),
nn.Tanh(),
nn.Linear(in_features=32, out_features=6),
nn.Tanh(),
)
def forward(self, x):
batch_size, c, h, w = x.size()
# 2行3列的矩阵
theta = self.fc(x.view(batch_size, -1)).view(batch_size, 2, 3)
# 仿射变换矩阵
grid = F.affine_grid(
theta, torch.Size((batch_size, c, h, w)), align_corners=False
)
# 开始采样
sample = F.grid_sample(x, grid, align_corners=False)
return sample