参考：

一、数据并行

单机单卡

基于torch.cuda.is_avaliable()来判断是否可用
可用GPU数量：gpus=torch.cuda.device_count()
限制GPU卡的使用：命令行CUDA_VISIBLE_DEVICES="0"/代码os.environ["CUDA_VISIBLE_DEVICES"]="0"
在训练过程，数据加载到GPU中：data=data.cuda() target=target.cuda()
模型保存和加载：

方式一：保存和加载模型参数和结构。
- 保存：模型训练完了，可以保存为.pt/.pth文件。torch.save(model, './model.pt')
- 加载：将本地的模型文件进行加载。model=torch.load('./model.pt', map_location=torch.device("cuda"/"cuda:index"/"cpu"))
方式二：只保存和加载模型参数。
- 保存：torch.save(model.state_dict(),'./model.pt')
- 加载：model.load_state_dict(torch.load('./model.pt', map_location=torch.device("cuda"/"cuda:index"/"cpu")))

当model.load_state_dict加载完模型后，模型拷贝到GPU上：model.cuda()

data/model.to(device) 与 data/model.cuda()两种方式效果完全等同。

单机多卡

DP：torch.nn.DataParallel（已经被淘汰）

基于torch.cuda.is_avaliable()来判断是否可用
可用GPU数量：gpus=torch.cuda.device_count()
设置GPU卡的使用：命令行CUDA_VISIBLE_DEVICES="0,1,2"/代码os.environ["CUDA_VISIBLE_DEVICES"]="0,1,2"
模型拷贝并放入DataParallel：model=torch.nn.DataParallel(model.cuda(), device_id=[0,1,2,3])
训练阶段，数据拷贝至GPU：data.cuda()
保存：模型训练完了，可以保存为.pt/.pth文件。torch.save(model, './model.pt')
加载：将本地的模型文件进行加载。model=torch.load('./model.pt', map_location=torch.device("cuda:0"))
模型加载完，模型拷贝至GPU：model.cuda()

模型加载中的参数 map_loaction：这个参数是将训练好的模型在加载时指定加载到哪一块上。

假设我们在cuda:0上训练好了一个模型。

（1）若不加这个参数，则默认加载到第0块GPU上：

加载模型：model=torch.load('./model.pt')

加载到哪块设备上？：print(next(model.parameters()).device)
结果为：cuda:0

（2）若指定加载到某块GPU上，则应为：

加载模型：model=torch.load('./model.pt',map_loaction={'cuda:0':'cuda:1'})

加载到哪块设备上？：print(next(model.parameters()).device)
结果为：cuda:1

（3）若指定错误：

加载模型：model=torch.load('./model.pt',map_loaction={'cuda:2':'cuda:3'})

由于模型在cuda:0上训练的，因此，将’cuda:2’上训练的模型加载到’cuda:3’上是不准确的。

（4）经常报的错误：RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. 多卡训练时将其他卡上训练的模型加载到卡 0上。

？？？？是不是模型训练可以多卡，但是模型加载只能加载到一张卡上。

解决方法：
model=torch.load('./model.pt',map_loaction={'cuda:1':'cuda:0'})

DDP：torch.nn.parallel.DistributeDataParallel

基本概念：

node：物理节点，1个node可以是一台机器，节点内部可以有多个gpu；nnodes是指物理节点数量；nproc_per_node是指每个物理节点上面的进程数量。
rank：所有node上，全局进程进行编号，rank 0为第一个进程或主进程，其余的为rank 1 2 3… torch.distributed.get_rank()
local_rank：当前node上，进程进行编号。torch.distributed.get_local_rank()
world_size：所有node上，全局进程总个数，即一个分布式任务中rank的数量。torch.distributed.get_world_size()
group：进程组，1个分布式任务就是1个group；这个可以先不用管。

我们运行一个分布式任务在3台机子上（你简单理解为一个多卡运行的YOLOv8的main.py程序）：咱3台机子=3个node，nnodes=3；每个node上有4个进程，每个node的进程数=nproc_per_node=4，全局进程总数=world_size=12，rank号=0,1,2,…11，每个node上local_rank=0,1,2,3。

注：1. 我强调了一个分布式任务，因为运行其他分布式任务，nproc_per_node\world_size\rank\local_rank各个分布式任务有各自的编号。
2. local_rank与GPU其实并没有太大关系，但是我们是可以看作local_rank对应为GPU的编号。

图片.png

配置local_rank参数：这个参数是torch.distributed.launch传递过来的，我们设置位置参数来接受，也可以通过torch.distributed.get_rank()获取进程id，两种方法都可以。

parser = argparse.ArgumentParser()

parser.add_argument('--local_rank', default=-1, type=int)

args = parser.parse_args()

local_rank = args.local_rank

初始化进程组：

torch.distributed.init_process_group("ucci", world_size=n_gpus, rank=args.local_rank)

“ucci”：GPU之间的通信方式。

设置GPU卡的使用（相当于CUDA_VISIBLE_DEVICES=”0,1,2″）：

local_rank = torch.distributed.get_rank()

# 配置每个进程的 GPU，根据local_rank来设定当前使用哪块GPU

torch.cuda.set_device(local_rank)

device = torch.device("cuda", local_rank)

模型拷贝并放入DistributeDataParallel：

model=torch.nn.parallel.DistributeDataParallel(model.cuda(args.local_rank), device_ids=[args.local_rank])

设置分布式采样器，将一个batch的数据分配到不同的GPU上，作为dataloader的参数：

train_sample=DistributeSampler(train_dataset)

创建dataloader：

train_dataloader=DataLoader(..., pin_memory=Ture, num_workers=args.workers, sampler=train_sampler)

注：1. num_workers=args.workers，数据集加载的时候，控制用于同时加载数据的线程数（默认为0，即在主线程读取）。2. 建议开启pin_memory=Ture，这样速度会更快，但如果开启后发现有卡顿现象或者内存占用过高，此时建议关闭。

将BN层替换成SyncBatchNorm：

model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model)

BatchNorm之类的层在其计算中使用了整个批次统计信息，因此无法仅使用一部分批次在每个GPU上独立进行操作。 PyTorch提供SyncBatchNorm作为BatchNorm的替换/包装模块，该模块使用跨GPU划分的整个批次计算批次统计信息。

训练阶段，数据拷贝至GPU：data=data.cuda(args.local_rank)
保存：torch.save(model.module.state_dict(), './model.pt')
加载：model.load_state_dict(torch.load('./model.pt', map_location=torch.device("cuda:0")))
命令行：python -m torch.distributed.launch --nproc_per_node=n_gpus train.py
在训练epoch那一层for中添加，为了让每张卡在每个周期中得到的数据是随机的，防止loss周期性增大：train_sample.set_epoch(epoch_index)

DDP：torch.nn.parallel.DistributeDataParrel单机多卡代码如下：

import torch
import torch.nn as nn
import argparse
from torch.autograd import Variable
from torch.utils.data import Dataset, DataLoader
import os
from torch.utils.data.distributed import DistributedSampler

# 1)添加参数  --local_rank
'''
每个进程分配一个 local_rank 参数, 表示当前进程在当前主机上的编号。这个参数是torch.distributed.launch传递过来的，我们设置位置参数来接受，local_rank代表当前程序进程使用的GPU标号
'''
parser = argparse.ArgumentParser()
parser.add_argument('--local_rank', default=-1, type=int,
                    help='node rank for distributed training')
args = parser.parse_args()

# 通过args接收 local_rank
local_rank = args.local_rank

# 通过 get_rank() 得到 local_rank，最好放在初始化之后使用
local_rank = torch.distributed.get_rank()

# 2)********************************************************
# 初始化使用nccl后端
torch.distributed.init_process_group(backend="nccl")

# 3)********************************************************
# 配置device
local_rank = torch.distributed.get_rank()
#配置每个进程的 GPU， 根据local_rank来设定当前使用哪块GPU
torch.cuda.set_device(local_rank)
device = torch.device("cuda", local_rank)

# 4)********************************************************
# 自己的数据获取
dataset = MyDataset(input_size, data_size)

# 使用 DistributedSampler
'''
使用 DistributedSampler 对数据集进行划分。它能帮助我们将每个 batch 划分成几个 partition，在当前进程中只需要获取和 rank 对应的那个 partition 进行训练
#注意 testset不用sampler
'''
train_sampler = torch.utils.data.distributed.DistributedSampler(dataset)

trainloader = DataLoader(dataset=dataset,
                         pin_memory=true,
                         shuffle=(train_sampler is None),   # 使用分布式训练 shuffle 应该设置为 False
                         batch_size=args.batch_size,
                         num_workers=args.workers,
                         sampler=train_sampler)



# 5)********************************************************
model = Model()
# 把模型移到对应的gpu
# 定义并把模型放置到单独的GPU上，需要在调用`model=DDP(model)`前做
model.to(device)

# 引入SyncBN，这句代码，会将普通BN替换成SyncBN。
model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model)

# GPU 数目大于 1 才有必要分布式训练
if torch.cuda.device_count() > 1:
    model = torch.nn.parallel.DistributedDataParallel(model,
                                                      device_ids=[local_rank],
                                                      output_device=local_rank)


# 6)********************************************************

for epoch in range(num_epochs):
    # 设置sampler的epoch，DistributedSampler需要这个来维持各个进程之间的相同随机数种子
    trainloader.sampler.set_epoch(epoch)
    # 后面这部分，则与原来完全一致了。
    for data, label in trainloader:
        prediction = model(data)
        loss = loss_fn(prediction, label)
        loss.backward()
        optimizer = optim.SGD(ddp_model.parameters(), lr=0.001)
        optimizer.step()

多机多卡

DDP：torch.nn.parallel.DistributeDataParallel

代码编写流程与单机多卡一致
命令行：python -m torch.distributed.lanch -- nprcc_per_node=n_gpus --nnoddes=2 --node_rank=0 --master_addr="主节点IP" --master_port="主节点端口号" train.py
模型保存与加载与单机多卡一致