需要将训练好的 pytorch model 移植到 ios上,需要转换成 coreml格式。
caffe2 可以在 coreml上使用,一种方式是 pytorch 转换成 caffe2, 再利用 caffe2的跨平台移植到移动端上;
此外,onnx 可以直接转换成 coreml 格式的。
需要转换的模型包含 conv
, batchnorm
, relu
, avgpool2d
, dropout
, linear
.
从源码编译 pytorch才能使用 onnx, 从源码编译才能使用 onnx.checker.编译完后先将pytorch model 转换成 onnx
格式的,
import torch
from model import TuneMobileNet, fcModel
from torch.autograd import Variable
model_name = 'mobile_dict.pt'
state_dict = torch.load(model_name)
cls_number = 17
model = TuneMobileNet(cls_number)
model.load_state_dict(state_dict)
x = Variable(torch.randn(1, 3, 224, 224), requires_grad=True)
torch_out = torch.onnx.export(model,
x,
'hard_mobile_convert.onnx',
verbose=True,
export_params=True)
在 avgpool2d的时候报错,padding size mismatch, 参考 issue [onnx] convered model the AveragePool bug in core-ml #3808, 和 pull requests fix pooling layer padding dim mistmatch bug #7。
解决了之后,再转换成 coreml格式的。
import onnx
import onnx_coreml
model = onnx.load('hard_mobile_convert.onnx')
cml = onnx_coreml.convert(model)
cml.save('hard_mobile.mlmodel')
接连报三个错误:
1. key error, 这是因为 pytorch 将 linear layer 转换成Transpose
和Gemm
两层,参考Why torch.nn.Linear is split into Transpose and Gemm layers in torch.onnx.export()? #3257, 和linear convert error #8, 就没有了transB
参数。
if node.attrs["broadcast"] != 1 or node.attrs["transB"] != 1:
raise ValueError(
"Gemm is supported only for inner_product layer"
)
key error: u'transB'
手动改为:
if node.attrs["broadcast"] != 1 or (hasattr(node, 'transB') and node.attrs["transB"] != 1):
bias
, 没有 权重w
。 w
, 结果在onnx_coreml.convert中提示图出错,这是 coremltools
自己的 check, debug 到了coremltools
导入了libcoremlpython.so
的动态库。只能将整个模型切分成两个部分,features
和 classifier
. features
转换成 coreml
识别的.mlmodel
格式,classifier
导出权重为 json文件,然后手动加载在 ios中。