摘要
本例提取了植物幼苗数据集中的部分数据做数据集，数据集共有 12 种类别，模型使用最经典的 resnet50，演示如何实现混合精度训练以及如何使用 DDP 的方式实现多卡并行训练。
通过本文你和学到：
1、如何使用混合精度训练？
2、如何制作 ImageNet 数据集？
3、如何使用 DDP 方式的进行多卡训练？
4、如何使用 Mixup 数据增强。
5、如何进行多卡 BN 同步？
6、如何使用余弦退火调整学习率？
7、如何使用 classification_report 实现对模型的评价。
8、预测的两种写法。
apex
使用 apex 实现混合精度训练，具体安装方法见：
https://blog.csdn.net/hhhhhhhhhhwwwwwwwwww/article/details/120839608
DP 和 DDP
pytorch 中的有两种分布式训练方式，一种是常用的 DataParallel(DP)，另外一种是 DistributedDataParallel(DDP)，两者都可以用来实现数据并行方式的分布式训练，DP 采用的是 PS 模式，DDP 采用的是 ring-all-reduce 模式，两种分布式训练模式主要区别如下：
1、DP 是单进程多线程的实现方式，DDP 是采用多进程的方式。
2、DP 只能在单机上使用，DDP 单机和多机都可以使用。
3、DDP 相比于 DP 训练速度要快。
Parameter Server 架构(PS 模式)
Parameter Server 架构(PS 模式)由 server 节点和 worker 节点组成。
server 节点的主要功能是初始化和保存模型参数、接受 worker 节点计算出的局部梯度、汇总计算全局梯度，并更新模型参数(DP)。
worker 节点的主要功能是各自保存部分训练数据，初始化模型，从 server 节点拉取最新的模型参数(pull)，再读取参数，根据训练数据计算局部梯度，上传给 server 节点(push)。
详细的计算过程如下：
PS 模式下的 DP，会造成负载不均衡，因为充当 server 的 GPU 需要一定的显存用来保存 worker 节点计算出的局部梯度；另外 server 还需要将更新后的模型参数 broadcast 到每个 worker，server 的带宽就成了 server 与 worker 之间的通信瓶颈，server 与 worker 之间的通信成本会随着 worker 数目的增加而线性增加。
ring-all-reduce 模式
ring-all-reduce 模式没有 server 节点，worker 与 worker 之间的通信构成一个环。
ring-all-reduce 模式下，所有 worker 只和自己相邻的两个 worker 进行通信，该工作模式分为两个工作阶段：
Scatter Reduce：在这个 Scatter Reduce 阶段，GPU 会逐步交换彼此的梯度并融合，最后每个 GPU 都会包含完整融合梯度的一部分
Allgather：GPU 会逐步交换彼此不完整的融合梯度，最后所有 GPU 都会得到完整的融合梯度
计算过程如下：
DDP 的基本用法 (代码编写流程)
使用 torch.distributed.init_process_group 初始化进程组
使用 torch.nn.parallel.DistributedDataParallel 创建 分布式模型
使用 torch.utils.data.distributed.DistributedSampler 创建 DataLoader
调整其他必要的地方(tensor 放到指定 device 上，S/L checkpoint，指标计算等)
使用 torch.distributed.launch / torch.multiprocessing 或 slurm 开始训练
Mixup
为了提高成绩我在代码中加入 Mixup 这种增强方式。使用到了 timm，安装命令：
导入包：from timm.data.mixup import Mixup，
定义 Mixup，和 SoftTargetCrossEntropy
mixup_fn = Mixup( mixup_alpha=0.8, cutmix_alpha=1.0, cutmix_minmax=None, prob=0.1, switch_prob=0.5, mode='batch', label_smoothing=0.1, num_classes=12) criterion_train = SoftTargetCrossEntropy()
复制代码
项目结构
resnet_demo├─data│ ├─Black-grass│ ├─Charlock│ ├─Cleavers│ ├─Common Chickweed│ ├─Common wheat│ ├─Fat Hen│ ├─Loose Silky-bent│ ├─Maize│ ├─Scentless Mayweed│ ├─Shepherds Purse│ ├─Small-flowered Cranesbill│ └─Sugar beet├─mean_std.py├─makedata.py├─train.py├─test1.py└─test.py
复制代码
mean_std.py：计算 mean 和 std 的值。
makedata.py：生成数据集。
计算 mean 和 std
为了使模型更加快速的收敛，我们需要计算出 mean 和 std 的值，新建 mean_std.py,插入代码：
from torchvision.datasets import ImageFolderimport torchfrom torchvision import transforms
def get_mean_and_std(train_data): train_loader = torch.utils.data.DataLoader( train_data, batch_size=1, shuffle=False, num_workers=0, pin_memory=True) mean = torch.zeros(3) std = torch.zeros(3) for X, _ in train_loader: for d in range(3): mean[d] += X[:, d, :, :].mean() std[d] += X[:, d, :, :].std() mean.div_(len(train_data)) std.div_(len(train_data)) return list(mean.numpy()), list(std.numpy())
if __name__ == '__main__': train_dataset = ImageFolder(root=r'data1', transform=transforms.ToTensor()) print(get_mean_and_std(train_dataset))
复制代码
数据集结构：
运行结果：
([0.3281186, 0.28937867, 0.20702125], [0.09407319, 0.09732835, 0.106712654])
把这个结果记录下来，后面要用！
生成数据集
我们整理还的图像分类的数据集结构是这样的
data├─Black-grass├─Charlock├─Cleavers├─Common Chickweed├─Common wheat├─Fat Hen├─Loose Silky-bent├─Maize├─Scentless Mayweed├─Shepherds Purse├─Small-flowered Cranesbill└─Sugar beet
复制代码
pytorch 和 keras 默认加载方式是 ImageNet 数据集格式，格式是
├─data│ ├─val│ │ ├─Black-grass│ │ ├─Charlock│ │ ├─Cleavers│ │ ├─Common Chickweed│ │ ├─Common wheat│ │ ├─Fat Hen│ │ ├─Loose Silky-bent│ │ ├─Maize│ │ ├─Scentless Mayweed│ │ ├─Shepherds Purse│ │ ├─Small-flowered Cranesbill│ │ └─Sugar beet│ └─train│ ├─Black-grass│ ├─Charlock│ ├─Cleavers│ ├─Common Chickweed│ ├─Common wheat│ ├─Fat Hen│ ├─Loose Silky-bent│ ├─Maize│ ├─Scentless Mayweed│ ├─Shepherds Purse│ ├─Small-flowered Cranesbill│ └─Sugar beet
复制代码
新增格式转化脚本 makedata.py,插入代码：
import globimport osimport shutil
image_list=glob.glob('data1/*/*.png')print(image_list)file_dir='data'if os.path.exists(file_dir): print('true') #os.rmdir(file_dir) shutil.rmtree(file_dir)#删除再建立 os.makedirs(file_dir)else: os.makedirs(file_dir)
from sklearn.model_selection import train_test_split
trainval_files, val_files = train_test_split(image_list, test_size=0.3, random_state=42)
train_dir='train'val_dir='val'train_root=os.path.join(file_dir,train_dir)val_root=os.path.join(file_dir,val_dir)for file in trainval_files: file_class=file.replace("\\","/").split('/')[-2] file_name=file.replace("\\","/").split('/')[-1] file_class=os.path.join(train_root,file_class) if not os.path.isdir(file_class): os.makedirs(file_class) shutil.copy(file, file_class + '/' + file_name)
for file in val_files: file_class=file.replace("\\","/").split('/')[-2] file_name=file.replace("\\","/").split('/')[-1] file_class=os.path.join(val_root,file_class) if not os.path.isdir(file_class): os.makedirs(file_class) shutil.copy(file, file_class + '/' + file_name)
复制代码
训练
完成上面的步骤后，就开始 train 脚本的编写，新建 train.py.
导入项目使用的库
import torchimport torch.nn as nnimport torch.nn.parallelimport torch.optim as optimimport torch.utils.dataimport torch.utils.data.distributedimport torchvision.datasets as datasetsimport torchvision.transforms as transformsfrom sklearn.metrics import classification_reportfrom timm.data.mixup import Mixupfrom timm.loss import SoftTargetCrossEntropyfrom torchvision.models.resnet import resnet50from apex import ampimport torch.distributed as distfrom torch.nn.parallel import DistributedDataParallel as DDPimport osfrom apex.parallel import convert_syncbn_model
复制代码
设置全局参数
设置学习率、BatchSize、epoch 等参数。
# 设置全局参数model_lr = 1e-4BATCH_SIZE = 256EPOCHS = 1000use_amp=False #是否使用混合精度classes=12CLIP_GRAD=5.0is_distributed=True
复制代码
model_lr：学习率，根据实际情况做调整。
BATCH_SIZE：batchsize，根据显卡的大小设置。
EPOCHS：epoch 的个数，一般 300 够用。
use_amp：是否使用混合精度。
classes：类别个数。
CLIP_GRAD：梯度的最大范数，在梯度裁剪里设置。
rank：默认是 0。
is_distributed：是否是分布式？
设置 distributed
使用 nccl 的方式初始化初始化进程组。
# 0. set up distributed devicerank = int(os.environ["RANK"])local_rank = int(os.environ["LOCAL_RANK"])torch.cuda.set_device(rank % torch.cuda.device_count())dist.init_process_group(backend="nccl")device = torch.device("cuda", local_rank)print(f"[init] == local rank: {local_rank}, global rank: {rank} ==")
复制代码
进程组的参数介绍：
GROUP：进程组，大部分情况下 DDP 的各个进程是在同一个进程组下。
WORLD_SIZE：总的进程数量 (原则上一个 process 占用一个 GPU 是较优的)。
RANK：当前进程的序号，用于进程间通讯，rank = 0 的主机为 master 节点。
LOCAL_RANK：当前进程对应的 GPU 号。
图像预处理与增强
数据处理比较简单，加入了 Cutout、做了 Resize 和归一化，定义 Mixup 函数。
# 数据预处理7transform = transforms.Compose([ transforms.Resize((224, 224)), transforms.ToTensor(), transforms.Normalize(mean=[0.51819474, 0.5250407, 0.4945761], std=[0.24228974, 0.24347611, 0.2530049])
])transform_test = transforms.Compose([ transforms.Resize((224, 224)), transforms.ToTensor(), transforms.Normalize(mean=[0.51819474, 0.5250407, 0.4945761], std=[0.24228974, 0.24347611, 0.2530049])])mixup_fn = Mixup( mixup_alpha=0.8, cutmix_alpha=1.0, cutmix_minmax=None, prob=0.1, switch_prob=0.5, mode='batch', label_smoothing=0.1, classes=12)
复制代码
读取数据
使用 pytorch 默认读取数据的方式，然后将 dataset_train.class_to_idx 打印出来，预测的时候要用到。
# 读取数据train_dataset = datasets.ImageFolder('data/train', transform=transform)train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset,shuffle=True)val_dataset = datasets.ImageFolder("data/val", transform=transform_test)print(train_dataset.class_to_idx)# 导入数据train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=BATCH_SIZE, sampler=train_sampler)test_loader = torch.utils.data.DataLoader(val_dataset, batch_size=BATCH_SIZE, shuffle=False)
复制代码
设置模型
设置 loss 函数，train 的 loss 为：SoftTargetCrossEntropy，val 的 loss：nn.CrossEntropyLoss()。
设置模型为 resnet50，预训练设置为 true，num_classes 设置为 12。
使用 convert_sync_batchnorm 函数实现多卡之间的 BN 同步。
创建 DDP 方式的多卡训练。
优化器设置为 adam。
学习率调整策略选择为余弦退火。
如果使用混合精度，则将 amp 初始化为“O1”。
# 实例化模型并且移动到GPUcriterion_train = SoftTargetCrossEntropy()criterion_val = torch.nn.CrossEntropyLoss()model_ft = resnet50(pretrained=True)num_ftrs = model_ft.fc.in_featuresmodel_ft.fc = nn.Linear(num_ftrs, classes)#model_ft=convert_syncbn_model(model_ft)#使用apex同步BNmodel_ft=torch.nn.SyncBatchNorm.convert_sync_batchnorm(model_ft)model_ft.to(device)print(model_ft)
# 选择简单暴力的Adam优化器，学习率调低optimizer = optim.Adam(model_ft.parameters(), lr=model_lr)cosine_schedule = optim.lr_scheduler.CosineAnnealingLR(optimizer=optimizer, T_max=20, eta_min=1e-6)if use_amp: model_ft, optimizer = amp.initialize(model_ft, optimizer, opt_level="O1") # 这里是“欧一”，不是“零一”# DistributedDataParallelmodel_ft = DDP(model_ft, device_ids=[local_rank], output_device=local_rank)
复制代码
定义训练和验证函数
定义训练函数和验证函数。在训练函数中：
首先 ，调用 set_epoch，每个 epoch 开始时调用 set_epoch() 方法，然后再创建 DataLoader 迭代器。
然后，初始化 loss，开始遍历 train_loader。
判断 batch 中有没有奇数的情况，如果有则减去一因为 MixUp 的 loss 需要偶数才行。
将数据放入转为 cuda，计算 mixup。
将 data 输入 model 得到输出结果，然后计算 loss。
判断是否使用混合精度，如果使用则调用 scaled_loss.backward()没有则调用 loss.backward()。
clip_grad_norm_()执行梯度裁剪，防止梯度爆炸。
在一个 epoch 完成后，使用 classification_report 计算详细的得分情况。
# 定义训练过程def train(model, device, train_loader, optimizer, epoch): if rank == 0: print(" ======= Training ======= \n") train_sampler.set_epoch(epoch) model.train() sum_loss = 0 total_num = len(train_loader.dataset) print(total_num, len(train_loader)) for batch_idx, (data, target) in enumerate(train_loader): if len(data)%2!=0: print(len(data)) data=data[0:len(data)-1] target=target[0:len(target)-1] print(len(data)) data, target = data.to(device, non_blocking=True), target.to(device, non_blocking=True) samples, targets = mixup_fn(data, target) output = model(samples) loss = criterion_train(output, targets) optimizer.zero_grad() if use_amp: with amp.scale_loss(loss, optimizer) as scaled_loss: scaled_loss.backward() grad_norm = torch.nn.utils.clip_grad_norm_(amp.master_params(optimizer), CLIP_GRAD) else: loss.backward() grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(),CLIP_GRAD) #torch.nn.utils.clipgrad_norm() 的使用应该在loss.backward() 之后，optimizer.step()之前. #注意这个方法只在训练的时候使用，在测试的时候验证和测试的时候不用。
optimizer.step() lr = optimizer.state_dict()['param_groups'][0]['lr'] print_loss = loss.data.item() sum_loss += print_loss if (batch_idx + 1) % 10 == 0: print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}\tLR:{:.9f}'.format( epoch, (batch_idx + 1) * len(data), len(train_loader.dataset), 100. * (batch_idx + 1) / len(train_loader), loss.item(), lr)) ave_loss = sum_loss / len(train_loader) print('epoch:{},loss:{}'.format(epoch, ave_loss))ACC = 0# 验证过程def val(model, device, test_loader): global ACC model.eval() test_loss = 0 correct = 0 total_num = len(test_loader.dataset) print(total_num, len(test_loader)) val_list = [] pred_list = [] with torch.no_grad(): for data, target in test_loader: for t in target: val_list.append(t.data.item()) data, target = data.to(device), target.to(device) output = model(data) loss = criterion_val(output, target) _, pred = torch.max(output.data, 1) for p in pred: pred_list.append(p.data.item()) correct += torch.sum(pred == target) print_loss = loss.data.item() test_loss += print_loss correct = correct.data.item() acc = correct / total_num avgloss = test_loss / len(test_loader) print('\nVal set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format( avgloss, correct, len(test_loader.dataset), 100 * acc)) if acc > ACC: if isinstance(model, torch.nn.parallel.DistributedDataParallel): torch.save(model.module, 'model_' + str(epoch) + '_' + str(round(acc, 3)) + '.pth') else: torch.save(model, 'model_' + str(epoch) + '_' + str(round(acc, 3)) + '.pth') ACC = acc return val_list, pred_list
# 训练is_set_lr = Falsefor epoch in range(1, EPOCHS + 1): train(model_ft, device, train_loader, optimizer, epoch) if epoch < 600: cosine_schedule.step() else: if is_set_lr: continue for param_group in optimizer.param_groups: param_group["lr"] = 1e-6 is_set_lr = True val_list, pred_list = val(model_ft, device, test_loader) print(classification_report(val_list, pred_list, target_names=train_dataset.class_to_idx))
复制代码
然后在命令中执行命令:
python -m torch.distributed.launch --nproc_per_node=2 --nnodes=1 --node_rank=0 --master_addr=localhost --master_port=22222 train.py
复制代码
nproc_per_node：表示每天机器上的显卡个数。
nnodes：节点个数。
node_rank：节点的 rank。
master_addr：主节点的 IP 地址。
master_port：主节点的端口号。
更多的设置方式如下：
例1: 1 node, 4 GPUs per node (4GPUs)>>> python -m torch.distributed.launch \ --nproc_per_node=4 \ --nnodes=1 \ --node_rank=0 \ --master_addr=localhost \ --master_port=22222 \ mnmc_ddp_launch.py
[init] == local rank: 3, global rank: 3 ==[init] == local rank: 1, global rank: 1 ==[init] == local rank: 0, global rank: 0 ==[init] == local rank: 2, global rank: 2 == 例2: 1 node, 2tasks, 4 GPUs per task (8GPUs)>>> CUDA_VISIBLE_DEVICES=0,1,2,3 python -m torch.distributed.launch \ --nproc_per_node=4 \ --nnodes=2 \ --node_rank=0 \ --master_addr="10.198.189.10" \ --master_port=22222 \ mnmc_ddp_launch.py
>>> CUDA_VISIBLE_DEVICES=4,5,6,7 python -m torch.distributed.launch \ --nproc_per_node=4 \ --nnodes=2 \ --node_rank=1 \ --master_addr="10.198.189.10" \ --master_port=22222 \ mnmc_ddp_launch.py
例3: 2 node, 8 GPUs per node (16GPUs)>>> python -m torch.distributed.launch \ --nproc_per_node=8 \ --nnodes=2 \ --node_rank=0 \ --master_addr="10.198.189.10" \ --master_port=22222 \ mnmc_ddp_launch.py
>>> python -m torch.distributed.launch \ --nproc_per_node=8 \ --nnodes=2 \ --node_rank=1 \ --master_addr="10.198.189.10" \ --master_port=22222 \ mnmc_ddp_launch.py
复制代码
注意：分布式训练只能在命令行中启动。
运行结果：
测试
我介绍两种常用的测试方式，第一种是通用的，通过自己手动加载数据集然后做预测，具体操作如下：
测试集存放的目录如下图：
第一步 定义类别，这个类别的顺序和训练时的类别顺序对应，一定不要改变顺序！！！！
第二步 定义 transforms，transforms 和验证集的 transforms 一样即可，别做数据增强。
第三步 加载 model，并将模型放在 DEVICE 里，
第四步 读取图片并预测图片的类别，在这里注意，读取图片用 PIL 库的 Image。不要用 cv2，transforms 不支持。
import torch.utils.data.distributedimport torchvision.transforms as transformsfrom PIL import Imagefrom torch.autograd import Variableimport osclasses = ('Black-grass', 'Charlock', 'Cleavers', 'Common Chickweed', 'Common wheat','Fat Hen', 'Loose Silky-bent', 'Maize','Scentless Mayweed','Shepherds Purse','Small-flowered Cranesbill','Sugar beet')transform_test = transforms.Compose([ transforms.Resize((224, 224)), transforms.ToTensor(), transforms.Normalize(mean=[0.51819474, 0.5250407, 0.4945761], std=[0.24228974, 0.24347611, 0.2530049])]) DEVICE = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")model = torch.load("model_2_0.418.pth",map_location='cuda:0')model.eval()model.to(DEVICE) path='data/test/'testList=os.listdir(path)for file in testList: img=Image.open(path+file) img=transform_test(img) img.unsqueeze_(0) img = Variable(img).to(DEVICE) out=model(img) # Predict _, pred = torch.max(out.data, 1) print('Image Name:{},predict:{}'.format(file,classes[pred.data.item()]))
复制代码
运行结果：
第二种 使用自定义的 Dataset 读取图片
import torch.utils.data.distributedimport torchvision.transforms as transformsfrom PIL import Imageimport os classes = ('Black-grass', 'Charlock', 'Cleavers', 'Common Chickweed', 'Common wheat','Fat Hen', 'Loose Silky-bent', 'Maize','Scentless Mayweed','Shepherds Purse','Small-flowered Cranesbill','Sugar beet')transform_test = transforms.Compose([ transforms.Resize((224, 224)), transforms.ToTensor(), transforms.Normalize(mean=[0.51819474, 0.5250407, 0.4945761], std=[0.24228974, 0.24347611, 0.2530049])]) DEVICE = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")model = torch.load("model.pth")model.eval()model.to(DEVICE) dataset_test =SeedlingData('data/test/', transform_test,test=True)print(len(dataset_test))# 对应文件夹的label for index in range(len(dataset_test)): item = dataset_test[index] img, label = item img.unsqueeze_(0) data = Variable(img).to(DEVICE) output = model(data) _, pred = torch.max(output.data, 1) print('Image Name:{},predict:{}'.format(dataset_test.imgs[index], classes[pred.data.item()])) index += 1
复制代码
运行结果：
完整的代码：https://download.csdn.net/download/hhhhhhhhhhwwwwwwwwww/85143391
评论