修改
- 损失函数
- 更新参数用Optimizer代替
代码
import torch
import torch.nn as nn
X = torch.tensor([1,2,3,4],dtype = torch.float32)
Y = torch.tensor([2,4,6,8],dtype = torch.float32)
w = torch.tensor(0.0,dtype = torch.float32, requires_grad = True)
def forward(x):
return w * X
loss = nn.MSELoss()
optimizer = torch.optim.SGD([w], lr=learning_rate)
print(f'Prediction before training: f(5) = {forward(5).item():.3f}')
learning_rate = 0.01
n_iters = 100
loss = nn.MSELoss()
optimizer = torch.optim.SGD([w], lr=learning_rate)
learning_rate = 1e-2
n_iters = 100
for epoch in range(n_iters):
y_pred = forward(X)
l = loss(Y,y_pred)
l.backward()
optimizer.step()
optimizer.zero_grad()
if epoch % 10 == 0:
print(f"epoch{epoch+1}: w = {w:.3f}, loss = {l:.8f}")
print(f'Prediction after training: f(5) = {forward(5).item():.3f}')
Prediction before training: f(5) = tensor([0., 0., 0., 0.], grad_fn=<MulBackward0>)
epoch1: w = 0.300, loss = 30.00000000
epoch11: w = 1.665, loss = 1.16278565
epoch21: w = 1.934, loss = 0.04506890
epoch31: w = 1.987, loss = 0.00174685
epoch41: w = 1.997, loss = 0.00006770
epoch51: w = 1.999, loss = 0.00000262
epoch61: w = 2.000, loss = 0.00000010
epoch71: w = 2.000, loss = 0.00000000
epoch81: w = 2.000, loss = 0.00000000
epoch91: w = 2.000, loss = 0.00000000
Prediction after training: f(5) = tensor([2.0000, 4.0000, 6.0000, 8.0000], grad_fn=<MulBackward0>)
|