本文档介绍如何把pytorch训练的LSTM网络转换成mindspore Lite支持的ms模型,并实现手机端的推理。
从mindspore Lite官网中的算子支持列表中发现,如果模型中用到了LSTM算子,只能从onnx转ms,目前主流的网络都是用pytorch框架训练,因此我们先用pytorch搭建LSTM模型,再调用torch.onnx把模型转换成onnx。
我们用LSTM搭建一个简单的分类网络模型,第一层是LSTM层,后续是两个全连接层,其网络结构用pytorch定义如下:
class LstmClassifier(nn.Module):
def __init__(self, input_dim = 6,hidden_dim=500, lstm_layers =2, number_of_classes=7356, data_fixed_length=100):
super(OcrLstm, self).__init__()
self.input_dim = input_dim
self.hidden_dim = hidden_dim
self.LSTM_layers = lstm_layers
self.number_of_classes = number_of_classes
# 直接定义lstm的层
self.dropout = nn.Dropout(0.1)
self.lstm = nn.LSTM(input_size=self.input_dim, hidden_size=self.hidden_dim, num_layers=self.LSTM_layers, batch_first=True)
self.fc1 = nn.Linear(in_features=self.hidden_dim, out_features=200)
self.fc2 = nn.Linear(200, number_of_classes)
def forward(self,x):
#out= self.LSTM(x)
h = torch.zeros((self.LSTM_layers, x.size(0), self.hidden_dim))
c = torch.zeros((self.LSTM_layers, x.size(0), self.hidden_dim))
out, _ = self.lstm(x, (h, c))
out = self.dropout(out)
out = torch.relu_(self.fc1(out[:, -1, :]))
out = self.dropout(out)
out = torch.softmax(self.fc2(out), dim=1)
return out
定义好模型之后,直接调用torch.onnx方法,把模型导出
x = torch.randn(1,100,6)
t = model(x)
torch.onnx.export(model,x,"LSTM.onnx")
使用Netron打开转换出来的模型,其结构如下:
使用mindspore的converter工具可以直接把onnx模型转换成ms模型,转换的命令如下
./converter_lite --fmk=ONNX --modelFile=OcrLSTM.onnx --outputFile=TEXT_REC_MODEL.onnx
转换过程的日志打印success:0表示转换成功
使用mindspore Lite可以把转换出来的模型部署到不同的终端,这里以部署到arm32的设备为例,使用C++代码,其余设备的与之相同,模型部署arm32的设备时,需要先下载对应的mindpore Lite依赖库,依赖库在这里下载下载MindSpore Lite — MindSpore Lite master documentation
由于我的设备是arm32,所以我直接下载列表中的第一个
下载好的依赖包中,runtime包含mindspore Lite运行所需的依赖
为了能让我们的C++程序成功调用mindspore Lite,我们需要修改我们C++工程中的CMakeLists.txt
在头文件中加入上述runtime目录中的include,同时加入动态链接库
target_include_directories(工程名 PRIVATE ${CMAKE_CURRENT_SOURCE_DIR} ${CMAKE_CURRENT_SOURCE_DIR} minspore目录/include )
target_link_libraries(tcdocformat PRIVATE D:/Project/teamproject/IdeaHubRecognition_latest/MeetingComponent_codehub/cmake/out/dist/armeabi/libmindspore-lite.so)
加入依赖后就可以在Cmake工程中直接用mindspore定义的接口调用mindspore Lite引擎了
这里参考官方文档中使用C++接口推理,只需要修改其中很小的一部分就可以实现端测推理LSTM,
模型初始化,加载等操作与上述文档保持一致,这里只讲如何定义输入,我们先获取模型的输入层,然后修改模型的尺寸,最后构造输入向量,把向量的地址传入输入层,具体代码如下
auto inputs = model->GetInputs();
// 把数据转换成连续内存
int rowNum = inputData.size();
int calNum = inputData.front().size();
int channel = 1;
// 输入resize成固定的大小
auto inputs = model->GetInputs();
vector resize_shape = {1, rowNum, calNum};
std::vector> new_shapes;
new_shapes.push_back(resize_shape);
model->Resize(inputs, new_shapes);
// 定义一个一维矩阵,把输入的值拷贝到该矩阵
float inputDataArr[rowNum * calNum];
int index;
for (int i = 0; i < rowNum; i++) {
for (int j = 0; j < calNum; j++) {
index = i * calNum + j;
inputDataArr[index] = inputData[i][j];
}
}
// 把一维矩阵的地址传入到模型的输入层inTensor
memcpy(inTensor.MutableData(), inputDataArr, channel * calNum * rowNum * sizeof(float));
最后调用predict方法即可完成推理
auto predict_ret = model->Predict(inputs, &outputs);