PyTorch CNN 训练后仅返回一个结果

Question

我正在训练 CNN 图像分类器。该网络将 255 x 255 RGB 图像分为五个类别，编号为

0

到

4

。

但是网络在训练过程中表现得很奇怪。尽管损失函数平滑下降，但模型在大多数情况下都会为批次中的所有样本返回相同的答案。更奇怪的是，最终它开始只回答 2。

这是一个典型的训练输出，包含 10 张图像。

LABELS                                 OUTPUT                                 CORRECT
tensor([2, 0, 2, 2, 2, 0, 2, 2, 2, 4]) tensor([0, 0, 0, 0, 0, 0, 0, 0, 0, 0]) 2 / 10
tensor([2, 2, 2, 2, 3, 4, 1, 2, 2, 2]) tensor([0, 0, 0, 0, 0, 0, 0, 0, 0, 0]) 0 / 10
tensor([2, 2, 2, 0, 2, 4, 3, 1, 2, 2]) tensor([0, 0, 0, 0, 0, 0, 0, 0, 0, 0]) 1 / 10
tensor([3, 4, 2, 2, 0, 4, 4, 3, 2, 0]) tensor([0, 0, 0, 0, 0, 0, 0, 0, 0, 0]) 2 / 10
tensor([1, 2, 2, 4, 2, 0, 1, 0, 0, 0]) tensor([0, 0, 0, 0, 0, 0, 0, 0, 0, 0]) 4 / 10
tensor([2, 2, 2, 3, 2, 0, 0, 1, 2, 2]) tensor([0, 0, 0, 0, 0, 0, 0, 0, 0, 0]) 2 / 10
tensor([1, 1, 0, 1, 2, 2, 1, 1, 0, 1]) tensor([0, 0, 0, 0, 0, 0, 0, 0, 0, 0]) 2 / 10
tensor([0, 2, 1, 3, 3, 2, 1, 0, 2, 2]) tensor([0, 0, 0, 0, 0, 0, 0, 0, 0, 0]) 2 / 10
tensor([2, 3, 2, 2, 3, 1, 0, 1, 0, 2]) tensor([0, 0, 0, 0, 0, 0, 0, 0, 0, 0]) 2 / 10
tensor([3, 2, 3, 1, 1, 2, 0, 4, 2, 2]) tensor([0, 0, 0, 0, 0, 0, 0, 0, 0, 0]) 1 / 10
tensor([2, 1, 0, 3, 1, 2, 2, 1, 2, 0]) tensor([2, 2, 2, 2, 2, 0, 2, 2, 0, 2]) 2 / 10
tensor([3, 0, 2, 1, 3, 1, 2, 4, 2, 2]) tensor([2, 2, 2, 2, 2, 2, 2, 2, 2, 2]) 4 / 10
tensor([2, 2, 1, 2, 1, 1, 1, 4, 3, 2]) tensor([2, 2, 2, 2, 2, 2, 2, 2, 2, 2]) 4 / 10

# Remaining predictions are always [2, 2, 2...]
# Loss function is not shown, but it declines smoothly and looks well behaved

虽然

2

是标签中最常见的类别（大约 50% 的图像），但我不明白为什么 CNN 应该“集中”在单个答案上（上面示例中的

0

）或始终预测

2

最后

即使精度不够好，我也希望在输出张量中获得更多不同的结果。我做错了什么？

这是我的网络代码...

class CNN(nn.Module):
    def __init__(self, n_layers=3, n_categories=5):
        super(CNN, self).__init__()
        self.conv1 = nn.Conv2d(n_layers, 6, 5)
        self.pool = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(6, 16, 5)
        self.conv3 = nn.Conv2d(16, 16, 5)
        self.fc1 = nn.Linear(16 * 28 * 28, 200)
        self.fc2 = nn.Linear(200, 84)
        self.fc3 = nn.Linear(84, n_categories)

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))        
        x = self.pool(F.relu(self.conv2(x)))
        x = self.pool(F.relu(self.conv3(x)))
        x = x.view(-1, 16 * 28 * 28)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

...优化器、损失函数和数据加载器...

model = CNN()

transforms = v2.Compose([
    v2.ToImageTensor(),
    v2.ConvertImageDtype(),
    v2.Resize((256, 256), antialias=True)
])

dataset = UBCDataset(transforms=transforms)
full_dataloader = DataLoader(dataset, batch_size=10, shuffle=False)

loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.001, momentum=0.9)

...以及产生上述输出的训练循环。损失函数没有显示出来，但正如预期的那样平稳下降。

batches = iter(full_dataloader)

print("LABELS                                 OUTPUT                                 CORRECT")
for X, y in batches:   
    model.train()    
    
    pred = model(X)
    loss = loss_fn(pred, y)
    
    loss.backward()
    optimizer.step()
    #optimizer.zero_grad()
    
    print(f"{y} {pred.argmax(1)} {int(sum(y == pred.argmax(1)))} / {len(y)} {loss.item()}")

欢迎任何意见。

Answer 1

事实证明问题出在我的损失函数和数据加载器中。我的数据不平衡并且没有正确标准化。

输入数据标准化并将权重添加到

CrossEntropyLoss

损失函数后，结果显着改善。