如何将Yolov5模型转换为tensorflow.js

Question

是否可以将YOLOv5 PyTorch模型转换为Tensorflow.js模型？

我正在开发一个对象检测web应用程序。因此，我已经使用Yolov5训练了数据，但现在我正在寻找将该模型转换为Tf.js的正确方法。

Answer 1

我想这就是你要找的https://github.com/zldrobit/tfjs-yolov5-example

在YoloV5存储库中，运行export.py命令。

python export.py --weights yolov5s.pt --include tfjs

然后cd进入上述链接的repo，并将weights文件夹复制到公共目录中：

cp ./yolov5s_web_model public/web_model

不要忘记，您必须更改src/index.js中的names数组以匹配您的自定义模型。

但不幸的是，在大约1-2秒的时间里，它似乎慢得令人痛苦。我不认为我能够让WebGL工作。

Answer 2

只需很少的临时步骤，但大多数情况下它是有效的：

1. 将PyTorch导出为ONNX
2. 将ONNX转换为TF保存的模型
3. 将TF保存的模型转换为TFJS图形模型

从ONNX转换到TF时，如果遇到不支持的操作，则可能需要调整目标版本。

此外，确保将输入分辨率设置为固定值，任何动态输入都会在这种多步转换中混乱。