| import unittest | |
| import torch | |
| from tensorboardX import SummaryWriter | |
| class ONNXGraphTest(unittest.TestCase): | |
| def test_onnx_graph(self): | |
| import subprocess | |
| zoo_address = 'https://onnxzoo.blob.core.windows.net/models/opset_8/mnist/mnist.tar.gz' | |
| res = subprocess.call(['wget', '-nc', zoo_address]) | |
| assert res == 0, 'cannot download example onnx model from the zoo' | |
| res = subprocess.call(['tar', 'xf', 'mnist.tar.gz', '-C', 'examples/', 'mnist/model.onnx']) | |
| with SummaryWriter() as w: | |
| w.add_onnx_graph('examples/mnist/model.onnx') | |