对比PyTorch和TensorFlow的自动差异和动态子类化模型

使用自定义模型类从头开始训练线性回归，比较PyTorch 1.x和TensorFlow 2.x之间的自动差异和动态模型子类化方法，

这篇简短的文章重点介绍如何在PyTorch 1.x和TensorFlow 2.x中分别使用带有模块/模型API的动态子类化模型，以及这些框架在训练循环中如何使用AutoDiff获得损失的梯度并从头开始实现 一个非常幼稚的渐变后代实现。

生成噪声的线性数据

为了专注于自动差异/自动渐变功能的核心，我们将使用最简单的模型，即线性回归模型，然后我们将首先使用numpy生成一些线性数据，以添加随机级别的噪声。

def generate_data(m=0.1, b=0.3, n=200):
  x = np.random.uniform(-10, 10, n)
  noise = np.random.normal(0, 0.15, n)
  y = (m * x + b ) + noise
  return x.astype(np.float32), y.astype(np.float32)
x, y = generate_data()
plt.figure(figsize = (12,5))
ax = plt.subplot(111)
ax.scatter(x,y, c = "b", label="samples")

模型

然后，我们将在TF和PyTorch中实现从零开始的线性回归模型，而无需使用任何层或激活器，而只需定义两个张量w和b，分别代表线性模型的权重和偏差，并简单地实现线性函数即可：y = wx + b

正如您在下面看到的，我们的模型的TF和PyTorch类定义基本上完全相同，但在一些api名称上只有很小的差异。

唯一值得注意的区别是，PyTorch明确地使用Parameter对象定义权重和要由图形“捕获”的偏置张量，而TF似乎在这里更“神奇”，而是自动捕获用于图形的参数。

确实在PyTorch参数中是Tensor子类，当与Module api一起使用时，它们具有非常特殊的属性，可以自动将自身添加到Module参数列表中，并会出现在在parameters（）迭代器中。

无论如何，两个框架都能够从此类定义和执行方法（__call__或 forward ），参数和图形定义中提取信息，以便向前执行图形执行，并且正如我们将看到的那样，通过自动可微分获得梯度功能，以便能够执行反向传播。

TensorFlow动态模型

class LinearRegressionKeras(tf.keras.Model):
  def __init__(self):
    super().__init__()
    self.w = tf.Variable(tf.random.uniform(shape=[1], -0.1, 0.1))
    self.b = tf.Variable(tf.random.uniform(shape=[1], -0.1, 0.1))
    
  def __call__(self,x): 
    return x * self.w + self.b

PyTorch动态模型

class LinearRegressionPyTorch(torch.nn.Module): 
  def __init__(self): 
    super().__init__() 
    self.w = torch.nn.Parameter(torch.Tensor(1, 1).uniform_(-0.1, 0.1))
    self.b = torch.nn.Parameter(torch.Tensor(1).uniform_(-0.1, 0.1))
  
  def forward(self, x):  
    return x @ self.w + self.b

训练循环，反向传播和优化器

现在我们已经实现了简单的TensorFlow和PyTorch模型，我们可以定义TF和PyTorch api来实现均方误差的损失函数，最后实例化我们的模型类并运行训练循环。

同样，本着眼于自动差异/自动渐变功能核心的目的，我们将使用TF和PyTorch特定的自动差异实现方式实现自定义训练循环，以便为我们的简单线性函数提供渐变并手动优化权重和偏差参数以及临时和朴素的渐变后代优化器。

在TensorFlow训练循环中，我们将特别明确地使用GradientTape API来记录模型的正向执行和损失计算，然后从该GradientTape中获得用于优化权重和偏差参数的梯度。

相反，在这种情况下，PyTorch提供了一种更“神奇”的自动渐变方法，隐式捕获了对参数张量的任何操作，并为我们提供了相同的梯度以用于优化权重和偏置参数，而无需使用任何特定的api。

一旦我们有了权重和偏差梯度，就可以在PyTorch和TensorFlow上实现我们的自定义梯度派生方法，就像将权重和偏差参数减去这些梯度乘以恒定的学习率一样简单。

此处的最后一个微小区别是，当PyTorch在向后传播中更新权重和偏差参数时，以更隐蔽和“魔术”的方式实现自动差异/自动graf时，我们需要确保不要继续让PyTorch从最后一次更新操作中提取grad，这次明确调用no_grad api，最后将权重和bias参数的梯度归零。

TensorFlow训练循环

def squared_error(y_pred, y_true):
  return tf.reduce_mean(tf.square(y_pred - y_true))

tf_model = LinearRegressionKeras()
[w, b] = tf_model.trainable_variables

for epoch in range(epochs):
  with tf.GradientTape() as tape:
    predictions = tf_model(x)
    loss = squared_error(predictions, y)
        
  w_grad, b_grad = tape.gradient(loss, tf_model.trainable_variables)
  w.assign(w - w_grad * learning_rate)
  b.assign(b - b_grad * learning_rate)

  if epoch % 20 == 0:
    print(f"Epoch {epoch} : Loss {loss.numpy()}")

PyTorch训练循环

def squared_error(y_pred, y_true):
  return torch.mean(torch.square(y_pred - y_true))

torch_model = LinearRegressionPyTorch()
[w, b] = torch_model.parameters()

for epoch in range(epochs):
  y_pred = torch_model(inputs)
  loss = squared_error(y_pred, labels)
  loss.backward()
  
  with torch.no_grad():
    w -= w.grad * learning_rate
    b -= b.grad * learning_rate
    w.grad.zero_()
    b.grad.zero_()
    
  if epoch % 20 == 0:
    print(f"Epoch {epoch} : Loss {loss.data}")

结论

正如我们所看到的，TensorFlow和PyTorch自动区分和动态子分类API非常相似，当然，两种模型的训练也给我们非常相似的结果。

在下面的代码片段中，我们将分别使用Tensorflow和PyTorch trainable_variables和parameters方法来访问模型参数并绘制学习到的线性函数的图。

绘制结果

[w_tf, b_tf] = tf_model.trainable_variables
[w_torch, b_torch] = torch_model.parameters()

with torch.no_grad():
  plt.figure(figsize = (12,5))
  ax = plt.subplot(111)
  ax.scatter(x, y, c = "b", label="samples")
  ax.plot(x, w_tf * x + b_tf, "r", 5.0, "tensorflow")
  ax.plot(x, w_torch * inputs + b_torch, "c", 5.0, "pytorch")
  ax.legend()
  plt.xlabel("x1")
  plt.ylabel("y",rotation = 0)

作者：Jacopo Mangiavacchi

本文代码：https://github.com/JacopoMangiavacchi/TF-VS-PyTorch

deephub翻译组