Update export.py (#7301)
* 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>pull/7311/head
parent
741fac815e
commit
d257c75c84
|
@ -480,15 +480,10 @@ def run(
|
|||
im, model = im.half(), model.half() # to FP16
|
||||
model.train() if train else model.eval() # training mode = no Detect() layer grid construction
|
||||
for k, m in model.named_modules():
|
||||
# if isinstance(m, Conv): # assign export-friendly activations
|
||||
# if isinstance(m.act, nn.SiLU):
|
||||
# m.act = SiLU()
|
||||
if isinstance(m, Detect):
|
||||
m.inplace = inplace
|
||||
m.onnx_dynamic = dynamic
|
||||
m.export = True
|
||||
if hasattr(m, 'forward_export'):
|
||||
m.forward = m.forward_export # assign custom forward (optional)
|
||||
|
||||
for _ in range(2):
|
||||
y = model(im) # dry runs
|
||||
|
|
Loading…
Reference in New Issue