Pytorch模型转Tensorflow

最新推荐文章于 2024-06-17 16:41:58 发布

一步一步HH

最新推荐文章于 2024-06-17 16:41:58 发布

阅读量1.6k

点赞数 2

分类专栏： tensorflow tfserving Python 文章标签： pytorch tensorflow 深度学习

本文链接：https://blog.csdn.net/m0_44982764/article/details/128386655

版权

本文介绍了如何将Pytorch模型转换成ONNX格式，然后进一步转化为Tensorflow的PB文件。在转换Yolov5模型时遇到因算子版本过高导致的错误，通过将`opset_version`设置为12解决了问题。

摘要由CSDN通过智能技术生成

1.将pytorch pt模型转onnx

def pttoonnx():
    model = torch.jit.load("./weights/crowdhuman_yolov5m.pt")
    img = torch.rand((1,3,640,640))
    torch.onnx.export(model, img, "./weights/yolov5.onnx")

2.将onnx模型转Tensorflow格式pb

def onnxtopb():
    onnx_model = onnx.load("./weights/yolov5.onnx")
    tf_rep = prepare(onnx_model, device="CPU")
    tf_rep.export_graph("./weights/yolov5")

yolov5模型onnx转pb的时候，出现一个错误，原因是转换算子版本太高，onnx_tf版本不支持，需要降低本版，torch.onnx.export(model, img, "./weights/yolov5.onnx", opset_version=12)，将opset_version设置为12

BackendIsNotSupposedToImplementIt: Unsqueeze version 13 is not implemented.