glenn-jocher commited on
Commit
53ed872
1 Parent(s): 095197b

Update export.py, yolo.py `sys.path.append()` (#3579)

Browse files
Files changed (2) hide show
  1. models/export.py +7 -5
  2. 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
- import models
 
 
 
 
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, models.common.Conv): # assign export-friendly activations
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, models.yolo.Detect):
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
- sys.path.append(Path(__file__).parent.parent.absolute().__str__()) # to run '$ python *.py' files in subdirectories
14
- logger = logging.getLogger(__name__)
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