PyTorch：控制流 + 权重共享

原文：https://pytorch.org/tutorials/beginner/examples_nn/dynamic_net.html#sphx-glr-beginner-examples-nn-dynamic-net-py

为了展示 PyTorch 动态图的强大功能，我们将实现一个非常奇怪的模型：一个从三到五阶（动态变化）的多项式，在每次正向传递中选择一个 3 到 5 之间的一个随机数，并将这个随机数作为阶数，第四和第五阶共用同一个权重来多次重复计算。

  1. import random
  2. import torch
  3. import math
  5. class DynamicNet(torch.nn.Module):
  6.     def __init__(self):
  7.         """
  8.         In the constructor we instantiate five parameters and assign them as members.
  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(()))
  15.         self.e = torch.nn.Parameter(torch.randn(()))
  17.     def forward(self, x):
  18.         """
  19.         For the forward pass of the model, we randomly choose either 4, 5
  20.         and reuse the e parameter to compute the contribution of these orders.
  22.         Since each forward pass builds a dynamic computation graph, we can use normal
  23.         Python control-flow operators like loops or conditional statements when
  24.         defining the forward pass of the model.
  26.         Here we also see that it is perfectly safe to reuse the same parameter many
  27.         times when defining a computational graph.
  28.         """
  29.         y = self.a + self.b * x + self.c * x ** 2 + self.d * x ** 3
  30.         for exp in range(4, random.randint(4, 6)):
  31.             y = y + self.e * x ** exp
  32.         return y
  34.     def string(self):
  35.         """
  36.         Just like any class in Python, you can also define custom method on PyTorch modules
  37.         """
  38.         return f'y = {self.a.item()} + {self.b.item()} x + {self.c.item()} x^2 + {self.d.item()} x^3 + {self.e.item()} x^4 ? + {self.e.item()} x^5 ?'
  40. # Create Tensors to hold input and outputs.
  41. x = torch.linspace(-math.pi, math.pi, 2000)
  42. y = torch.sin(x)
  44. # Construct our model by instantiating the class defined above
  45. model = DynamicNet()
  47. # Construct our loss function and an Optimizer. Training this strange model with
  48. # vanilla stochastic gradient descent is tough, so we use momentum
  49. criterion = torch.nn.MSELoss(reduction='sum')
  50. optimizer = torch.optim.SGD(model.parameters(), lr=1e-8, momentum=0.9)
  51. for t in range(30000):
  52.     # Forward pass: Compute predicted y by passing x to the model
  53.     y_pred = model(x)
  55.     # Compute and print loss
  56.     loss = criterion(y_pred, y)
  57.     if t % 2000 == 1999:
  58.         print(t, loss.item())
  60.     # Zero gradients, perform a backward pass, and update the weights.
  61.     optimizer.zero_grad()
  62.     loss.backward()
  63.     optimizer.step()
  65. print(f'Result: {model.string()}')

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

下载 Python 源码：dynamic_net.py

下载 Jupyter 笔记本：dynamic_net.ipynb

由 Sphinx 画廊生成的画廊