基于PyTorch的图像旋转角度检测模型小角度预测优化咨询

我正尝试构建一个基于PyTorch的图像旋转角度检测模型，数据集由1500份文档生成：先从[0, 90, -90, 180]中随机选取2个角度，再为每个角度添加random.uniform(-10, 10)的随机偏移，最终得到约4000张旋转图像。

我设计了一个预测目标角度正弦、余弦值的CNN模型，结构如下：

class CnnRotateRegression(nn.Module):
    def __init__(self):
        super(CnnRotateRegression, self).__init__()
        self.conv1 = nn.Conv2d(3, 64, kernel_size=3)
        self.conv2 = nn.Conv2d(64, 128, kernel_size=3)
        self.conv3 = nn.Conv2d(128, 256, kernel_size=3)
        self.conv4 = nn.Conv2d(256, 512, kernel_size=3)
        self.conv5 = nn.Conv2d(512, 512, kernel_size=3)
        
        self.bn1 = nn.BatchNorm2d(64)
        self.bn2 = nn.BatchNorm2d(128)
        self.bn3 = nn.BatchNorm2d(256)
        self.bn4 = nn.BatchNorm2d(512)
        self.bn5 = nn.BatchNorm2d(512)
        
        self.activation = nn.ReLU()
        self.pool = nn.AvgPool2d(kernel_size=2)
        self.pool2 = nn.AdaptiveAvgPool2d((8,8))

        self.linear_l1 = nn.Linear(512*8*8, 512)
        self.linear_l2 = nn.Linear(512, 256)
        self.linear_l3 = nn.Linear(256, 2) # sin + cos

    def forward(self, x):
        x = self.activation(self.pool(self.bn1(self.conv1(x))))
        x = self.activation(self.pool(self.bn2(self.conv2(x))))
        x = self.activation(self.pool(self.bn3(self.conv3(x))))
        x = self.activation(self.pool(self.bn4(self.conv4(x))))
        x = self.activation(self.pool(self.bn5(self.conv5(x))))
        
        x = self.pool2(x)
        x = x.view(x.size(0), -1)
        
        x = self.activation(self.linear_l1(x))
        x = self.activation(self.linear_l2(x))
        x = self.linear_l3(x)

        x = F.normalize(x, p=2, dim=1)

        return x

训练代码如下：

model = CnnRotateRegression()
model = model.to(device)

loss_function = nn.MSELoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)
num_of_epochs = 11


for epoch in range(num_of_epochs):
    model.train()
    running_loss = 0.0
    for images, labels in tqdm(train_Loader, desc="training loop"):

        images, labels = images.to(device), labels.to(device).float()

        angles = angle_to_sin_cos(labels)
        norm_angles = F.normalize(angles, p=2, dim=1)

        optimizer.zero_grad()
        outputs = model(images)
        loss = loss_function(outputs, norm_angles)
        loss.backward()
        optimizer.step()

        running_loss += loss.item()
    train_loss = running_loss / len(train_Loader)

角度与正弦、余弦值的转换函数：

def angle_to_sin_cos(angle):
    tensor_angle = angle.clone().detach()
    radian = tensor_angle * torch.pi / 180.0
    return torch.stack([torch.cos(radian), torch.sin(radian)], dim=1)

def sin_cos_to_angle(outputs):
    cos_val, sin_val = outputs[:, 0], outputs[:, 1]
    angle_rad = torch.atan2(sin_val, cos_val)
    angle_deg = angle_rad * (180 / torch.pi)
    return angle_deg

目前该模型在±10度的小角度范围内预测表现较差，请问有哪些改进方案可以提升小角度预测性能？

1. 调整数据集分布，强化小角度样本 #

当前数据集以[0,90,-90,180]核心角度为主，小角度仅作为偏移存在，模型对大角度特征的学习优先级远高于小角度。

• 提升小角度样本占比：直接生成[-10,10]区间内的随机角度样本，或修改生成策略，将原有的2个核心角度改为1个核心角度+更多小角度偏移样本，让小角度样本占比提升至30%-40%。
• 小角度专属数据增强：针对小角度图像添加二次随机旋转（如[-5,5]区间），强化模型对微小角度变化的感知能力。

2. 替换损失函数，放大小角度误差权重 #

MSE损失对所有角度的误差权重一致，但小角度对应的sin/cos值差异极小（比如0度和10度的cos值差异仅约0.015），模型难以感知这种细微变化。

• 使用角度损失：直接基于角度差计算损失，公式为loss = 1 - torch.cos(theta_pred_rad - theta_true_rad)，小角度的误差会被放大（10度差对应的loss约0.015，90度差对应的loss为1），迫使模型关注小角度预测误差。
• 加权MSE损失：为真实角度在±10度范围内的样本设置2-3倍的损失权重，让模型优先学习小角度特征。

3. 优化模型结构，保留细粒度特征 #

当前模型的5次AvgPool会丢失大量细粒度位置信息，而小角度旋转的特征变化恰恰体现在细微的位置偏移上。

• 减少池化次数：去掉1-2次AvgPool，或将部分AvgPool替换为Stride=2的卷积（降采样的同时保留更多特征）。
• 添加注意力机制：在模型中后层加入通道注意力（如SE模块）或空间注意力，引导模型关注文档边缘、文字排列等对小角度变化敏感的区域。
• 降低特征压缩程度：将AdaptiveAvgPool2d((8,8))调整为(16,16)，减少线性层的输入维度压缩，保留更多空间细节。

4. 调整训练策略，精细化小角度拟合 #

• 动态调整学习率：初始用0.001训练3-5轮，然后降至0.0001继续训练，让模型后期能精细调整小角度相关参数。
• 增加训练轮次：当前11轮训练不足以让模型充分学习小角度特征，可增加至20-30轮，同时加入早停机制防止过拟合。
• 小角度样本微调：常规训练完成后，用小角度样本集对模型进行1-2轮微调，固定卷积层参数，仅训练全连接层，强化模型对小角度的拟合能力。

5. 修正sin/cos转换细节 #

• 处理角度周期性：确保模型输出角度与真实角度处于同一周期区间（如[-180,180]），避免因周期问题导致的误差，可在sin_cos_to_angle函数中添加角度归一化逻辑。
• 移除冗余归一化：真实的sin/cos值本身就是单位向量，norm_angles = F.normalize(angles, p=2, dim=1)属于冗余操作，直接使用angles作为目标即可，避免不必要的计算误差。

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