适配样例(DDP场景(one NPU per process))
原生PyTorch框架的AMP支持3种多卡训练场景,NPU上适配的AMP当前仅支持一个进程在单个NPU上运行的DDP多卡训练场景。除了对样例代码添加AMP相关的改动外,还需要对样例代码进行必要的使能多卡场景改动。
导入AMP模块
导入AMP模块。进行多卡训练的初始化。
import time import torch import torch.nn as nn import torch_npu from torch_npu.npu import amp # 导入AMP模块 from torch.utils.data import Dataset, DataLoader import torchvision import os local_rank = int(os.environ["LOCAL_RANK"]) # 在shell脚本中循环传入local_rank变量作为指定的device device = torch.device('npu', local_rank) # local_rank用于自动获取device号 torch.distributed.init_process_group(backend="hccl",rank=local_rank) # 将通信方式设置为hccl torch_npu.npu.set_device(local_rank) ...... # 如样例代码所示,定义一个简单的神经网络
定义损失缩放参数
从torchvision中获取训练数据集,设置训练相关的参数batch_size、epochs,并在模型、优化器定义之后,定义AMP功能中的GradScaler。设置train_sampler并开启DDP模式。
# 数据集获取 train_data = torchvision.datasets.MNIST( root = 'mnist', download = True, train = True, transform = torchvision.transforms.ToTensor() ) train_sampler = torch.utils.data.distributed.DistributedSampler(train_data) #开启DDP模式 batch_size = 64 model = CNN().to(device) # 把模型放到指定NPU上 model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[local_rank], output_device=local_rank) # 开启DDP模式 train_dataloader = DataLoader(dataset=train_data, batch_size=batch_size, sampler=train_sampler) # 定义DataLoader loss_func = nn.CrossEntropyLoss().to(device) optimizer = torch.optim.SGD(model.parameters(), lr=0.1) # 定义优化器 scaler = amp.GradScaler() # 在模型、优化器定义之后,定义GradScaler epochs = 10 # 设置循环次数
适配AMP并训练
在训练代码中添加AMP功能相关的代码开启AMP。
for epoch in range(epochs): for imgs, labels in train_dataloader: imgs = imgs.to(device) # 把img数据放到指定NPU上 labels = labels.to(device) # 把label数据放到指定NPU上 with amp.autocast(): # 设置amp outputs = model(imgs) # 前向计算 loss = loss_func(outputs, labels) # 损失函数计算 optimizer.zero_grad() # 进行反向传播前后的loss缩放、参数更新 scaler.scale(loss).backward() # loss缩放并反向转播 scaler.step(optimizer) # 更新参数(自动unscaling) scaler.update() # 基于动态Loss Scale更新loss_scaling系数
执行以下命令拉起多卡训练。
torchrun --nproc_per_node=2 cnn_amp_ddp.py --standalone
脚本名称(cnn_amp_ddp.py)根据实际填写,--nproc_per_node参数表示多卡训练下使用卡的数量,--standalone参数表示使用独立模式运行,而不是与其他进程通信,避免运行卡死。
父主题: 混合精度适配(可选)