小参数网络中PyTorch训练速度慢于TensorFlow，求优化方法

阿华AIGC实验室

2026-6-26

PyTorch训练速度慢于TensorFlow的问题排查与优化

问题描述

我目前在使用一个小型结构化网络，虽然PyTorch因灵活性高是我的首选，但它的训练速度比TensorFlow慢很多——经过10次训练后，PyTorch的耗时几乎翻倍。是不是我的操作有问题？有没有办法提升PyTorch的计算效率？我复现了PyTorch的同款对比网络，代码如下：

PyTorch模型定义代码

import torch.nn as nn
import torch
import time
from torch.utils.data import DataLoader,TensorDataset
import numpy as np

## Parameters
learning_rate = 0.0005
num_of_epochs = 300
batch_size = 128

## Create Data
xx=torch.arange(40*640*60, dtype=torch.float32).view(640*60, 40)
ff=torch.arange(1*640*60, dtype=torch.float32).view(640*60, 1)
out=torch.arange(1*640*60, dtype=torch.float32).view(640*60, 1)

## Define the model
class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.lin = nn.Linear(40, 1, False)
        self.par = nn.Parameter(torch.Tensor(1))

    def forward(self, x, f):
        return torch.add(self.lin(x),self.par*f)

model = Model()
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
loss_fun = torch.nn.MSELoss()

dataset = TensorDataset(xx,ff,out)
train_dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=False)

PyTorch训练代码

%%time
## Training
model.train(True)
for i in range(n_train):
    start = time.time()
    for iter in range(num_of_epochs):
        train_loss = []
        for x, f, y in train_dataloader:
            # Zero your gradients for every batch!
            optimizer.zero_grad()

            # Make predictions for this batch
            out = model(x, f)

            # Compute the loss and its gradients
            loss = loss_fun(out, y)
            loss.backward()

            # Adjust learning weights
            optimizer.step()
            train_loss.append(loss.item())
        train_loss = np.mean(train_loss)
    end = time.time()
    exe_time['pytorch'].append(end-start)

问题排查与优化方案

1. 未启用CUDA加速（核心原因）

你的代码全程使用CPU训练，而TensorFlow在Colab环境会自动启用GPU，这是速度差距的主要来源。解决方法：

初始化模型时移到GPU：model = Model().to('cuda')

训练循环中把批次数据移到GPU：

for x, f, y in train_dataloader:
    x, f, y = x.to('cuda'), f.to('cuda'), y.to('cuda')
    # 后续训练步骤...

可选：DataLoader设置pin_memory=True，加速CPU到GPU的数据传输。

2. 训练循环的状态重置问题

外层n_train循环中，你没有重新初始化模型和优化器，导致后续训练时模型已收敛，梯度更新效率降低，甚至可能因参数饱和增加计算开销。如果是重复测试耗时，每次循环都要重置：

for i in range(n_train):
    # 每次训练都重新初始化模型和优化器
    model = Model().to('cuda')
    optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
    start = time.time()
    # 后续epoch训练代码...

3. DataLoader效率优化

开启多进程数据加载：设置num_workers=2（Colab环境建议2-4），减少数据加载等待时间：
```
train_dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=False, num_workers=2, pin_memory=True)
```