Выход модели имеет только одно измерение, он имеет размер [6] , но nn.CrossEntropyLoss
ожидает размер [batch_size, num_classes ] .
В модели вы сглаживаете вывод сверток. Вы должны сохранить размер партии, так как они независимы друг от друга, и их полное выравнивание объединит их в один. torch.flatten
принимает аргумент start_dim
(второй аргумент), который определяет, с какого измерения начинается сглаживание. Если установить для него значение 1, он начнется со второго измерения, а первое измерение (пакетное измерение) останется без изменений.
# Flatten everything but the first dimension
# From: [batch_size, channels, height, width] (4D)
# To: [batch_size, channels * height * width] (2D)
x = torch.flatten(nn.functional.relu(x), 1)
Выходные данные модели также должны иметь такое же количество классов, что и ваш набор данных. Поскольку у вас 120 классов, результат последнего линейного слоя должен быть 120.
self.fc3 = nn.Linear(200, 120)
Кроме того, метки должны находиться в диапазоне [0, 119] , потому что они индексы классов и, как и любая индексация в Python, отсчитываются от нуля. Если ваши метки находятся в диапазоне [1, 120] , вы можете просто вычесть из них единицу.