diff --git a/pt2tf/README.md b/pt2tf/README.md index 0c7bc4f..dac8db7 100755 --- a/pt2tf/README.md +++ b/pt2tf/README.md @@ -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 +``` diff --git a/pt2tf/pt2onnx.py b/pt2tf/pt2onnx.py index 114a107..9c5ee31 100644 --- a/pt2tf/pt2onnx.py +++ b/pt2tf/pt2onnx.py @@ -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)