Update export.py (#7301)
Browse files* Update export.py
Simplify code.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
export.py
CHANGED
@@ -480,15 +480,10 @@ def run(
|
|
480 |
im, model = im.half(), model.half() # to FP16
|
481 |
model.train() if train else model.eval() # training mode = no Detect() layer grid construction
|
482 |
for k, m in model.named_modules():
|
483 |
-
# if isinstance(m, Conv): # assign export-friendly activations
|
484 |
-
# if isinstance(m.act, nn.SiLU):
|
485 |
-
# m.act = SiLU()
|
486 |
if isinstance(m, Detect):
|
487 |
m.inplace = inplace
|
488 |
m.onnx_dynamic = dynamic
|
489 |
m.export = True
|
490 |
-
if hasattr(m, 'forward_export'):
|
491 |
-
m.forward = m.forward_export # assign custom forward (optional)
|
492 |
|
493 |
for _ in range(2):
|
494 |
y = model(im) # dry runs
|
|
|
480 |
im, model = im.half(), model.half() # to FP16
|
481 |
model.train() if train else model.eval() # training mode = no Detect() layer grid construction
|
482 |
for k, m in model.named_modules():
|
|
|
|
|
|
|
483 |
if isinstance(m, Detect):
|
484 |
m.inplace = inplace
|
485 |
m.onnx_dynamic = dynamic
|
486 |
m.export = True
|
|
|
|
|
487 |
|
488 |
for _ in range(2):
|
489 |
y = model(im) # dry runs
|