모든 준비가 완료되었기 때문에 학습 데이터를 이용하여 모델을 학습시킵니다. 이 부분은 앞에서 사용했던 코드 5-8과 동일합니다.
코드 5-11 모델 학습 및 성능 평가
num_epochs = 5
count = 0
loss_list = []
iteration_list = []
accuracy_list = []
predictions_list = []
labels_list = []
for epoch in range(num_epochs):
for images, labels in train_loader:
images, labels = images.to(device), labels.to(device)
train = Variable(images.view(100, 1, 28, 28))
labels = Variable(labels)
outputs = model(train)
loss = criterion(outputs, labels)
optimizer.zero_grad()
loss.backward()
optimizer.step()
count += 1
if not (count % 50):
total = 0
correct = 0
for images, labels in test_loader:
images, labels = images.to(device), labels.to(device)
labels_list.append(labels)
test = Variable(images.view(100, 1, 28, 28))
outputs = model(test)
predictions = torch.max(outputs, 1)[1].to(device)
predictions_list.append(predictions)
correct += (predictions == labels).sum()
total += len(labels)
accuracy = correct * 100 / total
loss_list.append(loss.data)
iteration_list.append(count)
accuracy_list.append(accuracy)
if not (count % 500):
print("Iteration: {}, Loss: {}, Accuracy: {}%".format(count, loss.data, accuracy))