我们学习的机器学习课程和华为合作,最后需要将训练完成模型部署到昇腾系列开发板上,这里记录一下踩过的坑。
模型转换
昇腾系列开发板支持的模型格式为Caffe(非Caffe2)和Tensorflow(后缀为.pb),通过华为自己开发的MindStudio 开发环境进行转换得到开发板支持的.om格式网络。
这里我一开始使用的是Pytorch框架进行自己模型的编写即训练(不得不说Pytorch是真的好用,比较起来Caffe和老版的Tensorflow太不方便了),然后通过一些方式转换成开发板支持的网络格式,主要事实为了偷懒,不太想再去学Tensorflow和Caffe。这里两种转换方式我都尝试过,下面依次讲一下踩过的坑。
Tensorflow
如果是在华为昇腾开发板上使用的话,这里不太推荐用Tensorflow去转,好像华为官方对Tensorflow(.pb)转换的支持还不是太好,可能会遇到某些算子不支持的情况,但这里还是记录一下。
可以通过ONNX(Open Neural Network Exchange)交换格式进行转换,即Pytorch转ONNX再转Tensorflow:
- Pytorch to ONNX
from torch.autograd import Variable
import torch.onnx
import torchvision
dummy_input = Variable(torch.randn(10, 3, 224, 224))