<조건 : nn_Module를 이용, 총 2000번의 에포크를 수행하고 200번마다 로그를 출력>
# 주어진 데이터
X = [[2.0, 1.0],
[3.0, 1.0],
[4.0, 1.0],
[3.0, 2.0],
[5.0, 1.0],
[7.0, 1.0],
[8.0, 2.0],
[9.0, 2.0],
[10.0, 3.0],
[11.0, 3.0]]
y = [0, 0, 0, 1, 1, 1, 2, 2, 2, 2]
답
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
# 주어진 데이터셋
x_train = torch.FloatTensor(X)
y_train = torch.LongTensor(y)
# 소프트맥스 회귀 모델 정의
class SoftmaxClassifierModel(nn.Module):
def __init__(self):
super().__init__()
self.linear = nn.Linear(2, 3)
def forward(self, x):
return self.linear(x)
model = SoftmaxClassifierModel()
# optimizer 설정
optimizer = optim.SGD(model.parameters(), lr=0.1)
nb_epochs = 2000
for epoch in range(nb_epochs + 1):
# H(x) 계산
prediction = model(x_train)
# cost 계산
cost = F.cross_entropy(prediction, y_train)
# cost로 H(x) 개선
optimizer.zero_grad()
cost.backward()
optimizer.step()
# 50번마다 로그 출력
if epoch % 200 == 0:
print('Epoch {:4d}/{} Cost: {:.6f}'.format(
epoch, nb_epochs, cost.item()
))