Проблема в том, что вы не применяете свои полностью подключенные слои fc1
и fc2
.
Ваш forward()
в настоящее время выглядит так:
def forward(self, x):
# x -> (batch_size, 784)
x = torch.relu(x)
# x -> (batch_size, 10)
x = torch.softmax(x, dim=1)
return x
Так что если вы измените его на:
def forward(self, x):
# x -> (batch_size, 784)
x = self.fc1(x) # added layer fc1
x = torch.relu(x)
# x -> (batch_size, 10)
x = self.fc2(x) # added layer fc2
x = torch.softmax(x, dim=1)
return x
Это должно работать.
Относительно ответа Уманга Гуптаса: На мой взгляд, вызов zero_grad()
перед вызовом backward()
, как это сделал мистер Робот, - это нормально. Это не должно быть проблемой.
Edit:
Итак, я сделал короткий тест - я установил итерации от 1000
до 10000
, чтобы увидеть более широкую картину, если она действительно уменьшается. (Конечно, я также загрузил данные в mnist_m
, поскольку они не были включены в опубликованный вами код)
Я добавил условие печати в код:
if epoch % 1000 == 0:
print('Epoch', epoch, '- Loss:', round(loss.item(), 3))
Что выводит потери каждые 1000
итераций:
Epoch 0 - Loss: 2.305
Epoch 1000 - Loss: 2.263
Epoch 2000 - Loss: 2.187
Epoch 3000 - Loss: 2.024
Epoch 4000 - Loss: 1.819
Epoch 5000 - Loss: 1.699
Epoch 6000 - Loss: 1.699
Epoch 7000 - Loss: 1.656
Epoch 8000 - Loss: 1.675
Epoch 9000 - Loss: 1.659
Протестировано с версией PyTorch 0.4.1
Итак, вы можете видеть, что с измененным forward()
сеть теперь учится, остальной код я оставил нетронутым.
Удачи дальше!