PyTorch: optim包

译者：@yongjay13、@speedmancs

校对者：@bringtree

本例中的全连接神经网络有一个隐藏层, 后接ReLU激活层, 并且不带偏置参数. 训练时通过最小化欧式距离的平方, 来学习从x到y的映射

在此实现中, 我们将弃用之前手工更新权值的做法, 转而用PyTorch的nn包来搭建神经网络. optim包则用来定义更新权值的优化器. optim包有众多深度学习常用的优化算法, 包括SGD+momentum, RMSProp, Adam等.

  1. import torch
  2. from torch.autograd import Variable
  4. # N 批量大小; D_in是输入尺寸;
  5. # H是隐藏尺寸; D_out是输出尺寸.
  6. N, D_in, H, D_out = 64, 1000, 100, 10
  8. # 创建随机张量来保存输入和输出,并将它们包装在变量中.
  9. x = Variable(torch.randn(N, D_in))
  10. y = Variable(torch.randn(N, D_out), requires_grad=False)
  12. # 使用nn包来定义我们的模型和损失函数.
  13. model = torch.nn.Sequential(
  14.     torch.nn.Linear(D_in, H),
  15.     torch.nn.ReLU(),
  16.     torch.nn.Linear(H, D_out),
  17. )
  18. loss_fn = torch.nn.MSELoss(size_average=False)
  20. # 使用优化包来定义一个优化器,它将为我们更新模型的权重.
  21. # 在这里,我们将使用 Adam;这个 optim 包包含许多其他优化算法.
  22. # Adam构造函数的第一个参数告诉优化器应该更新哪个Variables.
  23. learning_rate = 1e-4
  24. optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
  25. for t in range(500):
  26.     # 正向传递:通过将x传递给模型来计算预测的y.
  27.     y_pred = model(x)
  29.     # 计算和打印损失函数.
  30.     loss = loss_fn(y_pred, y)
  31.     print(t, loss.data[0])
  33.     # 在向后传递之前,使用优化器对象为其要更新的变量(这是模型的可学习权重）的所有梯度归零.
  34.     # 这是因为默认情况下,只要调用.backward(),梯度就会在缓冲区中累积(即不会被覆盖).
  35.     # 查看torch.autograd.backward的文档以获取更多详细信息.
  36.     optimizer.zero_grad()
  38.     # 向后传递:计算损失函数相对于模型参数的梯度
  39.     loss.backward()
  41.     # 在优化器上调用step函数会更新其参数
  42.     optimizer.step()