为何我的PyTorch代码运行速度比TensorFlow版本慢10倍？

针对你的PyTorch代码在CPU/GPU上训练速度远慢于TensorFlow的问题，以下是具体的优化点和修改建议：

1. 优化数据加载器的并行性 #

你的DataLoader默认使用num_workers=0，即主线程负责数据加载，这会导致GPU/CPU在训练时频繁等待数据，成为核心性能瓶颈。TensorFlow的model.fit默认使用异步数据加载，因此不会出现此类问题。

修改方式：创建train_loader和test_loader时，设置num_workers为CPU核心数的一半（你的i7-10为8核，建议设为4或8）：

train_loader = DataLoader(train, batch_size=BATCH_SIZE, shuffle=True, pin_memory=True, num_workers=4)
test_loader = DataLoader(test, batch_size=BATCH_SIZE, shuffle=False, pin_memory=True, num_workers=4)

注意：Windows系统下num_workers不宜设置过大，避免进程卡死，建议从2、4开始测试。

2. 关闭确定性模式，启用cudnn基准测试 #

你设置了torch.backends.cudnn.deterministic=True，这会强制cudnn使用确定性算法，牺牲性能来保证结果一致。而TensorFlow默认不开启严格确定性模式，可调整为：

if torch.cuda.is_available():
    torch.cuda.manual_seed(RANDOM_SEED)
    # 关闭确定性模式，开启基准测试以选择最优计算算法
    torch.backends.cudnn.deterministic = False
    torch.backends.cudnn.benchmark = True

这能让cudnn根据硬件自动选择最快的全连接/卷积算法，大幅提升GPU训练速度。

3. 启用混合精度训练（GPU环境） #

PyTorch可通过torch.cuda.amp启用混合精度训练，用半精度（FP16）计算部分张量，减少显存占用并提升计算速度。修改训练循环如下：

from torch.cuda.amp import GradScaler, autocast

# 初始化混合精度工具
scaler = GradScaler()

def train_loop(dataloader, model, loss_fn, optimizer, scaler):
    model.train()  # 显式设置模型为训练模式
    size = len(dataloader.dataset)
    avg_loss = 0
    for batch, (X, y) in enumerate(dataloader):
        X, y = X.to(device), y.to(device)
        # 混合精度计算上下文
        with autocast():
            pred = model(X)
            loss = loss_fn(pred, y)
        avg_loss += loss.item() * X.size(0)  # 修正损失累加逻辑：乘以batch size再求平均
        
        # 反向传播与优化
        optimizer.zero_grad()
        scaler.scale(loss).backward()
        scaler.step(optimizer)
        scaler.update()
    return avg_loss / size

# 训练循环调用时传入scaler
for t in range(EPOCHS):
    print(f"Epoch {t+1}...", end="")
    avg_loss = train_loop(train_loader, model, loss_fn, optimizer, scaler)
    print(f"Avg loss: {avg_loss:.4f}")

4. 使用PyTorch 2.0+的模型编译功能 #

如果你的PyTorch版本是2.0及以上，使用torch.compile可一键优化模型计算图，同时提升CPU和GPU的训练速度：

model = nn.Sequential(
    nn.Flatten(),
    nn.Linear(28*28, 256),
    nn.ReLU(),
    nn.Dropout(0.3),
    nn.Linear(256, 128),
    nn.ReLU(),
    nn.Linear(128, 10)
).to(device)
# 编译模型以优化性能
model = torch.compile(model)

5. 统一数据归一化方式（可选） #

你的TensorFlow代码将数据归一化到[0,1]，而PyTorch代码归一化到[-1,1]，虽不影响准确率，但统一归一化方式可让性能对比更严谨。修改transform为：

transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.0,), (1.0/255.0,))  # 等价于x/255.0
])

优化后预期效果 #

经过以上优化，GPU训练时间可大幅降低（预计能到10秒以内，接近甚至超过TensorFlow的CPU训练速度），CPU训练时间也会显著缩短。

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