【Pytorch学习】-- 补充 -- 使用Pytorch自带的Loss和Optimizer

修改

1. 损失函数
2. 更新参数用Optimizer代替

代码

import torch
import torch.nn as nn

# 已知公式形式： f = w * x，未知参数w的数值
# 目标 w = 2

# 初始化数据
# 数据集 X
X = torch.tensor([1,2,3,4],dtype = torch.float32)
# 验证集 Y
Y = torch.tensor([2,4,6,8],dtype = torch.float32)
# 初始 w
w = torch.tensor(0.0,dtype = torch.float32, requires_grad = True)

# 前向传播
def forward(x):
  return w * X

# 计算损失
# 损失函数使用方差(MSE)
loss = nn.MSELoss()
# 优化器，用于更新w
optimizer = torch.optim.SGD([w], lr=learning_rate)

print(f'Prediction before training: f(5) = {forward(5).item():.3f}')

# 2) Define loss and optimizer
learning_rate = 0.01
n_iters = 100

# callable function
loss = nn.MSELoss()

optimizer = torch.optim.SGD([w], lr=learning_rate)

# 训练
learning_rate = 1e-2
n_iters = 100

for epoch in range(n_iters):
  # 计算出prediction
  y_pred = forward(X)
  
  # 计算损失
  l = loss(Y,y_pred)

  # 计算梯度
  l.backward()

  # 更新 w
  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=)
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=)