为何遍历test_loader会影响PyTorch模型训练的loss值？

我发现只要添加for data in test_loader循环（哪怕循环内没有任何逻辑），模型训练输出的loss值就会发生变化。已经设置了所有相关随机种子，未添加该循环时重复运行程序，loss值完全一致。使用的PyTorch版本为1.9.0+cu111。

以下是复现代码：

import torch
import numpy as np
from torch.utils.data import DataLoader
from torchvision import transforms
from torchvision import datasets

batch_size = 64
learning_rate = 0.01
momentum = 0.5
EPOCH = 10
torch.manual_seed(0)
torch.cuda.manual_seed_all(0)
torch.cuda.manual_seed(0)
np.random.seed(0)

transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))])
train_dataset = datasets.MNIST(root='./data/mnist', train=True, transform=transform,download=True)
test_dataset = datasets.MNIST(root='./data/mnist', train=False, transform=transform,download=True)
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)

class Net(torch.nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = torch.nn.Sequential(
            torch.nn.Conv2d(1, 10, kernel_size=5),
            torch.nn.ReLU(),
            torch.nn.MaxPool2d(kernel_size=2),)
        self.conv2 = torch.nn.Sequential(
            torch.nn.Conv2d(10, 20, kernel_size=5),
            torch.nn.ReLU(),
            torch.nn.MaxPool2d(kernel_size=2),)
        self.fc = torch.nn.Sequential(
            torch.nn.Linear(320, 50),
            torch.nn.Linear(50, 10),)
    def forward(self, x):
        batch_size = x.size(0)
        x = self.conv1(x)
        x = self.conv2(x)
        x = x.view(batch_size, -1)
        x = self.fc(x)
        return x
model = Net()

criterion = torch.nn.CrossEntropyLoss() 
optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate, momentum=momentum) 

def train(epoch):
    running_loss = 0.0
    running_total = 0
    running_correct = 0
    for batch_idx, data in enumerate(train_loader, 0):
        optimizer.zero_grad()

        inputs, target = data
        outputs = model(inputs)
        loss = criterion(outputs, target)

        loss.backward()
        optimizer.step()

        running_loss += loss.item()
        _, predicted = torch.max(outputs.data, dim=1)
        running_total += inputs.shape[0]
        running_correct += (predicted == target).sum().item()

    print(epoch + 1, running_loss / len(train_loader))

def test():
    with torch.no_grad():    
        for data in test_loader:         # 注释/保留此行会导致loss变化
            a=0
    return 0

if __name__ == '__main__':
    acc_list_test = []
    for epoch in range(EPOCH):
        train(epoch)
        acc_test = test()

这是PyTorch 1.9版本中多进程DataLoader的随机状态污染问题：

• 默认情况下，DataLoader使用多进程（num_workers>0）加载数据，遍历DataLoader时会启动子进程。
• 子进程会继承主进程的随机状态，但旧版本中，子进程的初始化会意外修改主进程的随机数生成器状态。
• 你的训练DataLoader设置了shuffle=True，主进程随机状态被改变后，训练数据的打乱顺序会发生变化，最终导致训练loss改变。

方法1：固定子进程的随机种子 #

给DataLoader添加worker_init_fn，确保每个子进程的随机种子独立且固定，避免污染主进程状态：

def worker_init_fn(worker_id):
    # 给每个子进程分配独立的固定种子
    np.random.seed(0 + worker_id)
    torch.manual_seed(0 + worker_id)

# 修改DataLoader定义
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, worker_init_fn=worker_init_fn)
test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False, worker_init_fn=worker_init_fn)

方法2：禁用多进程数据加载 #

将num_workers设为0，避免子进程带来的随机状态干扰（适合调试场景）：

train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=0)
test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False, num_workers=0)

方法3：升级PyTorch版本 #

PyTorch在1.10及以后版本中修复了多进程DataLoader的随机种子继承问题，升级后该现象会自动消失。

内容的提问来源于stack exchange，提问作者zhihao Yao