TensorFlow与PyTorch同模型训练后结果不一致问题排查

我在TensorFlow和PyTorch中实现了结构完全一致的简单网络，通过将TensorFlow的初始权重复制到PyTorch模型中，确保两者初始状态完全一致，此时模型对训练数据的预测结果完全相同。但运行5轮训练后，两者的预测结果出现明显分歧——我使用了相同的超参数与优化器，这一现象的原因是什么？

加载数据 #

from sklearn.datasets import load_diabetes
X, y = load_diabetes(as_frame=True, return_X_y = True)
shape = X.to_numpy().shape[1]

初始化TensorFlow模型 #

import tensorflow as tf

feature_inputs = tf.keras.layers.Input(name="feature_input", shape=shape)

hidden_layer = tf.keras.layers.Dense(10)(feature_inputs)
activation = tf.keras.layers.Activation('relu')(hidden_layer)

output = tf.keras.layers.Dense(1)(activation)
model = tf.keras.Model(inputs=feature_inputs, outputs=output)

keras_weights = model.get_weights()

初始化PyTorch模型 #

import torch
from torch import nn
from torch.utils.data import DataLoader

class MLP(nn.Module):
    def __init__(self, input_shape, *args, **kwargs):
        super().__init__()
        self.layer1 = nn.Linear(input_shape, 10)
        self.layer2 = nn.Linear(10, 1)

        self.relu = nn.ReLU()

    def forward(self, x):
        x = self.layer1(x)
        x = self.relu(x)
        x = self.layer2(x)

        return x
    
class Dataset(torch.utils.data.Dataset):
    def __init__(self, X, y):
        if not torch.is_tensor(X) and not torch.is_tensor(y):
            self.X = torch.from_numpy(X)
            self.y = torch.from_numpy(y)

    def __len__(self):
        return len(self.X)

    def __getitem__(self, i):
        return self.X[i], self.y[i]


train_data = Dataset(X=X.to_numpy(), y=y.to_numpy())
trainloader = torch.utils.data.DataLoader(train_data, shuffle=False)
mlp = MLP(input_shape=shape)

torch_weights = mlp.state_dict()

复制初始权重 #

new_weights = {}

for idx, key in enumerate(torch_weights.keys()):
    if key.endswith('weight'):
        new_weights[key] = torch.Tensor(keras_weights[idx].T)
    else:
        new_weights[key] = torch.Tensor(keras_weights[idx])


mlp.load_state_dict(new_weights)

初始预测结果对比 #

此时模型对训练数据的预测结果完全一致：

mlp(torch.from_numpy(X.astype(np.float32).to_numpy()))[0]
# 返回 -0.0361
model.predict(X)[0]
# 返回 -0.0361313

训练实现与分歧现象 #

TensorFlow训练代码 #

optimizer = tf.keras.optimizers.SGD(learning_rate=1e-8)

model.compile(optimizer, loss="mean_squared_error")
history = model.fit(x=X, y=y,
                    epochs=5,
                    verbose=True,
                    shuffle=False)

PyTorch训练代码 #

loss_function = nn.MSELoss()
optimizer = torch.optim.SGD(mlp.parameters(), lr=1e-8)

for epoch in range(0, 5):
    for i, train_data in enumerate(trainloader, 0):
        inputs, targets = train_data
        inputs, targets = inputs.float(), targets.float()
        targets = targets.reshape((targets.shape[0], 1))

        optimizer.zero_grad()
        outputs = mlp(inputs)

        training_loss = loss_function(outputs, targets)

        training_loss.backward()

        optimizer.step()

训练后再次预测，结果出现明显分歧：

mlp(torch.from_numpy(X.astype(np.float32).to_numpy()))[0]
# 返回 -0.025996
model.predict(X)[0]
# 返回 -0.037605

核心原因分析 #

• 批量更新 vs 逐样本更新：这是最关键的差异。TensorFlow的model.fit默认使用批量梯度下降（Batch GD）——默认batch_size=32，即每计算32个样本的平均梯度后更新一次权重；而你的PyTorch代码中DataLoader未指定batch_size，默认值为1，属于随机梯度下降（Stochastic GD）——每个样本单独计算梯度并立即更新权重。两种更新方式的梯度方向和幅度差异会随着训练轮次快速放大，导致权重走向完全不同。
• 数值精度的细微差异：虽然初始权重一致，但在训练过程中，TensorFlow和PyTorch的浮点数计算逻辑（如累加顺序、数值舍入）存在细微差别，会在多次梯度更新后逐步累积，放大结果分歧。
• Loss计算的细节差异：PyTorch的nn.MSELoss默认对所有元素求平均，而TensorFlow的MeanSquaredError在默认情况下也是求平均，但如果样本维度处理有细微差别（比如是否保留维度），也可能导致梯度计算的微小差异。

要让两者结果一致，需要统一训练的批量大小：比如在PyTorch的DataLoader中设置batch_size=32（和TensorFlow默认一致），同时确保数据顺序完全相同（已设置shuffle=False），并统一数据类型精度（比如都使用float32）。

内容的提问来源于stack exchange，提问作者Nils Mackay