Update export.py, yolo.py `sys.path.append()` (#3579)
Browse files- models/export.py +7 -5
- models/yolo.py +4 -2
models/export.py
CHANGED
@@ -9,13 +9,15 @@ import sys
|
|
9 |
import time
|
10 |
from pathlib import Path
|
11 |
|
12 |
-
sys.path.append(Path(__file__).parent.parent.absolute().__str__()) # to run '$ python *.py' files in subdirectories
|
13 |
-
|
14 |
import torch
|
15 |
import torch.nn as nn
|
16 |
from torch.utils.mobile_optimizer import optimize_for_mobile
|
17 |
|
18 |
-
|
|
|
|
|
|
|
|
|
19 |
from models.experimental import attempt_load
|
20 |
from utils.activations import Hardswish, SiLU
|
21 |
from utils.general import colorstr, check_img_size, check_requirements, file_size, set_logging
|
@@ -56,12 +58,12 @@ def export(weights='./yolov5s.pt', # weights path
|
|
56 |
model.train() if train else model.eval() # training mode = no Detect() layer grid construction
|
57 |
for k, m in model.named_modules():
|
58 |
m._non_persistent_buffers_set = set() # pytorch 1.6.0 compatibility
|
59 |
-
if isinstance(m,
|
60 |
if isinstance(m.act, nn.Hardswish):
|
61 |
m.act = Hardswish()
|
62 |
elif isinstance(m.act, nn.SiLU):
|
63 |
m.act = SiLU()
|
64 |
-
elif isinstance(m,
|
65 |
m.inplace = inplace
|
66 |
m.onnx_dynamic = dynamic
|
67 |
# m.forward = m.forward_export # assign forward (optional)
|
|
|
9 |
import time
|
10 |
from pathlib import Path
|
11 |
|
|
|
|
|
12 |
import torch
|
13 |
import torch.nn as nn
|
14 |
from torch.utils.mobile_optimizer import optimize_for_mobile
|
15 |
|
16 |
+
FILE = Path(__file__).absolute()
|
17 |
+
sys.path.append(FILE.parents[1].as_posix()) # add yolov5/ to path
|
18 |
+
|
19 |
+
from models.common import Conv
|
20 |
+
from models.yolo import Detect
|
21 |
from models.experimental import attempt_load
|
22 |
from utils.activations import Hardswish, SiLU
|
23 |
from utils.general import colorstr, check_img_size, check_requirements, file_size, set_logging
|
|
|
58 |
model.train() if train else model.eval() # training mode = no Detect() layer grid construction
|
59 |
for k, m in model.named_modules():
|
60 |
m._non_persistent_buffers_set = set() # pytorch 1.6.0 compatibility
|
61 |
+
if isinstance(m, Conv): # assign export-friendly activations
|
62 |
if isinstance(m.act, nn.Hardswish):
|
63 |
m.act = Hardswish()
|
64 |
elif isinstance(m.act, nn.SiLU):
|
65 |
m.act = SiLU()
|
66 |
+
elif isinstance(m, Detect):
|
67 |
m.inplace = inplace
|
68 |
m.onnx_dynamic = dynamic
|
69 |
# m.forward = m.forward_export # assign forward (optional)
|
models/yolo.py
CHANGED
@@ -10,8 +10,8 @@ import sys
|
|
10 |
from copy import deepcopy
|
11 |
from pathlib import Path
|
12 |
|
13 |
-
|
14 |
-
|
15 |
|
16 |
from models.common import *
|
17 |
from models.experimental import *
|
@@ -25,6 +25,8 @@ try:
|
|
25 |
except ImportError:
|
26 |
thop = None
|
27 |
|
|
|
|
|
28 |
|
29 |
class Detect(nn.Module):
|
30 |
stride = None # strides computed during build
|
|
|
10 |
from copy import deepcopy
|
11 |
from pathlib import Path
|
12 |
|
13 |
+
FILE = Path(__file__).absolute()
|
14 |
+
sys.path.append(FILE.parents[1].as_posix()) # add yolov5/ to path
|
15 |
|
16 |
from models.common import *
|
17 |
from models.experimental import *
|
|
|
25 |
except ImportError:
|
26 |
thop = None
|
27 |
|
28 |
+
logger = logging.getLogger(__name__)
|
29 |
+
|
30 |
|
31 |
class Detect(nn.Module):
|
32 |
stride = None # strides computed during build
|