如何改进空间变换网络(STN)的旋转校正效果

• 任务背景：在自建靴子、鞋类样本数据集上应用空间变换网络（Spatial Transformation Network, STN），所有样本存在10°~30°范围的随机轻微旋转。模型基于FashionMNIST数据集完成训练与测试，预期STN模块可输出对齐后的图像，但实际变换结果未实现旋转对齐。
• STN变换效果：
  
• 训练过程指标变化：
  
• 现有实现代码：

class STN_CNN(nn.Module):
def __init__(self):
    super(STN_CNN, self).__init__()
    self.cnn = nn.Sequential(
        nn.Conv2d(1, 10, kernel_size=3, stride=1, padding=0),
        nn.MaxPool2d(2, stride=2),
        nn.ReLU(),
        nn.Conv2d(10, 16, kernel_size=3, stride=1, padding=0),
        nn.MaxPool2d(2, stride=2),
        nn.ReLU()
    )
    self.classifier = nn.Sequential(
        nn.Linear(16*2*2, 32),
        nn.ReLU(),
        nn.Linear(32, 10)
    )
    self.localization = nn.Sequential(
        nn.Conv2d(1, 20, kernel_size=5, stride=1, padding=0),
        nn.MaxPool2d(2, stride=2),
        nn.ReLU(),
        nn.Conv2d(20, 20, kernel_size=5, stride=1, padding=0),
        nn.ReLU()
    )
    self.fc_loc = nn.Sequential(
        nn.Linear(20*8*8, 20),
        nn.ReLU(),
        nn.Linear(20, 6)
    )
    self.AvgPool = nn.AvgPool2d(2, stride=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):
    x_loc = self.localization(x)
    x_loc = x_loc.view(-1, 20*8*8)
    theta = self.fc_loc(x_loc)
    theta = theta.view(-1, 2, 3)

    grid = F.affine_grid(theta, x.size())
    x = F.grid_sample(x, grid)
    x = self.AvgPool(x)

    return x

def forward(self, x):
    x = self.stn(x)
    x = self.cnn(x)
    x = x.view(-1, 16*2*2)
    x = self.classifier(x)

    return x

• 训练现状：模型累计训练100轮，旋转校正效果始终无提升，需要排查代码错误、给出STN旋转校正能力的优化方案。

显性代码错误 #

• 特征维度完全不匹配：输入是2828的FashionMNIST图像，STN分支经过grid_sample输出后你直接加了步长为2的AvgPool，输出尺寸变成1414，但后续CNN分支是针对2828输入设计的：两层卷积加池化后输出尺寸是55，对应特征维度应该是1655=400，和你写的16*2*2完全对不上。运行时靠view强行拉平特征，相当于把错位的特征喂给分类器，STN根本拿不到有效的分类梯度回传，自然学不会校正逻辑。
  修复方式：直接删掉STN分支里的AvgPool，重新计算CNN每层的输出维度，确保view操作的输入维度和卷积实际输出维度完全一致。
• 仿射变换参数无约束：你直接让网络输出6维完整仿射参数，自由度太高。对于只有小角度旋转的场景，网络很容易优先学到缩放、平移、剪切这类更易快速降低分类损失的变换，不会主动学习旋转校正。
• 接口参数不匹配：调用F.affine_grid和F.grid_sample时没有显式指定align_corners参数，PyTorch不同版本默认值不一致，很容易出现坐标偏移，导致变换结果错位。

训练策略问题 #

• 学习率设置不合理：STN定位分支和主分类网络用相同学习率，参数更新步长过大，很容易直接破坏初始的恒等变换状态，训练过程震荡，学不到稳定的校正逻辑。
• 训练节奏错误：从训练一开始就同时更新STN和主分类网络参数，初期分类网络特征提取能力极差，回传给STN的梯度噪声极大，很容易把STN参数带偏，后期也很难收敛回正确的校正方向。

• 限制变换参数自由度：针对仅存在小角度旋转的场景，不需要输出完整6维仿射矩阵。可以让fc_loc只输出1维旋转角度值，手动构造旋转矩阵：

  theta = torch.zeros(x.size(0), 2, 3, device=x.device)
  angle = self.fc_loc(x_loc).squeeze() # 输出维度为(batch, 1)
  theta[:, 0, 0] = torch.cos(angle)
  theta[:, 0, 1] = -torch.sin(angle)
  theta[:, 1, 0] = torch.sin(angle)
  theta[:, 1, 1] = torch.cos(angle)
  

  从根源上避免网络学到无关的缩放、剪切变换，强迫STN聚焦旋转校正任务。如果需要兼顾微小平移，可以让fc_loc输出3个值：旋转角、x方向平移、y方向平移，固定缩放项为1即可。
• 调整训练节奏：前5~10轮冻结STN分支所有参数，只训练主CNN分类网络，等分类网络具备基础特征提取能力后，再放开STN分支的参数更新。
• 差异化学习率设置：STN定位分支的学习率设置为主分类网络的1/5~1/10，避免参数更新过猛破坏训练稳定性。
• 增强定位网络能力：现有定位网络感受野不足，可以把第一层卷积核尺寸从5调整为7，增加一层BatchNorm提升训练稳定性，全连接层加入0.2~0.3的Dropout避免过拟合，提升角度估计的鲁棒性。
• 可选辅助损失：如果有旋转角度的标注，可以在STN输出的角度值上加L1回归损失，直接监督STN学习旋转角度，收敛速度会比仅靠分类损失监督快数倍。

内容的提问来源于stack exchange，提问作者zakaria14