zhengtao
Gitee
Regular → Executable
+143
-128
@@ -1,128 +1,143 @@
|
||||
## 工具使用说明与扩展性介绍
|
||||
|
||||
### 1.Pytorch有两种模型保存方法
|
||||
|
||||
##### 1.1 保存整个神经网络的结构信息
|
||||
|
||||
- 该方法保存的模型通过torch.load('.pth')，直接初始化新的神经网络对象;
|
||||
|
||||
``*#保存模型*`
|
||||
|
||||
`torch.save(model_object,'resnet.pth')`
|
||||
|
||||
`*#加载模型*`
|
||||
|
||||
`model=torch.load('resnet.pth')`
|
||||
|
||||
##### 1.1 保存整个神经网络的结构信息
|
||||
|
||||
- 该方法保存的方式：首先是导入对应的网络，再通过net.load_state_dict(torch.load(’.pth’))完成模型参数的加载；
|
||||
|
||||
`*#将my_resnet模型存储为my_resnet.pth*`
|
||||
|
||||
`torch.save(my_resnet.state_dict(),"my_resnet.pth")`
|
||||
|
||||
`*#加载resnet，模型存放在my_resnet.pth* my_resnet.load_state_dict(torch.load("my_resnet.pth"))`
|
||||
|
||||
`*#其中my_resnet是my_resnet.pth对应的网络结构；*`
|
||||
|
||||
### 2.Pytorch载入只含模型参数pth文件
|
||||
|
||||
pth文件只保存网络中的参数，具有速度快，占空间少的优点，网上Pytorch实现的可供下载的预训练模型一般也是这种吗，加载并导出为onnx格式时还需要在继承 nn.Module 实现网络各Layer层，例如，下面的示例中使用Pytorch实现了一个Net。
|
||||
|
||||
```
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
class CivilNet(nn.Module):
|
||||
def __init__(self):
|
||||
super(CivilNet, self).__init__()
|
||||
self.conv1 = nn.Conv2d(3, 6, 5)
|
||||
self.pool = nn.MaxPool2d(2, 2)
|
||||
self.conv2 = nn.Conv2d(6, 16, 5)
|
||||
self.fc1 = nn.Linear(16 * 5 * 5, 120)
|
||||
self.fc2 = nn.Linear(120, 84)
|
||||
self.fc3 = nn.Linear(84, 10)
|
||||
self.gemfield = "gemfield.org"
|
||||
self.syszux = torch.zeros([1,1])
|
||||
|
||||
def forward(self, x):
|
||||
x = self.pool(F.relu(self.conv1(x)))
|
||||
x = self.pool(F.relu(self.conv2(x)))
|
||||
x = x.view(-1, 16 * 5 * 5)
|
||||
x = F.relu(self.fc1(x))
|
||||
x = F.relu(self.fc2(x))
|
||||
x = self.fc3(x)
|
||||
return x
|
||||
```
|
||||
|
||||
##### 2.1 CivilNet模型的保存
|
||||
|
||||
如果我们要保存一个训练好的PyTorch模型的话，会使用下面的API：
|
||||
|
||||
```
|
||||
cn = CivilNet()
|
||||
......
|
||||
torch.save(cn.state_dict(), "your_model_path.pth")
|
||||
```
|
||||
|
||||
##### 2.1 CivilNet模型的加载
|
||||
|
||||
而如果我们要load一个pth模型来进行前向的时候，会使用下面的API：
|
||||
|
||||
```
|
||||
cn = CivilNet()
|
||||
|
||||
#参数反序列化为python dict
|
||||
state_dict = torch.load("your_model_path.pth")
|
||||
#加载训练好的参数
|
||||
cn.load_state_dict(state_dict)
|
||||
|
||||
#变成测试模式，dropout和BN在训练和测试时不一样
|
||||
#eval()会把模型中的每个module的self.training设置为False
|
||||
cn = cn.cuda().eval()
|
||||
```
|
||||
|
||||
### 3.pt2tf工具的使用简介
|
||||
|
||||
#1 建立虚拟环境 $ virtualenv .venv
|
||||
|
||||
\#2 激活虚拟环境 $ source .venv/bin/activate
|
||||
|
||||
\#3 安装依赖包 pipinstall−rrequirements.txtpipinstall−rrequirements.txt pip install -e onnx-tensorflow
|
||||
|
||||
\#4 生成onnx模型 $ python pt2onnx.py
|
||||
|
||||
\#5 生成pb模型 $ onnx-tf convert -i efficientnet-b3.onnx -o efficientnet-b3.pb
|
||||
|
||||
pth转pb文件的工具源码如下，开发者可以根据自己需要转换的模型进行改造，并将Pytorch中未内置而需自己实现的模型脚本上传到工程目录的models文件夹下
|
||||
|
||||
```
|
||||
import torch
|
||||
from efficientnet_pytorch import EfficientNet
|
||||
|
||||
# Specify which model to use
|
||||
model_name = 'efficientnet-b3'
|
||||
image_size = EfficientNet.get_image_size(model_name)
|
||||
print('Image size: ', image_size)
|
||||
|
||||
# Load model
|
||||
model = EfficientNet.from_pretrained(model_name)
|
||||
model.set_swish(memory_efficient=False)
|
||||
model.eval()
|
||||
print('Model image size: ', model._global_params.image_size)
|
||||
|
||||
# Dummy input for ONNX
|
||||
dummy_input = torch.randn(1, 3, 300, 300)
|
||||
|
||||
# Export with ONNX
|
||||
torch.onnx.export(model, dummy_input, f"{model_name}.onnx", verbose=True)
|
||||
```
|
||||
|
||||
- 第二行导入Pytorch中内置的网络模型EfficientNet(Pytorch内置模型中)
|
||||
- 若内置模型未实现，我们在models文件夹中继承nn.Module类实现我们的网络模型，可以参考第二章中的CivilNet网络样例
|
||||
- 通过模型脚本对象的from_pretrained接口来导入pth参数文件，加载模型与参数
|
||||
- 调用Pytorch的onnx模块将网络模型导出为onnx模型
|
||||
- 使用onnx-tensorflow模块将onnx模型转换为pb模型
|
||||
|
||||
中文|[EN](README_EN.md)
|
||||
|
||||
# pytorch模型转onnx工具
|
||||
|
||||
## 功能
|
||||
当前ATC工具只支持pb和caffe模型转om模型。如果需要使用pytorch模型转om模型，可以将pytorch模型转为onnx格式，再转为pb。本工具提供pytorch模型转onnx，以及onnx转pb功能。
|
||||
|
||||
## 使用环境
|
||||
1. 安装Ubuntu18.04的服务器或者虚拟机；
|
||||
|
||||
2. 服务器或者虚拟机内存大于等于4G；
|
||||
|
||||
3. 已经安装pip3。如未安装，可以执行如下命令安装：
|
||||
|
||||
```
|
||||
sudo apt-get install python3-pip
|
||||
sudo pip3 install --upgrade pip -i http://mirrors.aliyun.com/pypi/simple/ --trusted-host mirrors.aliyun.com
|
||||
```
|
||||
|
||||
4. 已经安装tensorflow、keras和pytorch
|
||||
|
||||
当前昇腾平台支持tensorflow 1.15，考虑后继pb模型转om，tensorflow版本推荐1.15及之前版本。tensorflow 1.15版本需要源码编译安装；使用pip命令直接安装时可以1.15之前的版本，以1.14为例：
|
||||
|
||||
```
|
||||
sudo pip3 install tensorflow==1.14.0 -i http://mirrors.aliyun.com/pypi/simple/ --trusted-host mirrors.aliyun.com
|
||||
```
|
||||
|
||||
对应的keras版本为2.2.5，安装命令：
|
||||
|
||||
```
|
||||
sudo pip3 install keras==2.2.5 -i http://mirrors.aliyun.com/pypi/simple/ --trusted-host mirrors.aliyun.com
|
||||
```
|
||||
|
||||
pytorch版本只要适配待转换的pytorch模型即可。pytorch的安装可以参考官网：https://pytorch.org/get-started/locally/
|
||||
|
||||
## 预置条件
|
||||
|
||||
1.pytorch模型文件。pytorch模型保存有两种，一种是保存有权重参数和网络结构，另外一种是指保存权重参数。本工具只支持保存权重参数的模型文件，模型保存接口示例：
|
||||
|
||||
```
|
||||
torch.save(my_resnet.state_dict(),"my_resnet.pth")
|
||||
```
|
||||
|
||||
2.模型实现代码。权重参数模型加载时，需要使用模型创建接口创建模型，作为模型加载的参数，所以需要模型实现代码。
|
||||
|
||||
## 工具获取
|
||||
|
||||
**方法1. 下载压缩包方式获取**
|
||||
|
||||
将 https://gitee.com/ascend/tools 仓中的脚本下载至服务器的任意目录。
|
||||
|
||||
例如存放路径为：$HOME/AscendProjects/tools。
|
||||
|
||||
**方法2. 命令行使用git命令方式获取**
|
||||
|
||||
在命令行中：$HOME/AscendProjects目录下执行以下命令下载代码。
|
||||
|
||||
git clone https://gitee.com/ascend/tools.git
|
||||
|
||||
|
||||
|
||||
## 使用方法
|
||||
|
||||
### 1. 安装工具依赖包
|
||||
|
||||
cd $HOME/AscendProjects/tools/pt2tf/
|
||||
sudo pip3 install -e onnx-tensorflow
|
||||
|
||||
### 2. pth模型文件转onnx
|
||||
1. 将pytorch模型和实现源码拷贝到pt2tf目录下
|
||||
|
||||
2. 使用vim或者文本工具打开pt2onnx.py，修改load_model函数。以resnet50模型为例，修改点如下：
|
||||
|
||||
（1）导入模型实现文件：
|
||||
|
||||
```
|
||||
#修改点1:导入模型代码.
|
||||
#例如:模型实现代码目录为./resnet50,网络实现在resnet.py的class ResNet50类
|
||||
from resnet50.resnet import ResNet50
|
||||
```
|
||||
|
||||
（2） 使用pytorch实例化模型对象
|
||||
|
||||
```
|
||||
#修改点2:创建模型对象
|
||||
model = ResNet50()
|
||||
```
|
||||
|
||||
（3）加载训练好的模型
|
||||
|
||||
```
|
||||
#修改点3:训练好的模型路径
|
||||
model.load_state_dict(torch.load(model_path))
|
||||
```
|
||||
|
||||
综上，完整的load_weight_model代码：
|
||||
|
||||
def load_model(model_path, input_shape):
|
||||
if not os.path.exists(model_path):
|
||||
print("The pytorch model is not exist")
|
||||
return None
|
||||
from resnet50.resnet import ResNet50
|
||||
|
||||
model = ResNet50()
|
||||
|
||||
model.load_state_dict(torch.load(model_path))
|
||||
|
||||
return model
|
||||
|
||||
3. 执行转换脚本
|
||||
|
||||
```
|
||||
python3 pt2onnx.py --model_path="./resnet50/models/resnet50_best.pth" --input_shape=1 3 224 224
|
||||
```
|
||||
|
||||
参数说明：
|
||||
|
||||
--model_type: 模型类别，默认值为0，表示完备信息模型；1: 仅包含权重参数的模型
|
||||
|
||||
--model_path: pytorch模型存放路径
|
||||
|
||||
--input_shape: 模型输入 shape
|
||||
|
||||
如果执行成功，将在pytorch目录下生成onnx文件，文件名和pytorch模型文件名一致，例如./resnet50/models/resnet50_best.onnx
|
||||
|
||||
### 3.使用 onnx-tf工具将onnx转为 pb
|
||||
|
||||
执行命令
|
||||
|
||||
onnx-tf convert -i ./resnet50/models/resnet50_best.onnx -o ./resnet50/resnet50_best.pb
|
||||
|
||||
参数说明：
|
||||
|
||||
-i：onnx文件路径
|
||||
|
||||
-o: 输出的pb模型文件
|
||||
|
||||
onnx-tf convert的参数说明详见帮助：
|
||||
|
||||
```
|
||||
onnx-tf convert --help
|
||||
```
|
||||
|
||||
|
||||
+56
-14
@@ -13,22 +13,64 @@
|
||||
# limitations under the License.
|
||||
#########################################################################
|
||||
|
||||
import os
|
||||
import torch
|
||||
from efficientnet_pytorch import EfficientNet
|
||||
import argparse
|
||||
|
||||
# Specify which model to use
|
||||
model_name = 'efficientnet-b3'
|
||||
image_size = EfficientNet.get_image_size(model_name)
|
||||
print('Image size: ', image_size)
|
||||
def parse_args():
|
||||
|
||||
# Load model
|
||||
model = EfficientNet.from_pretrained(model_name)
|
||||
model.set_swish(memory_efficient=False)
|
||||
model.eval()
|
||||
print('Model image size: ', model._global_params.image_size)
|
||||
parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
|
||||
|
||||
# Dummy input for ONNX
|
||||
dummy_input = torch.randn(1, 3, 300, 300)
|
||||
parser.add_argument('--model_path', default=None,
|
||||
help="""the pytorch model pth file path""")
|
||||
parser.add_argument('--input_shape', nargs='+', type=int,
|
||||
help="""the model input shape, e.g. 1 3 224 224""")
|
||||
|
||||
args, unknown_args = parser.parse_known_args()
|
||||
if len(unknown_args) > 0:
|
||||
for bad_arg in unknown_args:
|
||||
print("ERROR: Unknown command line arg: %s" % bad_arg)
|
||||
raise ValueError("Invalid command line arg(s)")
|
||||
|
||||
# Export with ONNX
|
||||
torch.onnx.export(model, dummy_input, f"{model_name}.onnx", verbose=True)
|
||||
return args
|
||||
|
||||
def load_model(model_path, input_shape):
|
||||
if not os.path.exists(model_path):
|
||||
print("The pytorch model is not exist")
|
||||
return None
|
||||
|
||||
#修改点1:放开导入模型的注释,并导入自己的模型实现接口.
|
||||
#例如:模型实现代码目录为./resnet50,网络实现在resnet.py的class ResNet50类
|
||||
#from resnet50.resnet import ResNet50
|
||||
|
||||
|
||||
#修改点2:放开创建模型对象注释,并根据自己的模型接口创建模型对象
|
||||
#model = ResNet50()
|
||||
|
||||
#修改点3:放开加载模型的注释
|
||||
#model.load_state_dict(torch.load(model_file))
|
||||
|
||||
return model
|
||||
|
||||
|
||||
def main():
|
||||
args = parse_args()
|
||||
print("model path ", args.model_path, ", shape ", args.input_shape)
|
||||
#加载模型
|
||||
model = load_model(args.model_path, args.input_shape)
|
||||
if model is None:
|
||||
print("Load model failed")
|
||||
return
|
||||
|
||||
#将模型切换到推理状态
|
||||
model.eval()
|
||||
#创建输入张量
|
||||
input = torch.randn(tuple(args.input_shape))
|
||||
#生成的onnx文件存放在pytorch模型同级目录下,文件名相同,后缀为onnx
|
||||
export_onnx_file = os.path.splitext(args.model_path)[0] + '.onnx'
|
||||
|
||||
# Export with ONNX
|
||||
torch.onnx.export(model, input, export_onnx_file, verbose=True)
|
||||
|
||||
if __name__== "__main__":
|
||||
main()
|
||||
|
||||
@@ -12,15 +12,15 @@ Keras-Preprocessing==1.1.2
|
||||
Markdown==3.2.2
|
||||
numpy==1.19.2
|
||||
onnx==1.7.0
|
||||
-e git+https://github.com/onnx/onnx-tensorflow.git@6b9e76d6fca74d5ddbf47049b9670120abaaf70f#egg=onnx_tf
|
||||
-e onnx-tensorflow
|
||||
opt-einsum==3.3.0
|
||||
Pillow==7.2.0
|
||||
protobuf==3.13.0
|
||||
PyYAML==5.3.1
|
||||
six==1.15.0
|
||||
tensorboard==1.15.0
|
||||
tensorflow==1.15.0
|
||||
tensorflow-estimator==1.15.1
|
||||
tensorboard==1.14.0
|
||||
tensorflow==1.14.0
|
||||
tensorflow-estimator==1.14.0
|
||||
termcolor==1.1.0
|
||||
torch==1.6.0
|
||||
torchvision==0.7.0
|
||||
|
||||
Reference in New Issue
Block a user