File size: 593 Bytes
c508d7f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 |
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')
|