PyTorch网络模型移植&训练指南

概述
约束与限制
迁移流程
快速上手
模型移植评估
环境准备
模型迁移
模型训练
性能调优和分析
精度调测
模型保存与转换
ShuffleNet模型调优示例
参考信息
FAQ

概述

当前阶段针对PyTorch框架实现的对接适配昇腾AI处理器的方案为在线对接方案。

方案特性及优点

昇腾AI处理器的加速实现方式是以各种算子为粒度进行调用（OP-based），即通过AscendCL调用一个或几个D亲和算子组合的形式，代替原有GPU的实现方式。其逻辑模型如图1所示。

图 1 逻辑模型

当前选择在线对接适配方案的主要原因有以下几点：

最大限度的继承PyTorch框架动态图的特性。
最大限度的继承GPU在PyTorch上的使用方式，可以使用户在将模型移植到昇腾AI处理器设备进行训练时，在开发方式和代码重用方面做到最小的改动。
最大限度的继承PyTorch原生的体系结构，保留框架本身出色的特性，比如自动微分、动态分发、Debug、Profiling、Storage共享机制以及设备侧的动态内存管理等。
扩展性好。在打通流程的通路之上，对于新增的网络类型或结构，只需涉及相关计算类算子的开发和实现。框架类算子，反向图建立和实现机制等结构可保持复用。
与GPU的使用方式和风格保持一致。用户在使用在线对接方案时，只需在Python侧和Device相关操作中，指定device为昇腾AI处理器，即可完成用昇腾AI处理器在PyTorch对网络的开发、训练以及调试，用户无需进一步关注昇腾AI处理器具体的底层细节。这样可以确保用户的最小化修改，迁移成本较低。

约束与限制

infershape阶段算子不支持unknowshape的推导。
cube计算的算子只支持float16。
不支持float16类型的inf/nan数据输入输出。
出现4D以上的format时不能降维。
Apex当前版本的实现方式为python实现，不支持APEX中的自定义优化CUDA Kernel。
Apex当前版本只支持适配昇腾AI处理器的混合精度计算和多种融合优化器功能，其他功能暂未支持。
集合通信约束：
- 数据并行模式中不同device上执行的图相同。
- 只支持1/2/4/8P粒度的分配。
- 只支持int8，int32，float16和float32数据类型。

迁移流程

模型迁移主要指将开源社区中实现过的模型迁移到昇腾AI处理器上，主要流程如图1所示。

图 1 迁移流程

表 1 迁移流程说明

场景	说明
模型选取	选取需要迁移的模型。
模型移植评估	详情请参见模型移植评估。
算子开发	详情请参见《PyTorch算子开发指南》。
环境准备	详情请参见环境准备。
模型迁移	详情请参见模型迁移。
模型训练	详情请参见模型训练。
错误分析	详情请参见《CANN 软件安装指南》中“日志参考>日志操作”章节设置日志级别和《CANN 软件安装指南》中“开发工具>AI Core Error分析工具”章节。
性能调优和分析	详情请参见性能调优和分析。
精度调测	详情请参见精度调测。
模型保存与转换	详情请参见模型保存与转换和《CANN 软件安装指南》中“ATC模型转换”章节。
应用软件开发	详情请参见《CANN 软件安装指南》中”应用开发（c++）“章节。
FAQ	主要涉及环境准备、模型迁移、模型调测和其他常见问题的解决方法。详情请参见FAQ。

快速上手

简介

对ResNet50模型进行迁移，帮助用户快速了解迁移过程。

模型选取

本样例基于PyTorch官网提供的Imagenet数据集训练模型main.py脚本进行适配昇腾910 AI处理器的迁移。

模型移植评估

模型是否可以迁移成功主要取决于模型算子是否支持昇腾AI处理器。故需要对模型算子对昇腾AI处理器的支持性进行评估，一般有两种方式评估算子支持性

模型迁移前，使用dump op方法获取算子信息，与《PyTorch API 支持清单》中自定义算子进行比较，确定是否支持。
模型迁移后，在昇腾设备上进行运行训练脚本，若存在不支持昇腾AI设备的算子，会提示报错信息。

若存在不支持算子，可以采用修该模型用等价支持的算子替换或者参考《PyTorch算子开发指南》进行算子开发。

ResNet50模型用到的算子已经在昇腾AI处理器上支持。

环境准备

请参见《PyTorch安装指南》进行CANN软件安装、PyTorch框架及混合精度模块安装，并配置环境变量。

参考PyTorch examples 准备模型运行所需要的Python环境及依赖。

硬件支持

aarch64架构推荐使用：Atlas800-9000+kunpeng920+Ascend910+NVMe 3.2T+750GRAM

X86_64架构推荐使用：Atlas800-9010+Intel Platinum8260+Ascend910+NVMe 3.2T+750GRAM

模型迁移

在main.py训练脚本的基础上进行修改，实现模型的单卡训练和单机多卡训练迁移。

单卡训练迁移

在main.py脚本中导入torch.npu模块。
```
import torch.npu
```
在main.py中定义训练设备。
```
CALCULATE_DEVICE = "npu:0"
```

修改参数以及判断选项，使其只在昇腾910 AI处理器上进行训练。

代码位置：main.py文件中的main_worker()函数：

def main_worker(gpu, ngpus_per_node, args):
    global best_acc1
    # 原代码为使用GPU进行训练，原代码如下：
    # args.gpu = gpu
    ############## npu modify begin #############
    args.gpu = None
    ############## npu modify end #############
    
    if args.gpu is not None:
        print("Use GPU: {} for training".format(args.gpu))
        
    if args.distributed:
        if args.dist_url == "env://" and args.rank == -1:
            args.rank = int(os.environ["RANK"])
        if args.multiprocessing_distributed:
            # For multiprocessing distributed training, rank needs to be the
            # global rank among all the processes
            args.rank = args.rank * ngpus_per_node + gpu
        dist.init_process_group(backend=args.dist_backend, init_method=args.dist_url,
                                world_size=args.world_size, rank=args.rank)
    # create model
    if args.pretrained:
        print("=> using pre-trained model '{}'".format(args.arch))
        model = models.__dict__[args.arch](pretrained=True)
    else:
        print("=> creating model '{}'".format(args.arch))
        model = models.__dict__[args.arch]()
    # 原代码中需要判断是否在GPU上进行训练，原代码如下：
    # if not torch.cuda.is_available():
        # print('using CPU, this will be slow')
    # elif args.distributed:
    ############## npu modify begin #############
    # 迁移后为直接判断是否进行分布式训练，去掉判断是否在GPU上进行训练
    if args.distributed:
    ############## npu modify end #############
        # For multiprocessing distributed, DistributedDataParallel constructor
        # should always set the single device scope, otherwise,
        # DistributedDataParallel will use all available devices.
        if args.gpu is not None:
           ......

将模型以及损失函数迁移到昇腾910 AI处理器上进行计算。

代码位置：main.py文件中的main_worker()函数：

elif args.gpu is not None:
        torch.cuda.set_device(args.gpu)
        model = model.cuda(args.gpu)
    else:
        # DataParallel will divide and allocate batch_size to all available GPUs
        if args.arch.startswith('alexnet') or args.arch.startswith('vgg'):
            model.features = torch.nn.DataParallel(model.features)
            model.cuda()
        else:
            # 原代码使用torch.nn.DataParallel()类来用多个GPU加速训练
            # model = torch.nn.DataParallel(model).cuda()
        ############## npu modify begin #############
            # 将模型迁移到NPU上进行训练。
           model = model.to(CALCULATE_DEVICE)
       ############## npu modify end #############
    # 原代码中损失函数是在GPU上进行计算
    # # define loss function (criterion) and optimizer
    # criterion = nn.CrossEntropyLoss().cuda(args.gpu)
    ############## npu modify begin #############
    # 将损失函数迁移到NPU上进行计算。
    criterion = nn.CrossEntropyLoss().to(CALCULATE_DEVICE)   
    ############## npu modify end #############

将数据集目标结果target修改成int32类型解决算子报错问题；将数据集迁移到昇腾910 AI处理器上进行计算。

代码位置：main.py文件中的train()函数：

   for i, (images, target) in enumerate(train_loader):
        # measure data loading time
        data_time.update(time.time() - end)

        if args.gpu is not None:
            images = images.cuda(args.gpu, non_blocking=True)
        # 原代码中训练数据集在GPU上进行加载计算，原代码如下：
        # if torch.cuda.is_available():
            # target = target.cuda(args.gpu, non_blocking=True)
        ############## npu modify begin #############
        # 将数据集迁移到NPU上进行计算并修改target数据类型，以提升性能
        if 'npu' in CALCULATE_DEVICE:     
            target = target.to(torch.int32)                      
        images, target = images.to(CALCULATE_DEVICE, non_blocking=True), target.to(CALCULATE_DEVICE, non_blocking=True)
        ############## npu modify end #############

代码位置：main.py文件中的validate()函数：

    with torch.no_grad():
        end = time.time()
        for i, (images, target) in enumerate(val_loader):
            if args.gpu is not None:
                images = images.cuda(args.gpu, non_blocking=True)
            # 原代码中训练数据集在GPU上进行加载计算，原代码如下：
            # if torch.cuda.is_available():
                # target = target.cuda(args.gpu, non_blocking=True)
            ############## npu modify begin #############
            # 将数据集迁移到NPU上进行计算并修改target数据类型
            if 'npu' in CALCULATE_DEVICE:
                target = target.to(torch.int32)
          images, target = images.to(CALCULATE_DEVICE, non_blocking=True), target.to(CALCULATE_DEVICE, non_blocking=True)
           ############## npu modify end #############

设置当前正在使用的device。

代码位置：main.py文件中的主函数入口：

if __name__ == '__main__':
    ############## npu modify begin #############
    if 'npu' in CALCULATE_DEVICE:
       torch.npu.set_device(CALCULATE_DEVICE)
    ############## npu modify begin #############
    main()

单机多卡训练修改

main.py增加头文件以支持基于PyTorch框架的模型在昇腾910 AI处理器上训练及进行混合精度训练。
```
import torch.npu
from apex import amp
```

参数设置增加以下参数，包括指定参与训练的昇腾910 AI处理器以及进行混合精度训练需要的参数。

parser.add_argument('--device', default='npu', type=str, help='npu or gpu')                        
parser.add_argument('--addr', default='10.136.181.115', type=str, help='master addr')                        
parser.add_argument('--device-list', default='0,1,2,3,4,5,6,7', type=str, help='device id list')
parser.add_argument('--amp', default=False, action='store_true', help='use amp to train the model')                    
parser.add_argument('--loss-scale', default=1024., type=float,
                    help='loss scale using in amp, default -1 means dynamic')
parser.add_argument('--opt-level', default='O2', type=str,
                    help='loss scale using in amp, default -1 means dynamic')

创建由device_id到process_id的映射函数，指定device进行训练。在main.py函数中增加以下接口。

def device_id_to_process_device_map(device_list):
    devices = device_list.split(",")
    devices = [int(x) for x in devices]
    devices.sort()

    process_device_map = dict()
    for process_id, device_id in enumerate(devices):
        process_device_map[process_id] = device_id

    return process_device_map

指定训练服务器的ip和端口。

代码位置：main.py文件中的主函数main()（修改部分为字体加粗部分）。

def main():
    args = parser.parse_args()
    ############## npu modify begin #############
    os.environ['MASTER_ADDR'] = args.addr 
    os.environ['MASTER_PORT'] = '29688'
    ############## npu modify end #############

创建由device_id到process_id的映射参数，获取单节点昇腾910 AI处理器数量。

代码位置：main.py文件中的主函数main()。

args.distributed = args.world_size > 1 or args.multiprocessing_distributed
############## npu modify begin #############
args.process_device_map = device_id_to_process_device_map(args.device_list)
if args.device == 'npu':
    ngpus_per_node = len(args.process_device_map)
else:
    ngpus_per_node = torch.cuda.device_count()
############## npu modify end #############
# 原代码如下：
# ngpus_per_node = torch.cuda.device_count()

获取进程process_id对应的昇腾910 AI处理器编号，指定在对应的昇腾910 AI处理器上进行训练。