PyTorch中不同DataLoader批量大小为何导致损失差异显著？

核心原因分析 #

• 参数更新步数差异：
  当使用全量数据（BATCH_SIZE=10000）训练时，每个epoch仅执行1次参数更新；而BATCH_SIZE=512时，每个epoch会执行10000//512≈19次更新。在最初20个epoch的设置下，全量数据的总更新次数只有20次，小批量的总更新次数却接近380次——参数迭代的次数差了一个数量级，模型优化的程度自然差距明显。这也是你将epoch增加到400后效果接近的核心原因：400次全量更新的总步数，和小批量20个epoch的更新步数基本对齐，参数得到了足够的优化。

• 梯度估计的特性差异：
  小批量的梯度是对全量真实梯度的随机近似，带有一定噪声，这种噪声有时能帮助模型避开局部最优；而全量梯度是精确计算的，收敛方向更稳定，但如果学习率或训练步数不足，容易卡在次优解。不过你已经固定了随机种子并关闭洗牌，小批量的梯度序列是固定的，这种情况下更新步数的差异是主导因素。

你的修改说明 #

你对原代码做了以下调整来控制变量：

• 固定所有随机种子、关闭数据洗牌，排除随机性干扰
• 将学习率调整为1e-4，样本量设为10**4，训练轮次初始设为20
• 添加绘图逻辑，直观展示不同批量下的近似效果差异

完整代码 #

import torch
import torch.nn as nn
import numpy as np
import matplotlib.pyplot as plt

# 固定随机种子
torch.manual_seed(42)
np.random.seed(42)

# 生成数据
X = torch.linspace(-np.pi, np.pi, 10**4).unsqueeze(1)
y = torch.sin(X)

# 定义模型
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.fc1 = nn.Linear(1, 128)
        self.fc2 = nn.Linear(128, 128)
        self.fc3 = nn.Linear(128, 1)
    
    def forward(self, x):
        x = torch.relu(self.fc1(x))
        x = torch.relu(self.fc2(x))
        x = self.fc3(x)
        return x

def train_model(batch_size, epochs):
    model = Net()
    criterion = nn.MSELoss()
    optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)
    
    # 关闭洗牌，固定数据顺序
    dataloader = torch.utils.data.DataLoader(
        torch.utils.data.TensorDataset(X, y),
        batch_size=batch_size,
        shuffle=False
    )
    
    for epoch in range(epochs):
        total_loss = 0.0
        for batch_X, batch_y in dataloader:
            optimizer.zero_grad()
            outputs = model(batch_X)
            loss = criterion(outputs, batch_y)
            loss.backward()
            optimizer.step()
            total_loss += loss.item()
        if (epoch+1) % 50 == 0:
            print(f"Batch Size {batch_size}, Epoch {epoch+1}, Loss: {total_loss/len(dataloader):.6f}")
    return model

# 训练不同批量的模型
model_512 = train_model(512, 20)
model_full = train_model(10000, 400)

# 绘图对比
plt.figure(figsize=(12, 6))
plt.plot(X.numpy(), y.numpy(), label='True sin(x)', color='blue')
plt.plot(X.numpy(), model_512(X).detach().numpy(), label='Batch Size 512 (20 Epochs)', color='orange', alpha=0.7)
plt.plot(X.numpy(), model_full(X).detach().numpy(), label='Batch Size 10000 (400 Epochs)', color='green', alpha=0.7)
plt.legend()
plt.xlabel('x')
plt.ylabel('y')
plt.title('Sin(x) Approximation with Different Batch Sizes')
plt.show()

内容的提问来源于stack exchange，提问作者Sagi Mann