如何在PyTorch中利用多CPU加速模型训练？

问题描述 #

无法使用GPU，希望将基于PyTorch的模型训练从单CPU切换为利用4个CPU满负荷运行。当前训练代码仅占用1个CPU，尝试过DistributedDataParallel但未生效，基础代码及Dataset/Network类见下文。

基础代码（用户提供） #

import pandas as pd

import torch
from torch import nn
from torch.utils.data import Dataset, DataLoader


df = pd.read_pickle('path/to/data')

X = df.drop(columns=['target'])
y = df[['target']]

train_data = CustomDataset(X, y)
train_loader = DataLoader(
    dataset=train_data,
    batch_size=64
)

model = Network(X.shape[-1])
criterion = nn.MSELoss()
optimizer = torch.optim.Adam(model.parameters())

epochs = 10

for e in range(1, epochs + 1):
    train_loss = .0

    model.train()
    for batch_id, (data, labels) in enumerate(train_loader):
        optimizer.zero_grad()
        target = model(data)
        loss = criterion(target, labels)
        loss.backward()
        optimizer.step()
        train_loss += loss.item()

    print('Epoch {}:	Train Loss: {:.4f}'.format(
        e,
        train_loss / len(train_loader)
    ))

附录：CustomDataset与Network类 #

class CustomDataset(Dataset):
    def __init__(
        self,
        X,
        y,
    ):
        self.X = torch.Tensor(X.values)
        self.y = torch.Tensor(y.values)

    def __getitem__(
        self,
        index
    ):
        return self.X[index], self.y[index]

    def __len__(self):
        return len(X)


class Network(nn.Module):
    def __init__(
        self,
        input_size
    ):
        super(Network, self).__init__()

        self.input_size = input_size

        self.linear_1 = nn.Linear(self.input_size, 32)
        self.linear_2 = nn.Linear(32, 1)


    def forward(self, data):
        output = self.linear_1(data)
        output = self.linear_2(output)

        return output

优化方案（分步骤实现最佳性能） #

1. 优化DataLoader多进程数据加载

这是见效最快的一步，让数据预处理/加载环节占用多个CPU，避免数据成为训练瓶颈。修改DataLoader参数：

train_loader = DataLoader(
    dataset=train_data,
    batch_size=64,
    num_workers=4,  # 启用4个进程并行加载数据
    pin_memory=True  # 加速数据从CPU到模型的传递（无GPU场景同样生效）
)

注意：Windows系统必须将训练代码包裹在if __name__ == '__main__':块内，否则num_workers会触发多进程报错；Linux/macOS建议添加该块以避免潜在问题。

2. 模型多CPU并行计算

单机器多CPU场景下，无需复杂的分布式配置，直接使用torch.nn.DataParallel即可实现模型并行：

model = Network(X.shape[-1])
# 指定使用4个CPU核心，也可省略device_ids让PyTorch自动使用所有可用CPU
model = torch.nn.DataParallel(model, device_ids=[0,1,2,3])

3. 额外性能优化（可选）

• 调整PyTorch底层运算线程数，让矩阵乘法等核心操作利用多线程：

torch.set_num_threads(4)  # 设置PyTorch内部运算线程数
torch.set_num_interop_threads(4)  # 设置跨操作调度线程数

• 增大batch_size：多CPU并行后，适当提升batch_size（如调整为128/256）可提高计算效率，同时建议将学习率同步放大（如乘以2）以保证训练效果。

修改后的完整可运行代码 #

import pandas as pd
import torch
from torch import nn
from torch.utils.data import Dataset, DataLoader

# 附录类定义
class CustomDataset(Dataset):
    def __init__(self, X, y):
        self.X = torch.Tensor(X.values)
        self.y = torch.Tensor(y.values)

    def __getitem__(self, index):
        return self.X[index], self.y[index]

    def __len__(self):
        return len(self.X)  # 修正原代码bug：原返回len(X)，应改为len(self.X)

class Network(nn.Module):
    def __init__(self, input_size):
        super(Network, self).__init__()
        self.input_size = input_size
        self.linear_1 = nn.Linear(self.input_size, 32)
        self.linear_2 = nn.Linear(32, 1)

    def forward(self, data):
        output = self.linear_1(data)
        output = self.linear_2(output)
        return output

if __name__ == '__main__':
    # 设置PyTorch线程数
    torch.set_num_threads(4)
    torch.set_num_interop_threads(4)

    df = pd.read_pickle('path/to/data')
    X = df.drop(columns=['target'])
    y = df[['target']]

    train_data = CustomDataset(X, y)
    train_loader = DataLoader(
        dataset=train_data,
        batch_size=64,
        num_workers=4,
        pin_memory=True
    )

    model = Network(X.shape[-1])
    # 启用多CPU并行
    model = torch.nn.DataParallel(model)

    criterion = nn.MSELoss()
    optimizer = torch.optim.Adam(model.parameters())

    epochs = 10

    for e in range(1, epochs + 1):
        train_loss = .0
        model.train()
        for batch_id, (data, labels) in enumerate(train_loader):
            optimizer.zero_grad()
            target = model(data)
            loss = criterion(target, labels)
            loss.backward()
            optimizer.step()
            train_loss += loss.item()

        print('Epoch {}: Train Loss: {:.4f}'.format(
            e,
            train_loss / len(train_loader)
        ))

关于DistributedDataParallel失效的说明 #

DistributedDataParallel针对分布式多节点/多GPU场景设计，单机器多CPU下需要额外初始化分布式环境（如设置rank、world_size），且需通过torch.distributed.launch启动脚本，配置复杂且无性能优势。单机器多CPU场景下，DataParallel是更简单高效的选择。

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