总体流程）
前言
承上
在上一篇博客里面我们简单的实现了：
从处理数据集 -> 加载数据集 -> 编写神经网络 -> 编写损失函数和后处理代码 -> 编写训练和测试代码 单机单 gpu 进行网络进行训练，
启下
在使用多 gpu 进行训练上，pytorch 不像 tensorflow 方便，具体 tensorflow 如何使用多 gpu 的方法，请参考另外一片博客，这里不再赘述。
在 pytorch 使用多 gpu 进行训练需要我们将模型和数据手动设置在多个的 gpu 上面，可采用的模式也是多种多样，此片博客依照 DDP 模式进行 pytorh 的多 gpu 进行训练。阅读此博客之前，需要去了解 PyTorch分布式训练、DP、DDP、自动混合精度训练MAP等知识点。下面给出参考阅读大神连接：
1、什么是 DDP， DP： https://zhuanlan.zhihu.com/p/356967195
2、什么是自动混合精度训练MAP： https://www.cnblogs.com/jimchen1218/p/14315008.html
3、PyTorch分布式训练 ： https://zhuanlan.zhihu.com/p/360405558
这里给出上面的一些知识点的一些粗略的定义：
1、DP、DDP模式的差别：
1、DP模式简介：以一张卡为服务器卡，其余的客户端卡，服务器负责整理各个客户端的梯度和数据进行汇总，再将相同信息分发下去给每一个客户端卡，所以DP受限与服务器卡的带宽，通信成本随着客户端卡的数量的增加而线程增长(服务端卡一次只能和一张客户端卡进行通信)
2、DDP模式简介：所有卡地位相同，每张卡都要逐渐收集到所有信息，但是因为每张卡可以接受信息(卡1)也可以发送信息(卡2)，且只发送部分信息，可以并行执行只要有一张卡收集完信息，再将信息传递出去，基本上带宽不是问题，随着 gpu 的增加，通信成本与(N-1)/N 成正比，也就是增比很小)
3、SYNC_BN简介：是否使用sync_bn, DDP模式下可以使用。BN 在分布式训练的一种升级版，数据并行，每个gpu上的batchsize就会变小，BN效果就会差一点，是否使用sync_bn，将各显卡上分配的数据先汇总在一起，然后进行规范化处理。提升 BN 效果。
注意：Windows 下系统默认使用DP模式调用所有显卡，不支持 DDP
和单卡训练的差别
和单卡训练相比，整体的代码的变化幅度不是很多，主要变化有下面两点：
1、在 train.py 中的 网络加载 的地方和 数据 分发的地方。
2、在 utils_train中单独的一个 epoch_train 的训练函数中 是否使用回合双精度进行训练的地方
接下来，我们按照上面的两个区别，依次解析和单卡训练的不同点。
网络加载和数据分发区别
首先要采用多卡训练的话，肯定要确保你的电脑上存在多张 gpu，这一点是毋庸置疑的。通常情况下，多少张 gpu 是一样可以看出来的，但是，哎嗨，为就是想通过程序确定，也行。
1、查看自己有多少张 gpu，并记录下来。
# train.py
import torch
if __name__ == "__main__":
ngpus_per_node = torch.cuda.device_count()
print("当前设备中存在的 GPU 的数量为：{}".format(gpus_per_node))
可能有些人对 per_node 的后缀名有点不理解，这里简单说一下。一般情况下，我们是在单机上进行单卡或者多卡进行训练，但是也是有在多机器上进行多卡进行训练的。如果采用多机器进行训练的话，每一个机器为一个节点，也就是一个node。
2、设置通信后端。
在多张 gpu 上面进行训练的时候，会出现卡与卡之间的数据的分发和拷贝的工作发生，这对这个工作，pytorch 可以设置不同的数据分发拷贝的后端，具体使用的通信后端： ncc1->NVIDIA、gloo->Facebook、mpi->OpenMPI。
为了区别多机和多卡上程序的 🆔 的问题(为什么会有多程序？ 在 DDP 模式下，采用的为多线程模式，也就是一张卡的训练是一张线程程序)，我们需要获取当前主机对应的 🆔 和当前主机上当前卡上运行程序的 🆔。
将上面两个放在一起的原因: 因为都是在你配置为 分布式 if config.DISTRIBUTED: 训练的条件下才会进行动作，
代码如下：
# train.py
import torch.distributed as dist
# 是否使用多卡训练
if __name__ == "__main__":
if config.DISTRIBUTED:
# 使用的通信后端 ncc1->NVIDIA、gloo->Facebook、mpi->OpenMPI
if dist.is_nccl_available():
dist.init_process_group(backend="nccl")
elif dist.is_gloo_available():
dist.init_process_group(backend="mpi")
elif dist.is_mpi_available():
dist.init_process_group(backend="mpi")
else:
raise TypeError("==> 显卡没有支持的通信后端，请重新选择训练模式. ")
# 获取当前进程的 rank(类似 pid)
local_rank = int(os.environ["LOCAL_RANK"])
# 获取当前主机的 rank (多主机的情况，主机的 id 号码)
rank = int(os.environ["RANK"])
# 设置当前进程所使用的 gpu, 这样设置为一个 进程 对应一个 gpu，也可以不这样设置
device = torch.device("cuda", local_rank)
# 只在最开始的进程中打印下面的信息
if local_rank == 0:
print(f"==> [{os.getpid()}]-> rank = {rank}, local_rank = {local_rank} training...")
print("==> GPU Device Count : {}".format(ngpus_per_node))
else:
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
local_rank = 0
3、加载网络模式
做完上面的一些判断之后，我们加载我们的网络模式。就和单机单卡的加载模式相同，这里不再赘述。
代码如下：
#---------------------------------------#
# 加载自己的网络模型
#---------------------------------------#
model = YOLOv3(config.model_config, num_classes=config.NUM_CLASSES)
if local_rank == 0:
print("==> 模型加载完毕... ")
#---------------------------------------#
# 加载优化器
#---------------------------------------#
if config.OPTIMIZER_TYPE == "adam":
optimizer = optim.Adam(
model.parameters(), lr=config.INIT_LEARNING_RATE, betas=(config.MOMENTUM, 0.999), weight_decay=config.WEIGHT_DECAY
)
if config.OPTIMIZER_TYPE == "sgd":
optimizer = optim.SGD(
model.parameters(), lr=config.INIT_LEARNING_RATE, momentum=config.MOMENTUM, weight_decay=config.WEIGHT_DECAY, nesterov=True
)
if local_rank == 0:
print("==> 优化器加载完毕... ")
#---------------------------------------#
# 判断是否需要加载之前保存的模型和相应的参数
# 注意加载模型的时候，使用的为原始的没有被分布式的模型
#---------------------------------------#
if config.LOAD_MODEL:
checkpoint_file = config.LOAD_WEIGHT_NAME
load_checkpoint(checkpoint_file, model, optimizer, device, config.INIT_LEARNING_RATE)
if local_rank == 0:
print("==> 权重加载完毕...")
else:
if local_rank == 0:
print("==> 没有加载权重...")
#---------------------------------------#
# 加载损失函数
#---------------------------------------#
loss_fn = YoloLoss()
if local_rank == 0:
print("==> 损失函数加载完毕... ")
#---------------------------------------#
# 加载 损失 和 map 记录器
#---------------------------------------#
if local_rank == 0:
time_str = datetime.datetime.strftime(datetime.datetime.now(),'%Y_%m_%d_%H_%M_%S')
log_dir = os.path.join(config.SAVE_DIR, "loss/loss_" + str(time_str))
loss_history = Loss_Map_History(log_dir=log_dir, model=model, input_shape=config.IMAGE_SIZE)
print("==> 损失函数日志记载函数加载完毕... ")
else:
loss_history = None
#---------------------------------------#
# 加载先验框 3 * 3 * 2
#---------------------------------------#
scaled_anchors = (torch.tensor(config.ANCHORS)).cuda(local_rank)
if local_rank == 0:
print("==> 先验框加载完毕... ")
解释一下 if local_rank == 0: 的含义：local_rank 是当前线程的 🆔， local_rank == 0 表示只在第一个线程程序中打印相关的值，不加限制的话，将会打印多条语句。
4、设置训练使用的精度
这里先粘贴上代码，下面进行简单解释， 具体的解释连接可以参考上面的链接：
#---------------------------------------#
# 是否使用混合精度进行训练
#---------------------------------------#
if config.FP16:
# 创建一个 scaler 用于先放大梯度，再缩小
scaler = torch.cuda.amp.GradScaler()
else:
scaler = None
我们训练网络的时候，默认创建的 tensor 都是 FloatTensor 类型，FloatTensor 当然也具有好多精度的不同类型，但是默认我们创建的 tensor 的精度是 32 bit 的，也就是我们不加以限制的话，网络模型的所有参数的的精度都是 32位 FloatTensor 类型。
高精度的参数，当然意为着训练的精度更高，网络的性能更好，但也带来了训练参数数量级的增加，使得网络训练速度变慢。我们希望加快网络的训练的话，要尽可能在不损失网络性能的前提下，减少计算量。
还有一个前情知识点：NVIDIA研究了一种用于混合精度训练的方法，该方法在训练网络时将单精度（FP32）与半精度(FP16)结合在一起，并使用相同的超参数实现了与FP32几乎相同的精度。
可能还有人疑惑为什么不单独使用Float16 的呢？原因： １．溢出错误；２．舍入误差；
１．溢出错误： FP16的动态范围比FP32位的狭窄很多，因此，在计算过程中很容易出现上溢出(太大)和下溢出(太小)，溢出之后就会出现"NaN"的问题。在深度学习中，由于激活函数的梯度往往要比权重梯度小，更易出现下溢出的情况
2.舍入误差：舍入误差指的是当梯度过小时，小于当前区间内的最小间隔时，该次梯度更新可能会失败。也就是权重比梯度的数量级大好多，相加的时候，程序自动舍掉了梯度，导致权重没有更新。
所以我们希望使用 NVIDIA 的这种做法，就出现的我们的 AMP(自动混合精度)。
AMP 包含两个关键词有：1、自动，2、混合精度。
1、自动：Tensor的dtype类型会自动变化，框架按需自动调整 tensor 的 dtype,当然有些地方还需手动干预。
2、混合精度：采用不止一种精度的 Tensor，torch.FloatTensor 和 torch.HalfTensor
pytorch 1.6 版本以上的的包中包含了 torch.cuda.amp，是ＮVIDIA开发人员贡献到pytorch里的。只有支持tensor core的CUDA硬件才能享受到AMP带来的优势。Tensor core是一种矩阵乘累加的计算单元，每个 tensor core 时针执行64个浮点混合精度操作（FP16矩阵相乘和FP32累加）。
为了消除torch.HalfTensor也就是FP16的问题，需要使用以下两种方法：
1、混合精度训练：在内存中用FP16做储存和乘法从而加速计算，而用FP32做累加避免舍入误差。混合精度训练的策略有效地缓解了舍入误差的问题。
2、损失放大（Loss scaling)：即使了混合精度训练，还是存在无法收敛的情况，原因是激活梯度的值太小，造成了溢出(乘法运算)。可以通过使用torch.cuda.amp.GradScaler，通过放大loss的值来防止梯度的 underflow（只在BP时传递梯度信息使用，真正更新权重时还是要把放大的梯度再 unscale 回去）；
也就是进行下面两步操作：
反向传播前： 将损失变化手动增大2^k倍，因此反向传播时得到的中间变量（激活函数梯度）则不会溢出；
反向传播后： 将权重梯度缩小2^k倍，恢复正常值。
5、加载训练数据集
在这一步我们需要将数据先按照正常的单机单卡的加载模式加载数据，但是在创建 数据迭代器 的时候，我们需要将 数据跌打器 设置为分布式迭代器，可以让多线程之间，同时加载不同的数据。
5.1、自定义处理数据的方法;
在这一步我们实例化我们处理数据的一个类，这个类可以通过输入索引的方式，对输入路径下的图片进行处理，这就为下面创建 数据迭代器 奠定了基础。
#---------------------------------------#
# 1、加载数据和验证集的迭代对象
#
# 2、分发数据
#---------------------------------------#
train_annotaion_lines, val_annotation_lines = get_anno_lines(train_annotation_path=config.TRAIN_LABEL_DIR, val_annotation_path=config.VAL_LABEL_DIR)
train_dataset = Dataset_loader(annotation_lines=train_annotaion_lines, input_shape=config.IMAGE_SIZE, anchors=config.ANCHORS,
transform=train_transforms, train = True, box_mode="midpoint")
val_dataset = Dataset_loader(annotation_lines=val_annotation_lines, input_shape=config.IMAGE_SIZE, anchors=config.ANCHORS,
transform=val_transforms, train = True, box_mode="midpoint")
5.2、分发数据的设置
在创建 数据迭代器 之前，需要确定自己多卡训练迭代取数据的 抽样方式(sampler)，当然在单卡的情况下，sample 的方式也是可以自己的定义。
torch.utils.data.distributed.DistributedSampler 是适用于多卡训练的 sample 方式，能够满足多线程的情况下，在 gpu 上进行分布式训练。
batch_size = config.BATCH_SIZE // ngpus_per_node 是因为，我们有多张卡，我们的目的是为了加快我们的训练速度，你可以将训练的 batch 调大一点，然后分发到每个 gpu 上面，这样可以加大训练的 batch 和训练速度。
if config.DISTRIBUTED:
# 使用分布式的数据的加载方式
train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset, shuffle=True)
val_sampler = torch.utils.data.distributed.DistributedSampler(val_dataset, shuffle=True)
batch_size = config.BATCH_SIZE // ngpus_per_node
config.SHUFFLR = False
# 训练分布，则学习率也以当适应改变
lr_limit_max = 1e-3 if config.OPTIMIZER_TYPE == "adam" else 5e-2
lr_limit_min = 3e-4 if config.OPTIMIZER_TYPE == "adam" else 5e-4
Init_lr_train = min(max(batch_size / config.BATCH_SIZE * config.INIT_LEARNING_RATE, lr_limit_min), lr_limit_max)
Min_lr_train = min(max(batch_size / config.BATCH_SIZE * config.MIN_LEARNING_RATE, lr_limit_min * 1e-2 ), lr_limit_min * 1e-2)
# 通过最大最小学习率获得学习率下降公式
lr_scheduler_func = get_lr_scheduler(config.LEARNING_RATE_DECAY_TYPE, Init_lr_train, Min_lr_train, config.NUM_EPOCHS)
else:
batch_size = config.BATCH_SIZE
train_sampler = None
val_sampler = None
config.SHUFFLR = True
5.3、创建 数据迭代器
再下面和正常的单卡训练的方式一样，我们创建一个 数据迭代器，将我们上面定义的多卡分布数的数据处理和加载方式放入到数据迭代器中即可。
train_loader = DataLoader(train_dataset, batch_size, config.SHUFFLR, num_workers=config.NUM_WORKERS,
pin_memory=config.PIN_MEMORY, drop_last=False, sampler=train_sampler)
val_loader = DataLoader(val_dataset, batch_size, config.SHUFFLR, num_workers=config.NUM_WORKERS,
pin_memory=config.PIN_MEMORY, drop_last=False,sampler=val_sampler)
if local_rank == 0:
print("==> 数据集迭代器加载完毕...")
5.4、分布式网络模式设置
在创建分布式网络模型的过程中，我们还需要了解一个小知识点，关于 多卡同步 BN 也就是 SYNC_BN 的简介，参考博客上面给出的定义。
注意：
1、在设置分布式网络之前，一定要确定为训练模式。
2、多卡同步 BN 一定要有多张 gpu
3、DDP 相比 DP 模式，需要手动填写参数。
代码中添加了很多注释，这里也不在赘述。
#---------------------------------------#
# 在设置分布式网络之前，一定要确定为训练模式
#---------------------------------------#
model_train = model.train()
#---------------------------------------#
# 是否使用多卡同步 BN
#---------------------------------------#
if config.SYNC_BN and ngpus_per_node > 1 and config.DISTRIBUTED:
model_train = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model_train)
else:
print("==> Sync_bn is not support in one gpu or not set distributed. ")
#---------------------------------------#
# 分布网络模型
#---------------------------------------#
if config.GPU:
if config.DISTRIBUTED:
# DDP 分布模式
model_train = model_train.cuda(local_rank)
# find_unused_parameters 是用来解决定义在 forward 函数中，作用没有用的网络层，引发错误的问题
model_train = torch.nn.parallel.DistributedDataParallel(model_train, device_ids=[local_rank], find_unused_parameters=True)
if local_rank == 0:
print("==> 分布式网络部署成功，使用 DDP 模式进行训练. ")
else:
# DP 分布模式
model_train = torch.nn.DataParallel(model)
cudnn.benckmark = True
model_train = model_train.cuda()
if local_rank == 0:
print("==> 分布式网络部署成功，使用 DP 模式进行训练. ")
if local_rank == 0:
print("==> 进入 epochs, 开始训练... ")
5.5、训练设置
训练的时候，和单卡训练的方式也是类似的，但是需要同步和等待多张卡在相同的 epoch 上同时训练完成，才可以进行下一个 epoch 的训练。
for epoch in range( config.NUM_EPOCHS ):
#---------------------------------------#
# 将每个 samper 设置相同的 epoch
#---------------------------------------#
if config.DISTRIBUTED:
train_sampler.set_epoch(epoch)
#---------------------------------------#
# 主训练函数
#---------------------------------------#
train_mean_loss = epoch_train(train_loader, model_train, optimizer, loss_fn, scaler, scaled_anchors, epoch, config.NUM_EPOCHS, local_rank=local_rank)
if local_rank == 0:
#---------------------------------------#
# 验证集测试函数
#---------------------------------------#
val_mean_loss, mapval = epoch_eval_loss_map(val_loader, model_train, epoch, config.NUM_EPOCHS, scaled_anchors, loss_fn, config.CONF_THRESHOLD, config.NMS_IOU_THRESH, local_rank=local_rank )
#---------------------------------------#
# 记录网络的 损失 和 map
#---------------------------------------#
loss_history.append_loss(epoch, train_mean_loss, val_mean_loss, mapval)
#---------------------------------------#
# 保存网络的权重
#---------------------------------------#
if config.SAVE_MODEL:
if (epoch + 1) % config.WEIGHT_SAVE_PERIOD == 0 or epoch + 1 == config.NUM_EPOCHS:
filename = os.path.join(config.SAVE_DIR, "checkpoint/ep%03d-train_loss%.3f-val_loss%.3f.pth"% (epoch, train_mean_loss, val_mean_loss))
save_checkpoint(model=model, optimizer=optimizer, filename=filename)
elif len(loss_history.val_loss) <= 1 or (val_mean_loss) <= min(loss_history.val_loss):
print('Save current best model to best_epoch_weights.pth')
filename = os.path.join(config.SAVE_DIR, "checkpoint/best_epoch_weights.pth")
save_checkpoint(model=model, optimizer=optimizer, filename=filename)
else: # 不然就是最后一个epoch了，保存最后一个epoch
filename = os.path.join(config.SAVE_DIR, "checkpoint/last_epoch_weights.pth")
save_checkpoint(model=model, optimizer=optimizer, filename=filename)
#---------------------------------------#
# 设置新的学习率
#---------------------------------------#
set_optimizer_lr(optimizer, lr_scheduler_func, epoch)
#---------------------------------------#
# 等待所有分布式训练走到相同的位置
#---------------------------------------#
if config.DISTRIBUTED:
dist.barrier()
if local_rank == 0:
loss_history.writer.close()
注意：
1、在训练之前，要将所有的卡设置到相同的 epoch 上面，在训练的后面，要设置为等待所有的卡训练完成才进行下议论的训练。
2、所有的打印 、保存，验证集函数，应该放在 线程🆔 为 0 的上面的进行，避免多次运行。
主训练函数的差别
截止到目前，我们将 train.py 中相较于单卡训练的不同大概的讲述的一遍，还剩下主训练函数中的一些细节没有讲述。下面我们将主训练函数的细节进行讲解。
函数中的大部分都给予了注释，这里也不再赘述，
在梯度求导的过程中，如果使用双精度，就需要将整个求导过程放入 with torch.cuda.amp.autocast():中，并将损失函数先放大，再将梯度缩小，来减少双精度带来的影响。
# utils_train.py
def epoch_train(train_loader, model, optimizer, loss_fn, scaler, scaled_anchors, cur_epoch, all_epoch, local_rank=0):
#---------------------------------------#
# 在主程序加载进度条
#---------------------------------------#
if local_rank ==0:
pmgressbar_train = tqdm(train_loader, desc=f"Train epoch {cur_epoch + 1}/{all_epoch}", postfix=dict, mininterval=0.5)
# 确保网络模式为可训练模式
model.train()
# 记录损失
train_losses = []
#---------------------------------------#
# 迭代抽取数据
#---------------------------------------#
for iteration, (images, targets) in enumerate(train_loader):
# 将数据放到 gpu 上面
with torch.no_grad():
if config.GPU:
images = images.cuda(local_rank)
targets_cuda = [target.cuda(local_rank) for target in targets]
# 清除梯度优化器中的值
optimizer.zero_grad()
if config.FP16:
# 使用混合双精进行训练
with torch.cuda.amp.autocast():
out = model(images)
loss = (
loss_fn(out[0], targets_cuda[0], scaled_anchors[0])
+ loss_fn(out[1], targets_cuda[1], scaled_anchors[1])
+ loss_fn(out[2], targets_cuda[2], scaled_anchors[2])
)
train_losses.append(loss.item())
# 反向传播
# 先将 loss 放大
scaler.scale(loss).backward()
# 对权重求导
scaler.step(optimizer)
# 将梯度更新，缩小梯度信息
scaler.update()
else:
# 否则正常训练
out = model(images)
loss = (
loss_fn(out[0], targets_cuda[0], scaled_anchors[0])
+ loss_fn(out[1], targets_cuda[1], scaled_anchors[1])
+ loss_fn(out[2], targets_cuda[2], scaled_anchors[2])
)
loss.backward()
optimizer.step()
# 计算平均 loss，并更新进度条
train_mean_loss = sum(train_losses) / len(train_losses)
if local_rank == 0:
pmgressbar_train.set_postfix(**{'train_loss' : train_mean_loss,
'lr' : get_lr(optimizer)})
pmgressbar_train.update()
if local_rank == 0:
pmgressbar_train.close()
print("一个epoch的训练集的训练结束. ")
return train_mean_loss
训练启动设置
和以往确定好训练参数，直接运行训练函数不同，在分布式训练下，我们还需要额外的一些输入参数，我们需要在终端中输入下面的参数：
CUDA_VISIBLE_DIVECES=0,1 python -m torch.distributed.launch --nproc_per_node=2 train.py
0,1 是GPU编号，nproc_per_node是指定用了哪些GPU, torch.distributed.launch会调用这个local_rank，依据 local_rank 来分布式训练，DDP 模式的开启就需要这个 lanuch 加载文件。
如果感觉每次都在终端写上面的这么长的 参数 太麻烦，可以建立一个 脚本文件 .sh， 在其中填入上面的运行参数即可。如下：
# train.sh
conda init pytorch
CUDA_VISIBLE_DIVECES=0,1 python -m torch.distributed.launch --nproc_per_node=2 train.py
使用 torch.distributed.launch 的记载方式，在加载的过程中，也可以在我们自己的文件中 加入自己的定义的 输入参数，然后在终端后面加上自己的关键字。
import argparse
parser = argparse.ArgumentParser()
parser.add_argument("--local_rank", type=int)
args = parser.parse_args()
local_rank = args.local_rank
当然，要是你的 pytorch 在 1.9 版本之上，想采用默认参数即可，还可以采用下面这种 pytorch 集中好的运行参数方法：
torchrun --nnodes=1 --standalone --nproc_per_node=2 train.py
#nproc_per_node是你单机上的 gpu 的数量，和上面一样的在终端，在 .sh 文件里面写上直接运行也行。
具体的 分布式运行方式可以参考下面的官方链接：
https://pytorch.org/docs/stable/elastic/run.html
总结
在上一篇文章那个的基础之上，我们一起学习了如何从单卡训练转换到多卡训练的过程，需要改变的地方可能没有很多，但是需要了解一些额外的知识点。大家一起加油。
代码链接：
后续：
1、使用 pytorch 实现 两阶段的目标检测算法。
2、学习并概括大致的语义分割的流程，使用 pytorch 实现 最新的语义分割复现和学习。
3、学习并概括使用 tensorflow 上实现 yolo 系列的大致流程。
1517
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
