pytorch训练打印loss pytorch load

转载

半夜未央好 2024-08-14 11:48:48

文章标签 pytorch训练打印loss 嵌入式 github 2d 文章分类 PyTorch 人工智能

工程的完整链接可以参考Github链接。

Pytorch以其动态图的调用方式，深得许多科研人员的喜爱，是许多人进行科研研究、算法预研的不二之选。本文我们跟大家讨论一下，如何使用Pytorch来进行嵌入式的算法部署。这里我们采用的离线训练框架为Pytorch，嵌入式端的推理框架为阿里巴巴近期开源的高性能推理框架MNN。下面我们将结合MNIST这个简单的分类任务来跟大家一步一步的完成嵌入式端的部署。

Pytorch的模型不能直接被MNN进行解析，所以我们这里需要选定一个媒介。参考之前专栏的一篇文章《整合mxnet和MNN的嵌入式部署流程》，这里也采用ONNX进行pytorch和MNN之间的桥梁。

模型的设计

模型的设计与《整合mxnet和MNN的嵌入式部署流程》文中的模型设计基本一样，大家可以看下面代码：

import

大家可以注意一下上述代码里的注释部分，这里我们进行一下介绍。由于Pytorch在实现矩阵乘法的时候，需要使用view进行数据的拉平，然后再进行matmul的操作，这几个操作MNN并没有进行支持（我第一次的实现是使用注释部分的代码，然后MNNConvert的时候报的错误，不支持一些op）。所以我们这里采用了一个4*4和1*1的con2d来代替全连接层。

2. 导出ONNX模型

我们需要导出另外一种可以被MNN解析的模型格式，这里我们选择的是ONNX。如下为导出ONNX的脚本文件:

import

3. 导出MNN模型

MNN提供了转换ONNX到MNN模型的工具，执行如下脚本即可，关于MNN转换工具编译可以参考Model Conversion。下面是转换脚本：

./

输出的结果如下：

Start to Convert Other Model Format To MNN Model...
[16:09:54] /Users/xindongzhang/MNN/tools/converter/source/onnx/onnxConverter.cpp:29: ONNX Model ir version: 3
Start to Optimize the MNN Net...
[16:09:54] /Users/xindongzhang/MNN/tools/converter/source/optimizer/optimizer.cpp:44: Inputs: 0
[16:09:54] /Users/xindongzhang/MNN/tools/converter/source/optimizer/optimizer.cpp:54: Outputs: 32, Type = Squeeze
Converted Done!

可以看出采用的ONNX IR版本为3，输入的节点名字为0，输出节点名字为32.

4. 在线部署

在线部署流程在这里，为使用MNN加载解析好的mnn模型参数进行inference等一系列业务操作。关于如何在android上面使用mnn进行部署，本专栏已经有好几篇介绍的文章，这里就不进行赘述了。完整的JNI业务代码可以参考如下链接JNI 业务代码。