updates
This commit is contained in:
parent
ce11ef28f8
commit
ff7ee7f1f1
|
@ -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
|
||||||
|
|
Loading…
Reference in New Issue