pytorch多机多卡

PyTorch支持多机多卡的训练，可以使用`torch.nn.DataParallel`或`torch.nn.DistributedDataParallel`来实现。这些模块可以帮助你在多个GPU或多台机器上进行并行计算。

使用`torch.nn.DataParallel`，你只需要将你的模型包装在这个模块中，并指定需要使用的GPU设备。例如：

import torch
import torch.nn as nn
from torch.nn.parallel import DataParallel

model = nn.Linear(10, 5)  # 一个简单的线性模型

# 将模型包装在DataParallel中，指定需要使用的GPU设备
model = DataParallel(model, device_ids=[0, 1])  # 在GPU 0和1上并行计算

input = torch.randn(20, 10)  # 输入数据
output = model(input)  # 并行计算输出

使用`torch.nn.DistributedDataParallel`，你需要使用`torch.distributed.launch`来启动多进程训练，并设置好分布式配置。例如：

import torch
import torch.nn as nn
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel

# 初始化分布式训练环境
dist.init_process_group(backend='nccl')

model = nn.Linear(10, 5)  # 一个简单的线性模型

# 将模型包装在DistributedDataParallel中
model = DistributedDataParallel(model)

input = torch.randn(20, 10)  # 输入数据
output = model(input)  # 分布式计算输出

以上是简单的示例代码，你可以根据你的具体需求进行修改和扩展。需要注意的是，多机多卡训练需要配置好分布式环境，并且确保每个GPU上的数据和模型参数都能正确同步。