Export with official `nn.SiLU()` (#7256)
Browse files- export.py +4 -7
- utils/general.py +1 -1
export.py
CHANGED
@@ -54,7 +54,6 @@ from pathlib import Path
|
|
54 |
|
55 |
import pandas as pd
|
56 |
import torch
|
57 |
-
import torch.nn as nn
|
58 |
from torch.utils.mobile_optimizer import optimize_for_mobile
|
59 |
|
60 |
FILE = Path(__file__).resolve()
|
@@ -64,10 +63,8 @@ if str(ROOT) not in sys.path:
|
|
64 |
if platform.system() != 'Windows':
|
65 |
ROOT = Path(os.path.relpath(ROOT, Path.cwd())) # relative
|
66 |
|
67 |
-
from models.common import Conv
|
68 |
from models.experimental import attempt_load
|
69 |
from models.yolo import Detect
|
70 |
-
from utils.activations import SiLU
|
71 |
from utils.datasets import LoadImages
|
72 |
from utils.general import (LOGGER, check_dataset, check_img_size, check_requirements, check_version, colorstr,
|
73 |
file_size, print_args, url2file)
|
@@ -474,10 +471,10 @@ def run(
|
|
474 |
im, model = im.half(), model.half() # to FP16
|
475 |
model.train() if train else model.eval() # training mode = no Detect() layer grid construction
|
476 |
for k, m in model.named_modules():
|
477 |
-
if isinstance(m, Conv): # assign export-friendly activations
|
478 |
-
|
479 |
-
|
480 |
-
|
481 |
m.inplace = inplace
|
482 |
m.onnx_dynamic = dynamic
|
483 |
if hasattr(m, 'forward_export'):
|
|
|
54 |
|
55 |
import pandas as pd
|
56 |
import torch
|
|
|
57 |
from torch.utils.mobile_optimizer import optimize_for_mobile
|
58 |
|
59 |
FILE = Path(__file__).resolve()
|
|
|
63 |
if platform.system() != 'Windows':
|
64 |
ROOT = Path(os.path.relpath(ROOT, Path.cwd())) # relative
|
65 |
|
|
|
66 |
from models.experimental import attempt_load
|
67 |
from models.yolo import Detect
|
|
|
68 |
from utils.datasets import LoadImages
|
69 |
from utils.general import (LOGGER, check_dataset, check_img_size, check_requirements, check_version, colorstr,
|
70 |
file_size, print_args, url2file)
|
|
|
471 |
im, model = im.half(), model.half() # to FP16
|
472 |
model.train() if train else model.eval() # training mode = no Detect() layer grid construction
|
473 |
for k, m in model.named_modules():
|
474 |
+
# if isinstance(m, Conv): # assign export-friendly activations
|
475 |
+
# if isinstance(m.act, nn.SiLU):
|
476 |
+
# m.act = SiLU()
|
477 |
+
if isinstance(m, Detect):
|
478 |
m.inplace = inplace
|
479 |
m.onnx_dynamic = dynamic
|
480 |
if hasattr(m, 'forward_export'):
|
utils/general.py
CHANGED
@@ -738,7 +738,7 @@ def non_max_suppression(prediction,
|
|
738 |
# min_wh = 2 # (pixels) minimum box width and height
|
739 |
max_wh = 7680 # (pixels) maximum box width and height
|
740 |
max_nms = 30000 # maximum number of boxes into torchvision.ops.nms()
|
741 |
-
time_limit = 0.
|
742 |
redundant = True # require redundant detections
|
743 |
multi_label &= nc > 1 # multiple labels per box (adds 0.5ms/img)
|
744 |
merge = False # use merge-NMS
|
|
|
738 |
# min_wh = 2 # (pixels) minimum box width and height
|
739 |
max_wh = 7680 # (pixels) maximum box width and height
|
740 |
max_nms = 30000 # maximum number of boxes into torchvision.ops.nms()
|
741 |
+
time_limit = 0.1 + 0.03 * bs # seconds to quit after
|
742 |
redundant = True # require redundant detections
|
743 |
multi_label &= nc > 1 # multiple labels per box (adds 0.5ms/img)
|
744 |
merge = False # use merge-NMS
|