知乎链接：寒武纪MLU PyTorch框架的Yolox移植教程 - 知乎 (zhihu.com)

寒武纪MLU PyTorch框架的Yolox移植教程

主要记录对Yolox官方源码在寒武纪平台下的移植过程。最终实现的效果是，能够在MLU PyTorch上正确跑通Yolox的量化模型并生成对应离线模型。

环境准备：

1. Linux系统（教程中用的是ubuntu18.04）

2. MLU驱动以及对应CNToolkit，cambricon_pytorch（教程中采用的是MLU270，相关文件的下载和安装参见寒武纪官网开发社区）

3. 高版本Pytorch version>=1.7（高版本Pytorch用来加载Yolox的pretrained model的权重）

跑通Yolox源码：

首先需要在高版本的PyTorch环境下把Yolox跑起来（非必要），本步骤的核心是load高版本下Yolox的权重，然后重新储存。因为MLU PyTorch的版本是1.3，而PyTorch较高版本权重的储存方式采用了压缩的方式，因此需要重新load权重，然后再重新将它以非压缩的方式储存。

1.下载官方开源的Yolox源码，按照网页上的Installation步骤安装。

2.下载训练好的权重，本教程以YOLOX-s为例，跑通demo.py。

3.修改demo.py，在模型加载权重后的代码后model.load_state_dict(ckpt["model"])，添加一行代码用于重新保存训练好的权重。

torch.save(model.state_dict(), "yolox_s_unzip.pt", _use_new_zipfile_serialization=False)

在MLU PyTorch环境下，在CPU上跑通Yolox：

1.进入对应Docker并激活MLU PyTorch环境。

2.下载官方开源的Yolox源码，按照网页上的Installation步骤安装。执行pip3 install -v -e . # or python3 setup.py develop 前需要修改一下setup.py文件，把pytorch版本的要求改一下。

# assert torch_ver >= [1, 7], "Requires PyTorch >= 1.7"assert torch_ver >= [1, 3], "Requires PyTorch >= 1.3"

3.新建一个文件夹weights，将之前新储存的权重文件copy进来。

4.修改Yolox模型中的激活函数，MLU PyTorch中没有自带torch.nn.SiLU激活函数。在/your/path/YOLOX-main/yolox/models/network_blocks.py中添加SiLU，并在get_activation中替换。

class SiLU(nn.Module):
    """export-friendly version of nn.SiLU()"""
    @staticmethod
    def forward(x):
        return x * torch.sigmoid(x)def get_activation(name="silu", inplace=True):
    if name == "silu":
        # module = nn.SiLU(inplace=inplace)
        module = SiLU()
    elif name == "relu":
        module = nn.ReLU(inplace=inplace)
    elif name == "lrelu":
        module = nn.LeakyReLU(0.1, inplace=inplace)
    else:
        raise AttributeError("Unsupported act type: {}".format(name))
    return module

5.复制demo.py为test_yolox_s.py，设置device为CPU，在YOLOX-main下打开终端，运行下面代码。

python tools/test_yolox_s.py image -f exps/default/yolox_s.py -c ./weights/yolox_s_unzip.pt --path assets/dog.jpg --conf 0.25 --nms 0.45 --tsize 640 --save_result --device cpu

6.正常出现如下信息，且输出文件夹下生成的检测图片正常即代表跑通。

| INFO     | __main__:main:335 - Args: Namespace(camid=0, cfg='cpu', ckpt='./weights/yolox_s_unzip.pt', conf=0.25, demo='image', device='cpu', exp_file='exps/default/yolox_s.py', experiment_name='yolox_s', fp16=False, fuse=False, legacy=False, name=None, nms=0.45, path='assets/dog.jpg', save_result=True, trt=False, tsize=640)
| INFO     | __main__:main:345 - Model Summary: Params: 8.97M, Gflops: 26.81
| INFO     | __main__:main:358 - loading checkpoint
| INFO     | __main__:main:363 - loaded checkpoint done.
| INFO     | __main__:inference:243 - Infer time: 0.2545s
| INFO     | __main__:image_demo:282 - Saving detection result in ./YOLOX_outputs/yolox_s/vis_res/

将Yolox量化：

Yolox的demo.py里构造了一个类Predictor来做推理，因此主要修改main方法和Predictor的__init__和inference部分。首先在main方法中添加mlu量化的代码，此时用来构造Predictor的就是量化后的模型了。

global quantized_modelif args.cfg == 'qua':
    # 为了保证精度，采用int16量化
    quantized_model = mlu_quantize.quantize_dynamic_mlu(model, dtype='int16', gen_quant=True)
    predictor = Predictor(
        quantized_model, exp, COCO_CLASSES, trt_file, decoder,
        args.device, args.fp16, args.legacy, args.cfg
    )

然后在__init__里新增一个属性self.cfg用于区分cpu运行，qua量化运行和mlu在线运行。

def __init__(
        self,
        model,
        exp,
        cls_names=COCO_CLASSES,
        trt_file=None,
        decoder=None,
        device="cpu",
        fp16=False,
        legacy=False,
        cfg="cpu"
    ):
        self.model = model
        self.cls_names = cls_names
        self.decoder = decoder
        self.num_classes = exp.num_classes
        self.confthre = exp.test_conf
        self.nmsthre = exp.nmsthre
        self.test_size = exp.test_size
        self.device = device
        self.fp16 = fp16
        self.preproc = ValTransform(legacy=legacy)
        self.cfg = cfg

接下来，在inference方法里，保存量化的权重。

with torch.no_grad():
    t0 = time.time()
    if self.cfg == "qua":
        outputs = self.model(img)
        torch.save(self.model.state_dict(), './yolox_s_int16_2022XXXX.pt')
        print('run qua!')
    else:
        outputs = self.model(img)

如果输出结果正常，说明量化正确。

将量化后的模型放到MLU上运行：

首先，同样需要修改main方法和Predictor的inference方法。其次，因为MLU PyTorch支持的算子有限，需要修改yolox结构中的部分代码。

1.在main方法中，给模型加载之前量化保存的权重，并将模型放到MLU上。

global quantized_modelglobal quantized_netif args.cfg == 'qua':
    quantized_model = mlu_quantize.quantize_dynamic_mlu(model, dtype='int16', gen_quant=True)
    predictor = Predictor(
        quantized_model, exp, COCO_CLASSES, trt_file, decoder,
        args.device, args.fp16, args.legacy, args.cfg
        )elif args.cfg == 'mlu':
    quantized_net = mlu_quantize.quantize_dynamic_mlu(model, dtype='int16', gen_quant=True)
    state_dict = torch.load('./yolox_s_int16_2022XXXX.pt')  # 用量化保存的权重文件
    quantized_net.load_state_dict(state_dict)
    quantized_net_mlu = quantized_net.to(ct.mlu_device())
    predictor = Predictor(
        quantized_net_mlu, exp, COCO_CLASSES, trt_file, decoder,
        args.device, args.fp16, args.legacy, args.cfg
    )else:
    predictor = Predictor(
        model, exp, COCO_CLASSES, trt_file, decoder,
        args.device, args.fp16, args.legacy,
    )

2.在Predictor的inference方法中，将输入图片放到MLU上。

if self.device == "gpu":
    img = img.cuda()
    if self.fp16:
        img = img.half()  # to FP16elif self.cfg == "mlu":
    img = img.type(torch.FloatTensor).to(ct.mlu_device())

3.添加mlu推理的选项。

with torch.no_grad():
    t0 = time.time()
    if self.cfg == "qua":
        outputs = self.model(img)
        torch.save(self.model.state_dict(), './yolox_s_int16_2022XXXX.pt')
        print('run qua!')
    elif self.cfg == "mlu":
        outputs = self.model(img)
        outputs = outputs.cpu()  # 模型输出的postprocess需要在CPU上做，一些算子mlu不支持
    else:
        outputs = self.model(img)

4.修改YOLOX-main/yolox/models/yolo_head.py文件中，YOLOXHead的decode_outputs方法。原因有三个，第一，在对模型输出进行解码时，新创建了两个变量grids和strides，因此需要将他们放到MLU上，才能保证mlu正常计算。第二，在这里发现对MLU上的变量进行切片赋值操作有问题，改用cat操作代替。第三，mlu不支持arange，meshgrid和fill算子，需要用常规算子替代。

def decode_outputs(self, outputs, dtype):
    grids = []
    strides = []
    
    for (hsize, wsize), stride in zip(self.hw, self.strides):
        # 修改arange和meshgrid算子
        def meshgrid_selfdefine(h_size: int, w_size: int):
            yv = torch.zeros(h_size, w_size)
            xv = torch.zeros(h_size, w_size)
            for i in range(h_size):
                yv[i, :] = i
            for j in range(w_size):
                xv[:, j] = j
            return yv, xv
        # yv, xv = torch.meshgrid([torch.arange(hsize), torch.arange(wsize)])
        yv, xv = meshgrid_selfdefine(hsize, wsize)
        yv = yv.type(torch.int32)
        xv = xv.type(torch.int32)
        grid = torch.stack((xv, yv), 2).view(1, -1, 2)
        # 添加类型转换
        grid = grid.type(dtype)
        stride = float(stride)  # 8,16,32
        grids.append(grid)
        shape = grid.shape[:2]
        # 修改full算子
        # strides.append(torch.full((*shape, 1), stride))
        tmp_stride = torch.zeros((*shape, 1)).type(torch.float)
        tmp_stride[...] = stride
        strides.append(tmp_stride)
   
    grids = torch.cat(grids, dim=1)
    strides = torch.cat(strides, dim=1)
    if str(outputs.device) != "cpu":
        grids = grids.to(ct.mlu_device())
        strides = strides.to(ct.mlu_device())
    # MLU切片赋值有问题，用cat代替
    # outputs[..., :2] = (outputs[..., :2] + grids) * strides
    # outputs[..., 2:4] = torch.exp(outputs[..., 2:4]) * strides
    outputs = torch.cat(((outputs[..., :2] + grids) * strides, outputs[..., 2:]), 2)
    outputs = torch.cat((outputs[..., :2], torch.exp(outputs[..., 2:4]) * strides, outputs[..., 4:]), 2)
    return outputs

这里还需要注意，按照上述修改，是可以用MLU PyTorch在MLU上正常运行的，但如果要保存为离线模型，就会出现问题，因为这个方法里出现了loop的控制流，在用mlu的jit.trace进行追踪生成离线模型时会出现警告或报错。由于我们的离线模型在部署的时候，输入的大小往往是固定的，因此，可以提前计算出grids和strides，用于decode的处理。

在__init__方法里添加

self.hw_fixed = [(80,80), (40,40), (20,20)]self.strides_fixed = [8, 16, 32]self.grids_fixed = []self.strides_fixed_next = []for (hsize, wsize), stride in zip(self.hw_fixed, self.strides_fixed):
    yv, xv = torch.meshgrid([torch.arange(hsize), torch.arange(wsize)])
    yv = yv.type(torch.int32)
    xv = xv.type(torch.int32)
    grid = torch.stack((xv, yv), 2).view(1, -1, 2)
    grid = grid.type(torch.float)
    stride = float(stride)  # 8,16,32
    self.grids_fixed.append(grid)
    shape = grid.shape[:2]
    self.strides_fixed_next.append(torch.full((*shape, 1), stride))self.grids_fixed = torch.cat(self.grids_fixed, dim=1)self.strides_fixed_next = torch.cat(self.strides_fixed_next, dim=1)

将decode_outputs方法更新为

def decode_outputs(self, outputs, dtype):
    if str(outputs.device) != "cpu":
        grids = self.grids_fixed
        strides = self.strides_fixed_next
        grids = grids.to(ct.mlu_device())  #MLU上运行报错数据类型不匹配，将其转化为float
        strides = strides.to(ct.mlu_device())
    else:
        grids = []
        strides = []
        for (hsize, wsize), stride in zip(self.hw, self.strides):
            yv, xv = torch.meshgrid([torch.arange(hsize), torch.arange(wsize)])
            yv = yv.type(torch.int32)
            xv = xv.type(torch.int32)
            grid = torch.stack((xv, yv), 2).view(1, -1, 2)
            grid = grid.type(dtype)
            stride = float(stride)  # 8,16,32
            grids.append(grid)
            shape = grid.shape[:2]
            strides.append(torch.full((*shape, 1), stride))
    grids = torch.cat(grids, dim=1)
    strides = torch.cat(strides, dim=1)
    outputs[..., :2] = (outputs[..., :2] + grids) * strides
    outputs[..., 2:4] = torch.exp(outputs[..., 2:4]) * strides
    return outputs

5.最后运行程序，设置cfg='mlu'，如果输出结果正常，则说明在MLU上运行成功。

保存离线模型

修改main方法，采用jit.trace追踪模型即可，要注意模型中是否存在控制流。

with torch.no_grad():
    t0 = time.time()
    if self.cfg == "qua":
        outputs = self.model(img)
        torch.save(self.model.state_dict(), './yolox_s_int16_2022XXXX.pt')
        print('run qua!')
    elif self.cfg == "mlu":
        core_number = 4
        ct.set_core_number(core_number)
        ct.set_input_format(0)
        ct.set_core_version('MLU270')
        ct.save_as_cambricon('yolox_s_MLU270_xxxx')
        net_offline = torch.jit.trace(self.model, img, check_trace=False)
        net_offline(img)
        torch_mlu.core.mlu_model.save_as_cambricon("")
        print('offline model save over!')
        outputs = self.model(img)
        outputs = outputs.cpu()  # 模型输出的postprocess需要在CPU上做，一些算子mlu不支持
    else:
        outputs = self.model(img)

至此，整个yolox在MLU PyTorch上的移植就做完了。