x_data = torch.tensor([[1.0],[2.0],[3.0]])
y_data = torch.tensor([[2.0],[4.0],[6.0]])

class LinearModel(torch.nn.Module):
? ? def __init__(self):
? ? ? ? super(LinearModel,self).__init__()
? ? ? ? self.linear = torch.nn.Linear(1,1)
? ? ? ??
? ? def forward(self, x):
? ? ? ? y_pred = self.linear(x)
? ? ? ? return y_pred
? ? ? ??
model = LinearModel() ?

criterion = torch.nn.MSELoss(reduction='sum')
optimizer = torch.optim.SGD(model.parameters(),lr=0.01)

epoch_list = []
loss_list = []
w_list = []
b_list = []
for epoch in range(1000):
? ? y_pred = model(x_data)?? ??? ??? ??? ??? ? ?# 計(jì)算預(yù)測(cè)值
? ? loss = criterion(y_pred, y_data)?? ?# 計(jì)算損失
? ? print(epoch,loss)
? ??
? ? epoch_list.append(epoch)
? ? loss_list.append(loss.data.item())
? ? w_list.append(model.linear.weight.item())
? ? b_list.append(model.linear.bias.item())
? ??
? ? optimizer.zero_grad() ? # 梯度歸零
? ? loss.backward() ? ? ? ? # 反向傳播
? ? optimizer.step() ? ? ? ?# 更新

# 輸出權(quán)重和偏置
print('w = ',model.linear.weight.item())
print('b = ',model.linear.bias.item())

# 測(cè)試模型
x_test = torch.tensor([[4.0]])
y_test = model(x_test)
print('y_pred = ',y_test.data)

y_pred = ?tensor([[7.9997]])

# 二維曲線圖
plt.plot(epoch_list,loss_list,'b')
plt.xlabel('epoch')
plt.ylabel('loss')
plt.show()

# 三維散點(diǎn)圖
fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
ax.scatter(w_list,b_list,loss_list,c='r')
#設(shè)置坐標(biāo)軸
ax.set_xlabel('weight')
ax.set_ylabel('bias')
ax.set_zlabel('loss')
plt.show()