Python下多用户类实例并行训练PyTorch神经网络的实现方案

前置核心说明 #

• 必须采用多进程实现：Python GIL会限制CPU/GPU密集型的PyTorch训练并行，多线程无法实现真并行效果
• 多进程默认是复制主进程的对象到子进程，子进程对自身实例的修改不会自动同步回主进程，必须通过返回值主动同步参数
• 所有传入进程的user实例必须可被pickle序列化（Windows/macOS默认spawn启动模式要求），只要user类没有持有不可序列化的对象（如打开的文件句柄、运行中的锁等），原生实现即可满足要求
• 若使用GPU训练，需给每个进程分配独立的CUDA设备，避免显存冲突

快速改造实现（基于ProcessPoolExecutor） #

1. 依赖导入 #

import os
import torch
from concurrent.futures import ProcessPoolExecutor

2. 改造user的train方法，增加参数返回 #

原有训练逻辑全部保留，最后返回需要同步到主进程的训练后模型参数：

def train(self):
    self.current_model.train()
    optim = self.optim(self.current_model.parameters(), lr=self.lr)
    criterion = self.criterion()
    for epoch in range(self.epochs):
        loss_per_epoch = 0
        counter = 0
        for i, data in enumerate(self.dataloader):
            x, y = data
            fx = self.current_model(x.unsqueeze(1))
            loss = criterion(fx, y)
            optim.zero_grad()
            loss.backward(retain_graph=True)
            optim.step()
            loss_per_epoch += loss.item()
            counter += 1
    # 新增：返回训练后的模型参数，用于同步回主进程
    return self.current_model.state_dict()

3. 改造_train_users_locally实现并行调度 #

def _train_users_locally(self, max_workers=None):
    # 最大并行数默认取CPU物理核心数，可手动指定
    max_workers = max_workers or os.cpu_count()
    user_instances = list(self.users.values())

    with ProcessPoolExecutor(max_workers=max_workers) as executor:
        # 存future和对应user的映射，方便后续同步参数
        future_map = {}
        for user in user_instances:
            future = executor.submit(user.train)
            future_map[future] = user
        
        # 等待所有任务完成，同步训练后的参数到主进程的user实例
        for future, user in future_map.items():
            trained_state = future.result()
            user.current_model.load_state_dict(trained_state)

注意：Windows/macOS环境下必须加主入口保护，否则会递归启动进程：

if __name__ == "__main__":
    # 主程序逻辑（实例化管理类、加载users数据、调用训练方法等）
    manager = YourManagerClass()
    manager._train_users_locally()

GPU训练适配方案 #

如果模型运行在CUDA上，需要给每个进程分配独立GPU避免显存冲突，改造点如下：

1. 调整train方法，支持指定设备 #

def train(self, gpu_id=None):
    # 绑定当前进程对应的设备
    device = f"cuda:{gpu_id}" if gpu_id is not None and torch.cuda.is_available() else "cpu"
    self.current_model = self.current_model.to(device)

    self.current_model.train()
    optim = self.optim(self.current_model.parameters(), lr=self.lr)
    criterion = self.criterion()
    for epoch in range(self.epochs):
        loss_per_epoch = 0
        counter = 0
        for i, data in enumerate(self.dataloader):
            x, y = data
            # 把数据搬到对应设备
            x = x.unsqueeze(1).to(device)
            y = y.to(device)
            fx = self.current_model(x)
            loss = criterion(fx, y)
            optim.zero_grad()
            loss.backward(retain_graph=True)
            optim.step()
            loss_per_epoch += loss.item()
            counter += 1
    return self.current_model.state_dict()

2. 调度时传入对应GPU编号 #

def _train_users_locally(self, gpu_ids=None):
    # 无GPU时用CPU并行，有GPU时按照传入的GPU列表分配
    if gpu_ids is None or not torch.cuda.is_available():
        max_workers = os.cpu_count()
        gpu_ids = [None]
    else:
        max_workers = len(gpu_ids)

    user_instances = list(self.users.values())
    with ProcessPoolExecutor(max_workers=max_workers) as executor:
        future_map = {}
        for idx, user in enumerate(user_instances):
            # 轮询分配GPU
            assigned_gpu = gpu_ids[idx % len(gpu_ids)]
            future = executor.submit(user.train, gpu_id=assigned_gpu)
            future_map[future] = user
        
        for future, user in future_map.items():
            trained_state = future.result()
            user.current_model.load_state_dict(trained_state)

性能优化可选方案 #

• 可替换为torch.multiprocessing.Pool替代原生ProcessPoolExecutor，PyTorch对张量序列化做了底层优化，能大幅减少大模型并行时的内存开销，调用逻辑和上述实现完全一致
• 原有多进程下print输出会互相覆盖，建议替换为日志写入文件，或给每个进程的输出增加user标识
• 并行数不要超过物理核心数过多，否则上下文切换开销会抵消并行收益

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