如何使用PyTorch multiprocessing独立训练多个神经网络模型

场景适配结论 #

torch.multiprocessing完全适配你的使用场景。你需要的是基于同一数据集并行运行多个互不干扰的训练任务，不需要进程间参数同步，不需要用到分布式训练的通信机制，直接用该模块启动独立子进程即可，能充分利用多核CPU、多GPU资源大幅降低整体训练耗时。

核心注意事项 #

• 启动方式显式指定为spawn，可规避Windows、macOS平台的兼容性问题
• 主进程提前初始化的数据集在Linux系统下会通过写时复制机制共享，不会重复占用内存；Windows系统会为每个子进程复制一份数据集，内存充足的情况下不影响使用
• 可根据硬件资源给每个子进程分配独立的GPU，避免资源竞争

简易可运行代码示例 #

import torch
import torch.multiprocessing as mp
from torch.utils.data import DataLoader, TensorDataset
import torch.nn as nn
import torch.optim as optim

# 自定义神经网络结构，可替换为你自己的模型
class SimpleMLP(nn.Module):
    def __init__(self, input_dim=10, hidden_dim=32, output_dim=1):
        super().__init__()
        self.layers = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, output_dim),
            nn.Sigmoid()
        )
    
    def forward(self, x):
        return self.layers(x)

# 单模型训练逻辑，每个子进程独立执行，参数完全不共享
def train_single_model(model_idx, dataset, device_id, epoch_num=10):
    device = torch.device(f"cuda:{device_id}" if torch.cuda.is_available() else "cpu")
    # 每个进程初始化独立的模型、优化器、损失函数
    model = SimpleMLP().to(device)
    criterion = nn.BCELoss()
    optimizer = optim.Adam(model.parameters(), lr=1e-3)
    dataloader = DataLoader(dataset, batch_size=32, shuffle=True)
    
    print(f"模型{model_idx}开始训练，运行设备: {device}")
    for epoch in range(epoch_num):
        total_loss = 0.0
        for x, y in dataloader:
            x, y = x.to(device), y.to(device)
            optimizer.zero_grad()
            pred = model(x)
            loss = criterion(pred, y)
            loss.backward()
            optimizer.step()
            total_loss += loss.item()
    
    # 训练完成后单独保存当前模型权重
    torch.save(model.state_dict(), f"./model_{model_idx}.pth")
    print(f"模型{model_idx}训练完成，权重已保存")

if __name__ == "__main__":
    # 显式指定进程启动方式，避免跨平台报错
    mp.set_start_method("spawn", force=True)
    
    # 主进程只加载一次公共数据集
    sample_num = 10000
    input_dim = 10
    x = torch.randn(sample_num, input_dim)
    y = torch.randint(0, 2, (sample_num, 1)).float()
    shared_dataset = TensorDataset(x, y)
    
    # 训练配置：训练4个模型，2块GPU每块跑2个任务，可根据自身硬件调整
    model_count = 4
    available_devices = [0, 0, 1, 1]
    
    # 批量启动子进程
    process_list = []
    for idx in range(model_count):
        device_id = available_devices[idx]
        p = mp.Process(
            target=train_single_model,
            args=(idx, shared_dataset, device_id)
        )
        p.start()
        process_list.append(p)
    
    # 等待所有训练任务完成
    for p in process_list:
        p.join()
    
    print("全部模型训练完成")

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