Glenn Jocher 4 лет назад
Родитель
Сommit
a9d20eba3e
1 измененных файлов: 5 добавлений и 2 удалений
  1. +5
    -2
      models/export.py

+ 5
- 2
models/export.py Просмотреть файл

@@ -30,7 +30,7 @@ if __name__ == '__main__':

# TorchScript export
try:
print('\nStarting TorchScript export with torch %s...' % torch.__version)
print('\nStarting TorchScript export with torch %s...' % torch.__version__)
f = opt.weights.replace('.pt', '.torchscript') # filename
ts = torch.jit.trace(model, img)
ts.save(f)
@@ -52,6 +52,9 @@ if __name__ == '__main__':
onnx_model = onnx.load(f) # load onnx model
onnx.checker.check_model(onnx_model) # check onnx model
print(onnx.helper.printable_graph(onnx_model.graph)) # print a human readable model
print('ONNX export success, saved as %s\nView with https://github.com/lutzroeder/netron' % f)
print('ONNX export success, saved as %s' % f)
except Exception as e:
print('ONNX export failed: %s' % e)

# Finish
print('\nExports complete. Visualize with https://github.com/lutzroeder/netron.')

Загрузка…
Отмена
Сохранить