Torch模型转tensorflow模型踩坑记录

尝试在windows上将torch模型转换为tensorflow模型失败。于是转为在linux上转换,可以成功。记录一下操作过程。

本机环境win7_64位 conda4.3+tensorflow1.12+torch1.3.0+python3.6

在windows下进行转换失败,只能进行到将porch转换为onnx模型。将onnx转换为tf模型报错:

Torch模型转tensorflow模型踩坑记录

该错误表示没有tensorflow-addons模块,安装该模块提示需要tensorflow版本>=2.1.0

Torch模型转tensorflow模型踩坑记录

也就是说Tensorflow 1.x版本不支持该操作,自建一个虚拟环境,下载相关pytorch,tensorflow2.0等模块。不敢直接升级已有的tensorflow。因为tensorflow1.x和2.x API相差比较大,很多模块位置出现变化。为防止升级以后出现导包不成功(出现 no module name ‘xxx’)。还是自建一个虚拟环境比较好。

#创建虚拟环境,virtualenv_name为虚拟环境名称,该环境会保存在anaconda安装目录下的envs下,python版本为3.6

conda create -n virtualenv_name python=3.6

 

#安装tensorflow

pip install tensorflow==2.1.0

#如果在线安装不成功,就在pypi中下载tensorflow安装包继续离线安装

pip install tensorflow_file_path

 

#安装tensorflow-addons

pip install tensorflow-addons

 

#安装torch和torchvision ,到torch官网找到匹配的安装脚本

pip install torch==1.3.1+cpu torchvision==0.4.2+cpu -f https://download.pytorch.org/whl/torch_stable.html

 

#安装onnx

pip install onnx

 

#安装onnx-tensorflow,官网没有onnx-tensorflow的安装包,有git版本https://github.com/onnx/onnx-tensorflow。下载以后安装

pip install  -e onnx-tensorflow/

#pip list 查看安装的模块发现onnx==1.6,onn-tf==1.5。不过没关系,差别不大。也可以把onnx版本降为1.5

#运行转换脚本,脚本内容

Torch模型转tensorflow模型踩坑记录

在windows上报错

Torch模型转tensorflow模型踩坑记录

看不懂这是缺了个什么模块,查看很久的资料也没能解决。所以放弃在windows上转换。

转战linux

将安装模块的过程重新走一遍

然后运行转换脚本也报错

Torch模型转tensorflow模型踩坑记录

这个报错是说妄图占用不能使用的内存段,绝望了。

于是换了个思路,先把pth转换为onnx,再用onnx-tf 把onnx转为pb。

pb转onnx脚本内容

Torch模型转tensorflow模型踩坑记录

onnx转pb

onnx-tf convert -t tf -i modelname.onnx, -o result_name.pb

在windows上试过,只能把pth转换为onnx。onnx转为pb还是报错:

Torch模型转tensorflow模型踩坑记录

在linux上

Torch模型转tensorflow模型踩坑记录

#该错误是说不认识-t tf参数,于是乎去掉-t参数,再存运行

onnx-tf convert -i modelname.onnx, -o result_name.pb

Torch模型转tensorflow模型踩坑记录

可以看到该目录下生成了pb文件。大功告成,但这只是转换模型成功,还不知道能不能成功加载。

 

还试过pytorch2keras,不过该项目能支持转换的torch模块和方法有限,复杂的模型可能没办法转换。具体可以参考https://github.com/nerox8664/pytorch2keras中的readme,md文档