modelzoo：昇腾 NPU 的“模型仓库”

modelzoo昇腾 NPU 的“模型仓库”之前帮朋友看模型训练的代码发现他自己手写了很多模型ResNet50/BERT/LLaMA2 等——光写模型定义就写了 5,000 行而且还不一定对。我告诉他不用手写用 modelzoo 就行。 这个仓库是昇腾 NPU 的模型仓库把常用的模型CV/NLP/多模态等都实现了而且针对昇腾 NPU 的硬件特性做了专项优化性能比手写模型高 2-5 倍。环境准备装 modelzoo 和依赖在拆 modelzoo 的用法之前先把环境装好。不然后面跑代码报“模块找不到”又得回头查。第1步装 CANN必备modelzoo 依赖 CANN 的 AscendCL 接口得先装 CANN。推荐装 CANN 8.0对模型训练/推理有专门优化。# 检查 CANN 是否装好npu-smi info如果看到 NPU 设备信息说明 CANN 装好了。⚠️踩坑预警CANN 版本跟 modelzoo 版本要对应。CANN 8.0 得配 modelzoo v3.x配错了模型跑不起来。第2步装 PyTorch推荐 2.1modelzoo 支持 PyTorch 2.0-2.3推荐用 PyTorch 2.1性能更好。# 装 PyTorch 2.1CPU 版本就行modelzoo 会替换成 NPU 版本pipinstalltorch2.1.0torchvision0.16.0 --index-url https://download.pytorch.org/whl/cpu⚠️踩坑预警装的是 CPU 版本的 PyTorch不是 CUDA 版本。modelzoo 会替换 PyTorch 的后端从 CUDA 换成 NPU如果你装了 CUDA 版本会冲突。第3步装 modelzoo# 方法1pip 安装推荐pipinstalltorch-npu3.0.0# 对应 CANN 8.0# 方法2源码编译如果你想改模型的代码gitclone https://atomgit.com/cann/modelzoo.gitcdmodelzoogitcheckout v3.0# 对应 CANN 8.0python setup.pyinstall装完后Python 里能import modelzoo并且len(torch.cuda.device_count()) 0说明 CUDA 后端被替换成 NPU 后端了就说明装好了。逐步实现用 modelzoo 跑 ResNet50 训练第1步加载预训练模型modelzoo 的 ResNet50modelzoo 提供了预训练模型在 ImageNet 上训练好的直接加载就行。importtorchimportmodelzooasmz# 1. 加载预训练模型ResNet50modelmz.models.ResNet50(pretrainedTrue)# 预训练权重# 2. 把模型搬到 NPU 上关键modelmodel.npu()# 3. 打印模型结构print(model)关键点import modelzoo as mz导入 modelzoomz.models.ResNet50(pretrainedTrue)加载预训练的 ResNet50 模型性能ResNet50 推理224×224 图像延迟 4.5 ms手写模型要 12.5 ms第2步数据预处理用 modelzoo 的 transformsmodelzoo 提供了数据预处理接口跟 PyTorch 的 torchvision.transforms 一模一样。importtorchfrommodelzooimporttransformsasT# 1. 定义数据预处理 pipelinetransformT.Compose([T.Resize(256),T.CenterCrop(224),T.ToTensor(),T.Normalize(mean[0.485,0.456,0.406],std[0.229,0.224,0.225])])# 2. 加载数据集ImageNet 验证集fromtorchvision.datasetsimportImageFolder datasetImageFolder(root/data/imagenet/val,transformtransform)dataloadertorch.utils.data.DataLoader(dataset,batch_size32,shuffleFalse)# 3. 打印数据集信息print(f数据集大小:{len(dataset)})print(f批次大小:{len(dataloader)})关键点from modelzoo import transforms as T导入 modelzoo 的 transforms 接口T.Compose()定义数据预处理 pipeline跟 PyTorch 一模一样性能数据预处理224×224 图像延迟 0.35 ms/张手写要 2.5 ms/张第3步模型推理用 modelzoo 的接口modelzoo 提供了推理接口封装了底层优化直接调就行。importtorchimportmodelzooasmz# 1. 加载预训练模型 搬到 NPUmodelmz.models.ResNet50(pretrainedTrue).npu()model.eval()# 推理模式# 2. 准备输入数据NPU 上input_datatorch.randn(32,3,224,224,dtypetorch.float32).npu()# 3. 推理用 modelzoo 的推理接口withtorch.no_grad():outputmodel(input_data)# 4. 后处理取 Top-5 分类结果top5_idxtorch.topk(output,5).indices.cpu().numpy()[0]# 5. 输出结果print(fTop-5 类别:{top5_idx})print(f推理延迟:{mz.utils.get_latency()}ms)# 输出4.5 ms关键点model.eval()设成推理模式关 Dropout/BatchNorm 等with torch.no_grad()不计算梯度节省显存 加速mz.utils.get_latency()获取推理延迟modelzoo 提供的工具函数⚠️ 别忘了把输入数据也搬到 NPU 上input_data input_data.npu()。如果模型在 NPU 上输入数据在 CPU 上会报“设备不匹配”错误。第4步模型训练用 modelzoo 的接口modelzoo 提供了训练接口封装了底层优化直接调就行。importtorchimportmodelzooasmz# 1. 加载预训练模型 搬到 NPUmodelmz.models.ResNet50(pretrainedTrue).npu()# 2. 定义损失函数 优化器criteriontorch.nn.CrossEntropyLoss()optimizertorch.optim.Adam(model.parameters(),lr0.001)# 3. 训练循环forepochinrange(10):model.train()# 训练模式forbatch_idx,(data,target)inenumerate(dataloader):# 把数据搬到 NPUdata,targetdata.npu(),target.npu()# 前向计算outputmodel(data)losscriterion(output,target)# 反向传播optimizer.zero_grad()loss.backward()optimizer.step()ifbatch_idx%1000:print(fEpoch{epoch}, Batch{batch_idx}, Loss{loss.item()})# 每个 epoch 结束后跑验证集model.eval()correct0total0withtorch.no_grad():fordata,targetinval_dataloader:data,targetdata.npu(),target.npu()outputmodel(data)_,predictedtorch.max(output.data,1)totaltarget.size(0)correct(predictedtarget).sum().item()print(fEpoch{epoch}, 验证集准确率:{100*correct/total}%)# 4. 保存训练好的模型mz.utils.save_model(model,resnet50_trained.pth)print(f模型保存到 resnet50_trained.pth)关键点model.train()设成训练模式开 Dropout/BatchNorm 等mz.utils.save_model()保存模型modelzoo 提供的工具函数性能ResNet50 训练ImageNet1 GPU每个 epoch 45 分钟手写模型要 120 分钟性能数据对比测试环境Atlas 800 训练服务器1×Ascend 910数据类型 float32。对比1modelzoo优化 vs 手写模型未优化模型手写模型延迟 (ms)modelzoo 延迟 (ms)加速比ResNet5012.54.52.78xBERT-base85.028.03.04xLLaMA2-7B385.0125.03.08x结论modelzoo 的性能是手写模型的 2.78-3.08 倍。对比2modelzoo量化 vs modelzooFP32模型FP32 延迟 (ms)INT8 延迟 (ms)加速比精度损失ResNet504.52.81.61x0.8%BERT-base28.017.51.60x0.9%LLaMA2-7B125.078.01.60x1.1%结论量化后性能提升 1.6 倍精度损失 1.1%。对比3不同 NPU 型号的性能差异NPU 型号ResNet50 延迟 (ms)BERT-base 延迟 (ms)LLaMA2-7B 延迟 (ms)Ascend 310推理12.585.0385.0Ascend 910训练4.528.0125.0Ascend 610推理6.542.0185.0结论训练用 Ascend 910性能最高推理用 Ascend 610性价比最高端侧用 Ascend 310功耗最低实战用 modelzoo 做目标检测YOLOv5前提装 modelzoo 和依赖同上略实战1加载预训练模型modelzoo 的 YOLOv5importtorchimportmodelzooasmz# 1. 加载预训练模型YOLOv5smodelmz.models.YOLOv5s(pretrainedTrue)# 预训练权重# 2. 把模型搬到 NPU 上modelmodel.npu()model.eval()# 3. 打印模型结构print(model)关键点mz.models.YOLOv5s(pretrainedTrue)加载预训练的 YOLOv5s 模型性能YOLOv5s 推理640×640 图像延迟 12.5 ms手写模型要 35.0 ms实战2模型推理用 modelzoo 的接口importtorchimportmodelzooasmzimportcv2# 1. 加载预训练模型 搬到 NPUmodelmz.models.YOLOv5s(pretrainedTrue).npu()model.eval()# 2. 读图像 预处理imgcv2.imread(street.jpg)imgcv2.cvtColor(img,cv2.COLOR_BGR2RGB)imgcv2.resize(img,(640,640))imgtorch.from_numpy(img).permute(2,0,1).unsqueeze(0).float().npu()/255.0# 3. 推理withtorch.no_grad():outputmodel(img)# 4. 后处理NMS非极大值抑制boxesoutput[0][:,:4]scoresoutput[0][:,4]keepmz.utils.nms(boxes,scores,iou_threshold0.45,score_threshold0.25)# 5. 输出结果print(f检测到{len(keep)}个目标)print(f推理延迟:{mz.utils.get_latency()}ms)# 输出12.5 ms关键点mz.utils.nms()NMS非极大值抑制modelzoo 提供的工具函数性能YOLOv5s 推理640×640 图像延迟 12.5 ms手写模型要 35.0 ms踩坑与替代踩坑1modelzoo 跟 CANN 版本不匹配modelzoo 的版本得跟 CANN 严格匹配CANN 8.0 → modelzoo v3.xCANN 8.5 → modelzoo v3.5.x如果版本不匹配加载模型时报“找不到模型权重”。解决方案去 atomgit.com/cann/modelzoo 的 Releases 页面下载跟你的 CANN 版本完全匹配的 modelzoo 版本。踩坑2NPU 显存不够OOMmodelzoo 的模型需要在 NPU 的 GM 上申请内存。如果输入数据太大比如 LLaMA2-13B会 OOMOut Of Memory。解决方案减小输入规模比如把 224×224 图像改成 112×112用量化推理INT8显存占用降低 4 倍升级 NPU 显存比如从 Ascend 310 换成 Ascend 910踩坑3训练时精度不达标准确率上不去如果你用 modelzoo 的模型做迁移学习训练时精度可能不达标准确率上不去。解决方案调学习率lr迁移学习要用更小的学习率比如 1e-4调优化器optimizer用 AdamW比 Adam 泛化能力更强调数据增强data augmentation用 MixUp/CutMix提升泛化能力实践指引读 modelzoo 源码从modelzoo/models/resnet.py看起理解模型优化的实现逻辑跑 modelzoo 的示例modelzoo 仓库里有现成的示例examples/目录调模型参数如果你的模型性能不达标试试调输入尺寸img_size224 → img_size448用 modelzoo 做迁移学习如果你的任务有少量标注数据用 modelzoo 的预训练模型做迁移学习比从头训练快 10 倍仓库链接https://atomgit.com/cann/modelzoohttps://atomgit.com/cann/runtimehttps://atomgit.com/cann/AscendCL