PyTorch：自定义nn模块

原文：https://pytorch.org/tutorials/beginner/examples_nn/polynomial_module.html#sphx-glr-beginner-examples-nn-polynomial-module-py

经过训练的三阶多项式，可以通过最小化平方的欧几里得距离来预测y = sin(x)从-pi到pi。

此实现将模型定义为自定义Module子类。 每当您想要一个比现有模块的简单序列更复杂的模型时，都需要以这种方式定义模型。

  1. import torch
  2. import math
  4. class Polynomial3(torch.nn.Module):
  5.     def __init__(self):
  6.         """
  7.         In the constructor we instantiate four parameters and assign them as
  8.         member parameters.
  9.         """
  10.         super().__init__()
  11.         self.a = torch.nn.Parameter(torch.randn(()))
  12.         self.b = torch.nn.Parameter(torch.randn(()))
  13.         self.c = torch.nn.Parameter(torch.randn(()))
  14.         self.d = torch.nn.Parameter(torch.randn(()))
  16.     def forward(self, x):
  17.         """
  18.         In the forward function we accept a Tensor of input data and we must return
  19.         a Tensor of output data. We can use Modules defined in the constructor as
  20.         well as arbitrary operators on Tensors.
  21.         """
  22.         return self.a + self.b * x + self.c * x ** 2 + self.d * x ** 3
  24.     def string(self):
  25.         """
  26.         Just like any class in Python, you can also define custom method on PyTorch modules
  27.         """
  28.         return f'y = {self.a.item()} + {self.b.item()} x + {self.c.item()} x^2 + {self.d.item()} x^3'
  30. # Create Tensors to hold input and outputs.
  31. x = torch.linspace(-math.pi, math.pi, 2000)
  32. y = torch.sin(x)
  34. # Construct our model by instantiating the class defined above
  35. model = Polynomial3()
  37. # Construct our loss function and an Optimizer. The call to model.parameters()
  38. # in the SGD constructor will contain the learnable parameters of the nn.Linear
  39. # module which is members of the model.
  40. criterion = torch.nn.MSELoss(reduction='sum')
  41. optimizer = torch.optim.SGD(model.parameters(), lr=1e-6)
  42. for t in range(2000):
  43.     # Forward pass: Compute predicted y by passing x to the model
  44.     y_pred = model(x)
  46.     # Compute and print loss
  47.     loss = criterion(y_pred, y)
  48.     if t % 100 == 99:
  49.         print(t, loss.item())
  51.     # Zero gradients, perform a backward pass, and update the weights.
  52.     optimizer.zero_grad()
  53.     loss.backward()
  54.     optimizer.step()
  56. print(f'Result: {model.string()}')

脚本的总运行时间：（0 分钟 0.000 秒）

下载 Python 源码：polynomial_module.py

下载 Jupyter 笔记本：polynomial_module.ipynb

由 Sphinx 画廊生成的画廊