onnx2caffe的一些经验

1 背景

因为上海思的npu，CNN部分需要转换为caffe

2 环境

2.1 onnx转换环境

python：3.6
pytorch：1.4
onnx：1.7

2.2 caffe转换环境

python：2.7
onnx：1.7
caffe：1.0

3 步骤

3.1 pytorch2onnx

转换后用onnx的推理一下，验证转换的onnx是否正确

3.2 onnx瘦身

onnxsimplifier

pip install onnx-simplifier
python -m onnxsim srcXXX.onnx dstXXX.onnx

3.3 onnx2caffe

4 经验

4.1 opset_version=11

torch1.4版本，在torch.onnx.export可以选择opset_version=11,可以解决后面很多问题

// torch.onnx.export
torch.onnx.export(..., opset_version=11)

4.2 model.module

如果torch在训练的时候使用了nn.DataParallel，模型加载的时候也要用nn.DataParallel加载权重，在export的时候使用model.module。否则onnx导出的结果不正确（也有其他的办法处理权重）

model = nn.DataParallel(model)
torch.onnx.export(model.module, ...)

4.3 op：mul在转换时的问题

MT的轮子在onnx2caffe将mul转换为eltwise使用elstwise_param是PROD。在需要给某一层乘以一个固定的数时，这样的转换会报错，找不到其中一个bottom层，mul在转换的时候会将固定的数转换为一个layer（下面修改前中的 bottom: “135”）
注意：后面的解决办法只是在遇到这个问题的时候使用，不具有普遍适用性

4.3.1 解决方法：将mul层手动修改为Scale层

修改前

layer {
  name: "136"
  type: "Eltwise"
  bottom: "134"
  bottom: "135"
  top: "136"
  eltwise_param {
    operation: PROD
  }
}

修改后

layer {
  name: "136"
  type: "Scale"
  bottom: "134"
  top: "136"
  param {
    lr_mult: 0
  }
  scale_param {
    filler {
      type: "constant"
      value: 0.1
    }
  }
}

4.3.2 解决方法： 修改MT的轮子(TODO)

1.将mul的转换增加一个分支专门处理Scale的情况

_ONNX_NODE_REGISTRY = {
    ...
    "Mul": _convert_Mul
    ...
    }    

2.应该是可以参照BatchNorm的scale层处理

4.4 op：Sub在转换时的问题

4.4.1 解决方法1：增加和Add相同的Sub层，手动修改参数

_weightloader.py和_operators.py

_ONNX_NODE_REGISTRY = {
    ...
    "Add": _convert_Add,
    "Sub": _convert_Add,
    ...
    }

*.prototxt(转换后的caffetxt)
修改前

layer {
  name: "128"
  type: "Eltwise"
  bottom: "124"
  bottom: "127"
  top: "128"
  eltwise_param {
    operation: SUM
  }
}

修改后

layer {
  name: "128"
  type: "Eltwise"
  bottom: "124"
  bottom: "127"
  top: "128"
  eltwise_param {
    operation: SUM
    coeff: 1
    coeff: -1
  }
}

4.4.2 解决方法2：修改MT的轮子(TODO)

增加Sub层的处理函数

_ONNX_NODE_REGISTRY = {
    ...
    "Sub": _convert_Sub
    ...
    }

4.5 op：Squeeze在转换时的问题

Squeeze没有参数，可以在训练代码直接去掉该op（仅限op是CNN的最后一层，中间层当然不能去掉）

4.6 onnx瘦身

onnxsimplifier非常重要

pip install onnx-simplifier
python -m onnxsim org.onnx org_sim.onnx 

5 TODO

5.1 pytorch2caffe
直接转换，github上有轮子
5.2 caffe修改的时候还有更简单的方法

TODO的问题没时间整（万恶的美帝限制了菊花，把海思搞的风雨飘摇，要备货其他的npu了，有空再搞这个，或者 pr给MT）