相同网络与超参下TensorFlow和PyTorch的MSE结果差异原因

在神经网络结构、超参数配置完全对齐的前提下，仅将实现框架从TensorFlow切换为PyTorch，训练与测试阶段得到的均方误差（MSE）存在小幅差异。本次实验的具体配置如下：

• 网络结构：输入维度70，包含2个各有40个神经元的隐藏层，输出维度为2
• 训练超参数：Adam优化器，学习率0.001，epsilon取值1e-08，训练轮次共600轮
  核心疑问：
• 该差异产生的原因是什么？
• 两种框架是否本应输出完全一致的结果？

TensorFlow 版本 #

import tensorflow as tf
from tensorflow import keras
import pandas as pd
import numpy as np
from tensorflow.keras import layers
from tensorflow.keras import losses
from tensorflow.keras.optimizers import Adam

train_df = pd.read_csv('./train.csv')

inputs = keras.Input(shape=(70,))
x = layers.Dense(40, activation='relu')(inputs)
x = layers.Dense(40, activation='relu')(x)

# Output
outputs = layers.Dense(2, )(x)

model = keras.Model(inputs=inputs, outputs=outputs, name='my_model')
new_model = model.compile(loss=losses.MeanSquaredError(), optimizer=Adam(lr=0.001, epsilon=1e-08))

inputs = train_df.iloc[:, 0:70].to_numpy()
print(type(inputs))
targets = train_df.iloc[:, 70:72].to_numpy()
print(type(targets))
epochs = 600

for epoch in range(0, epochs):
    print("Epoch: " + str(epoch + 1))
    model.fit(inputs, targets, epochs=1, )

test_df = pd.read_csv('./test.csv')

inputTest = test_df.iloc[:, 0:70].to_numpy()
targetTest = test_df.iloc[:, 70:72].to_numpy()

outputTest = model.predict(inputTest)
mse = np.mean(np.power(targetTest - outputTest, 2), axis=1)
totalMse = mse.mean()
print("The mean square error is : " + str(totalMse))

PyTorch 版本 #

import pandas as pd
import numpy as np

import numpy as np
import torch
import torch.optim as optim
import torch.nn.functional as F
import torch.nn as nn

train_df = pd.read_csv('./train.csv')

class newModel(nn.Module):
    def __init__(self):
        super(newModel, self).__init__()
        self._layers = 40
        self._output_dim = 2
        self._linear_one = nn.Linear(70, self._layers)
        self._linear_two = nn.Linear(self._layers, self._layers)
        self._linear_three = nn.Linear(self._layers, self._output_dim)

    def forward(self, x):
        x = F.relu(self._linear_one(x))
        x = F.relu(self._linear_two(x))
        output = (self._linear_three(x))
        return output


device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
criterion = nn.MSELoss()
model = newModel()
learning_rate = 0.001
num_epochs = 600
optimizer = optim.Adam(model.parameters(), lr=learning_rate)

inputs = train_df.iloc[:, 0:70].to_numpy()
targets = train_df.iloc[:, 70:72].to_numpy()

inputs = torch.from_numpy(inputs.astype(np.float32))
targets = torch.from_numpy(targets.astype(np.float32))
inputs = torch.reshape(inputs, (-1, 70)).to(device)
targets = torch.reshape(targets, (-1, 2)).to(device)

for epoch in range(num_epochs):
    # forward feed
    y_pred = model.forward(inputs)

    loss = criterion(y_pred, targets)
    loss.backward()
    optimizer.step()
    optimizer.zero_grad()

    print('epoch {}, loss {}'.format(epoch, loss.item()))

test_df = pd.read_csv('./test.csv')

inputs = test_df.iloc[:, 0:70].to_numpy()
targets = test_df.iloc[:, 70:72].to_numpy()

inputs = torch.from_numpy(inputs.astype(np.float32))
inputs = torch.reshape(inputs, (-1, 70)).to(device)

outputs = model.forward(inputs).detach().numpy()

mse = np.mean(np.power(targets - outputs, 2), axis=1)
totalMse = mse.mean()
print("The mean square error is : " + str(totalMse))

首先给出明确结论：两种框架本身就不保证输出完全一致的结果，观察到的小幅MSE差异是正常现象，而且贴出的两份代码本身就存在未对齐的训练配置，进一步放大了差异。

具体原因可以分为三类：

1. 代码层面未对齐的显式配置
   • 训练批量大小不一致：Keras的model.fit()如果不手动指定batch_size参数，默认会用32的小批量做训练；而PyTorch代码是把全部训练样本一次性输入网络，做全批量梯度下降，两种梯度下降方式的参数更新轨迹本身就存在差异。
   • 默认权重初始化规则不同：Keras的Dense层默认采用Glorot（Xavier）均匀分布初始化权重，偏置默认初始化为0；PyTorch的nn.Linear层默认采用Kaiming（He）均匀分布初始化权重，偏置默认从范围为$[-1/\sqrt{fan_in}, 1/\sqrt{fan_in}]$的均匀分布采样初始化。初始权重从训练第一步就不一样，后续结果自然会产生分叉。
2. 框架底层实现的隐式差异
   • Adam优化器的实现细节存在区别：两个框架的Adam优化器在epsilon的作用位置、动量更新的计算顺序、权重衰减的默认处理逻辑上都有细微差别，这些差异会在多轮训练中逐步累计。
   • 训练流程的默认行为不同：Keras的fit方法默认会在每个epoch前对训练样本做shuffle，给出的PyTorch代码完全没有做样本打乱，每轮迭代的样本顺序固定，也会带来训练轨迹的差异。
3. 浮点计算的固有非确定性
   哪怕把上述所有配置全部对齐，也不可能得到逐比特完全一致的结果：
   • GPU上的矩阵乘法、激活函数等算子本身存在非确定性，不同框架调用的底层CUDA kernel实现不同，浮点数的累加顺序差异会带来微小的计算误差，float32精度本身只有7位左右有效数字，误差经过600轮迭代累计后会产生可观测的MSE差异。
   • 两个框架的自动求导实现逻辑不同，梯度计算的数值精度也会存在微小偏差。

如果需要把差异缩小到浮点误差级别，可以做如下对齐操作：

• 两边固定相同的随机种子，手动给网络层加载完全一致的初始权重
• 统一训练的batch size、样本shuffle规则、数据预处理逻辑
• 开启框架的确定性计算模式：TensorFlow调用tf.config.experimental.enable_op_determinism()，PyTorch调用torch.use_deterministic_algorithms(True)
• 逐项对齐Adam优化器的所有参数，确认beta1、beta2、epsilon、权重衰减等配置完全一致，且实现逻辑匹配

内容的提问来源于stack exchange，提问作者Andres TM