support yolov5

[wwdok]

hello, I want to draw the graph of yolov5 by folowing code:

import torch

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

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

but it reports error:

c:\ProgramData\Anaconda3\lib\site-packages\hiddenlayer\pytorch_builder.py in import_graph(hl_graph, model, args, input_names, verbose)
     69     # Run the Pytorch graph to get a trace and generate a graph from it
     70     trace, out = torch.jit._get_trace_graph(model, args)
---> 71     torch_graph = torch.onnx._optimize_trace(trace, torch.onnx.OperatorExportTypes.ONNX)
...
RuntimeError: Unsupported: ONNX export of index_put in opset 9. Please try opset version 11.

some issues (https://github.com/onnx/onnx/issues/3057, https://github.com/pytorch/pytorch/issues/46237) say place opset_version=11 in torch.onnx.export(), but here i can not find torch.onnx.export(), so i don't know how to fix this error.

[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