schihei/ascend-tools
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

pt2tf工具评审意见修改

Browse Source
This commit is contained in:
zxros10
2020-10-12 15:37:38 +08:00
parent f673126e32
commit c213852ac0
2 changed files with 60 additions and 71 deletions
Download Patch File Download Diff File Expand all files Collapse all files
pt2tf/README.md
+51 -44
View File
@@ -17,25 +17,34 @@
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/
## 预置条件
pytorch模型pth文件。pytorch的模型文件有两种：
1. 模型保存有网络结构和权重参数。需要在训练时使用如下接口保存模型：
1.pytorch模型文件。pytorch模型保存有两种，一种是保存有权重参数和网络结构，另外一种是指保存权重参数。本工具只支持保存权重参数的模型文件，模型保存接口示例：
```
torch.save(model_object,'resnet.pth')
torch.save(my_resnet.state_dict(),"my_resnet.pth")
```
2. 只保存模型权重参数。在训练时使用如下接口保存模型：
```
torch.save(my_resnet.state_dict(),"my_resnet.pth")
```
本工具两种模型的转换都支持，但是如果模型只有权重参数，则在转换时还需要完整的模型实现代码
2.模型实现代码。权重参数模型加载时，需要使用模型创建接口创建模型，作为模型加载的参数，所以需要模型实现代码。
## 工具获取
**方法1. 下载压缩包方式获取**
将 https://gitee.com/ascend/tools 仓中的脚本下载至服务器的任意目录。
@@ -55,31 +64,12 @@ pytorch模型pth文件。pytorch的模型文件有两种：
### 1. 安装工具依赖包
cd $HOME/AscendProjects/tools/pt2tf/
sudo pip3 install -r requirements.txt
sudo pip3 install -e onnx-tensorflow
### 2. pth模型文件转onnx
pt2tf工具对pytorch的两种模型转onnx都支持。如果是包含完备信息（网络结构和权重参数）的模型，仅仅需要模型文件即可；如果是仅包含权重参数的模型，则还需要模型的实现代码。
#### 2.1 包含网络结构和权重参数的模型转onnx
在pt2tf工具目录下执行pt2onnx.py脚本，例如:
```
python3 pt2onnx.py --model_path="./resnet50_model.pth" --input_shape=1 3 224 224
```
参数说明：
--model_path：pytorch模型路径
--input_shape: 模型输入 shape
执行脚本后，会在pytorch模型同一目录下生成onnx文件，文件名和pytorch模型名一致, 后缀为onnx
#### 2.2 权重参数模型文件转onnx
1. 将pytorch模型和实现源码拷贝到pt2tf目录下
2. 使用vim或者文本工具打开pt2onnx.py，修改load_weight_model函数。以resnet50模型为例，修改点如下：
2. 使用vim或者文本工具打开pt2onnx.py，修改load_model函数。以resnet50模型为例，修改点如下：
（1）导入模型实现文件：
@@ -100,24 +90,27 @@ pt2tf工具对pytorch的两种模型转onnx都支持。如果是包含完备信
```
#修改点3:训练好的模型路径
model.load_state_dict(torch.load(model_file))
model.load_state_dict(torch.load(model_path))
```
综上，完整的load_weight_model代码：
def load_weight_model(model_file):
from resnet50.resnet import ResNet50
model = ResNet50()
model.load_state_dict(torch.load(model_file))
return 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_type=1 --model_path="./resnet50_model.pth" --input_shape=1 3 224 224
python3 pt2onnx.py --model_path="./resnet50/models/resnet50_best.pth" --input_shape=1 3 224 224
```
​ 参数说明：
@@ -128,9 +121,23 @@ pt2tf工具对pytorch的两种模型转onnx都支持。如果是包含完备信
​ --input_shape: 模型输入 shape
如果执行成功，将在pytorch目录下生成onnx文件，文件名和pytorch模型文件名一致，例如./resnet50/models/resnet50_best.onnx
### 3.使用 onnx-tf工具将onnx转为 pb
执行命令
onnx-tf convert -i ./resnet50/model_resnet.onnx -o ./resnet50/model_resnet.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
```
pt2tf/pt2onnx.py
+9 -27
View File
@@ -21,12 +21,9 @@ def parse_args():
parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
parser.add_argument('--model_type', default=0, type=int,
help="""the type of pytorch model,
0:model wiht net and weight; 1:only weight""")
parser.add_argument('--model_path', default=None,
help="""the pytorch model pth file path""")
parser.add_argument('--input_shape', default=[], type=list,
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()
@@ -37,12 +34,15 @@ def parse_args():
return args
def load_weight_model(model_file):
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
model = None
#修改点2:放开创建模型对象注释,并根据自己的模型接口创建模型对象
#model = ResNet50()
@@ -52,30 +52,12 @@ def load_weight_model(model_file):
return model
def load_complete_mode(model_file):
return torch.load(model_file)
def load_model(model_type, model_path):
if not os.path.exists(model_path):
print("The pytorch model is not exist")
return None
model = None
if model_type == 0:
model = load_complete_mode(model_path)
elif model_type == 1:
model = load_weight_model(model_path)
else:
print("Unknow model type %d, please "
"execute --help to obtain help"%(model_type))
return model
def main():
args = parse_args()
print("model path ", args.model_path, ", shape ", args.input_shape)
#加载模型
model = load_model(args.model_type, args.model_path)
model = load_model(args.model_path, args.input_shape)
if model is None:
print("Load model failed")
return
@@ -85,7 +67,7 @@ def main():
#创建输入张量
input = torch.randn(tuple(args.input_shape))
#生成的onnx文件存放在pytorch模型同级目录下,文件名相同,后缀为onnx
export_onnx_file = os.path.splitext(model_file)[0] + '.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)