support yolov5

[wwdok]

support yolov5

[craymichael]

I just found a fix (for my use case at least). Before the hiddenlayer graph-building, add this to set the ONNX opset version globally:

from torch.onnx.symbolic_helper import _set_opset_version
_set_opset_version(11)

...

hl.build_graph(model, x)

[wwdok]

@craymichael 's solution is workable ! Now i can successfully get the yolov5 graph, the code is:

import torch
from torch.onnx.symbolic_helper import _set_opset_version
_set_opset_version(11)

model = torch.hub.load('ultralytics/yolov5', 'yolov5n', pretrained=True)
# print("==>> model: ", model)

import hiddenlayer as hl
graph = hl.build_graph(model, torch.zeros([1, 3, 512, 512]))

# To export graph, you need to install Graphviz excutable, and add it to environ path
import os 
os.environ["PATH"] += os.pathsep + "D:/Dependency/Graphviz/bin/"  # change this to your actual path
graph.save('yolov5.pdf') 

The generated graph is:
yolov5.pdf