TensorFlow神经网络无法收敛，PyTorch却正常？线性数据集排查

背景 #

生成了一个存在精确线性关系的小型数据集，使用PyTorch搭建单线性层模型，搭配SGD优化器与MSE损失函数，在学习率1e-1、50个epoch的配置下，模型能近似收敛到真实参数值。但使用TensorFlow搭建完全相同逻辑的模型时，50个epoch后权重与真实值差异显著，损失远未接近0，只有将epoch增加到1000才能接近真实参数。

数据集生成代码 #

import numpy as np

def gen_data(n, k):
    np.random.seed(5711)
    beta = np.random.uniform(0, 1, size=(k, 1))
    print("beta is:", beta)
    X = np.random.normal(size=(n, k))
    y = X.dot(beta).reshape(-1, 1)
    D = np.concatenate([X, y], axis=1)
    return D.astype(np.float32)

PyTorch实现代码（可正常收敛） #

import torch
from torch.utils.data import DataLoader, Dataset
from torch import nn
from sklearn.model_selection import train_test_split

n = 10
k = 2
device =  "cpu"

class Daten(Dataset):

    def __init__(self, df):
        self.df = df
        self.ycol = df.shape[1] - 1

    def __getitem__(self, index):
        return self.df[index, :self.ycol], self.df[index, self.ycol:]

    def __len__(self):
        return self.df.shape[0]

def split_into(D, batch_size=64, **kwargs):
    D_train, D_test = train_test_split(D, **kwargs)
    df_train, df_test = Daten(D_train), Daten(D_test)
    dl_train, dl_test = DataLoader(df_train, batch_size=batch_size), DataLoader(df_test, batch_size=batch_size)
    return dl_train, dl_test

D = gen_data(n, k)
dl_train, dl_test = split_into(D, test_size=0.2)

class NeuralNetwork(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = nn.Sequential(
            nn.Linear(k, 1)
        )

    def forward(self, x):
        ypred = self.linear(x)
        return ypred


model = NeuralNetwork().to(device)
print(model)
loss_fn = nn.MSELoss()
optimizer = torch.optim.SGD(model.parameters(), lr=1e-1)

def train(dataloader, model, loss_fn, optimizer):
    size = len(dataloader.dataset)
    model.train()
    for batch, (X, y) in enumerate(dataloader):
        X, y = X.to(device), y.to(device)
        print(y.shape)

        # 计算预测误差
        pred = model(X)
        loss = loss_fn(pred, y)

        # 反向传播
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()

        if batch % 100 == 0:
            loss, current = loss.item(), (batch + 1) * len(X)
            print(f"loss: {loss:>7f}  [{current:>5d}/{size:>5d}]")

epochs = 50
for t in range(epochs):
    print(f"Epoch {t + 1}\n-------------------------------")
    train(dl_train, model, loss_fn, optimizer)
print("Done!")

TensorFlow实现代码（原收敛异常版本） #

import keras.layers
from sklearn.model_selection import train_test_split
from keras.models import Sequential
import tensorflow as tf

n = 10
k = 2
X = gen_data(n, k)
D_train, D_test = train_test_split(X, test_size=0.2)
X_train, y_train = D_train[:,:k], D_train[:,k:]
X_test, y_test = D_test[:,:k], D_test[:,k:]

model = Sequential([keras.layers.Dense(1)])
model.compile(optimizer=tf.keras.optimizers.SGD(lr=1e-1), loss=tf.keras.losses.mean_squared_error)
model.fit(X_train, y_train, batch_size=64, epochs=50)

问题排查与解决方案 #

核心原因 #

1. 多余的偏置参数：生成的数据集是纯线性关系（y = X.dot(beta)），没有偏置项，但TensorFlow的Dense层默认会添加偏置参数，这额外增加了优化维度，拖慢收敛速度；而PyTorch虽然也默认带偏置，但初始化和更新节奏使其能更快将偏置收敛到0附近。
2. 权重初始化差异：PyTorch的nn.Linear默认采用Kaiming均匀初始化，更适合小样本下的快速收敛；TensorFlow的Dense层默认采用Glorot均匀初始化，初始权重分布对线性回归任务的适配性稍差。

修复后的TensorFlow代码 #

修改Dense层去掉偏置项，即可让模型在50个epoch内快速收敛：

import keras.layers
from sklearn.model_selection import train_test_split
from keras.models import Sequential
import tensorflow as tf

n = 10
k = 2
X = gen_data(n, k)
D_train, D_test = train_test_split(X, test_size=0.2)
X_train, y_train = D_train[:,:k], D_train[:,k:]
X_test, y_test = D_test[:,:k], D_test[:,k:]

# 去掉偏置项，匹配数据集的生成逻辑
model = Sequential([keras.layers.Dense(1, use_bias=False)])
model.compile(optimizer=tf.keras.optimizers.SGD(lr=1e-1), loss=tf.keras.losses.mean_squared_error)
model.fit(X_train, y_train, batch_size=64, epochs=50)

# 查看权重
print("Learned weights:", model.get_weights())

进一步对齐PyTorch初始化（可选） #

如果需要完全对齐PyTorch的初始化逻辑，可以手动指定Kaiming初始化：

# 使用He均匀初始化（对应PyTorch的Kaiming初始化）
initializer = tf.keras.initializers.HeUniform()
model = Sequential([keras.layers.Dense(1, use_bias=False, kernel_initializer=initializer)])

内容的提问来源于stack exchange，提问作者P.Jo