This commit is contained in:
Glenn Jocher 2020-01-30 12:12:04 -08:00
parent ce11ef28f8
commit ff7ee7f1f1
1 changed files with 3 additions and 1 deletions

View File

@ -36,14 +36,16 @@ def detect(save_img=False):
# Fuse Conv2d + BatchNorm2d layers # Fuse Conv2d + BatchNorm2d layers
# model.fuse() # model.fuse()
# torch_utils.model_info(model, report='summary') # 'full' or 'summary'
# Eval mode # Eval mode
model.to(device).eval() model.to(device).eval()
# Export mode # Export mode
if ONNX_EXPORT: if ONNX_EXPORT:
model.fuse()
img = torch.zeros((1, 3) + img_size) # (1, 3, 320, 192) img = torch.zeros((1, 3) + img_size) # (1, 3, 320, 192)
torch.onnx.export(model, img, 'weights/export.onnx', verbose=False, opset_version=10) torch.onnx.export(model, img, 'weights/export.onnx', verbose=False, opset_version=11)
# Validate exported model # Validate exported model
import onnx import onnx