Update TorchScript suffix to `*.torchscript` (#5856)
Browse files- detect.py +4 -4
- export.py +3 -3
- models/common.py +6 -6
- utils/activations.py +2 -2
- val.py +5 -5
detect.py
CHANGED
@@ -81,18 +81,18 @@ def run(weights=ROOT / 'yolov5s.pt', # model.pt path(s)
|
|
81 |
imgsz = check_img_size(imgsz, s=stride) # check image size
|
82 |
|
83 |
# Half
|
84 |
-
half &= (pt or engine) and device.type != 'cpu' # half precision only supported by PyTorch on CUDA
|
85 |
-
if pt:
|
86 |
model.model.half() if half else model.model.float()
|
87 |
|
88 |
# Dataloader
|
89 |
if webcam:
|
90 |
view_img = check_imshow()
|
91 |
cudnn.benchmark = True # set True to speed up constant image size inference
|
92 |
-
dataset = LoadStreams(source, img_size=imgsz, stride=stride, auto=pt
|
93 |
bs = len(dataset) # batch_size
|
94 |
else:
|
95 |
-
dataset = LoadImages(source, img_size=imgsz, stride=stride, auto=pt
|
96 |
bs = 1 # batch_size
|
97 |
vid_path, vid_writer = [None] * bs, [None] * bs
|
98 |
|
|
|
81 |
imgsz = check_img_size(imgsz, s=stride) # check image size
|
82 |
|
83 |
# Half
|
84 |
+
half &= (pt or jit or engine) and device.type != 'cpu' # half precision only supported by PyTorch on CUDA
|
85 |
+
if pt or jit:
|
86 |
model.model.half() if half else model.model.float()
|
87 |
|
88 |
# Dataloader
|
89 |
if webcam:
|
90 |
view_img = check_imshow()
|
91 |
cudnn.benchmark = True # set True to speed up constant image size inference
|
92 |
+
dataset = LoadStreams(source, img_size=imgsz, stride=stride, auto=pt)
|
93 |
bs = len(dataset) # batch_size
|
94 |
else:
|
95 |
+
dataset = LoadImages(source, img_size=imgsz, stride=stride, auto=pt)
|
96 |
bs = 1 # batch_size
|
97 |
vid_path, vid_writer = [None] * bs, [None] * bs
|
98 |
|
export.py
CHANGED
@@ -5,7 +5,7 @@ Export a YOLOv5 PyTorch model to other formats. TensorFlow exports authored by h
|
|
5 |
Format | Example | Export `include=(...)` argument
|
6 |
--- | --- | ---
|
7 |
PyTorch | yolov5s.pt | -
|
8 |
-
TorchScript | yolov5s.torchscript
|
9 |
ONNX | yolov5s.onnx | 'onnx'
|
10 |
CoreML | yolov5s.mlmodel | 'coreml'
|
11 |
TensorFlow SavedModel | yolov5s_saved_model/ | 'saved_model'
|
@@ -19,7 +19,7 @@ Usage:
|
|
19 |
|
20 |
Inference:
|
21 |
$ python path/to/detect.py --weights yolov5s.pt
|
22 |
-
yolov5s.torchscript
|
23 |
yolov5s.onnx
|
24 |
yolov5s.mlmodel (under development)
|
25 |
yolov5s_saved_model
|
@@ -66,7 +66,7 @@ def export_torchscript(model, im, file, optimize, prefix=colorstr('TorchScript:'
|
|
66 |
# YOLOv5 TorchScript model export
|
67 |
try:
|
68 |
LOGGER.info(f'\n{prefix} starting export with torch {torch.__version__}...')
|
69 |
-
f = file.with_suffix('.torchscript
|
70 |
|
71 |
ts = torch.jit.trace(model, im, strict=False)
|
72 |
d = {"shape": im.shape, "stride": int(max(model.stride)), "names": model.names}
|
|
|
5 |
Format | Example | Export `include=(...)` argument
|
6 |
--- | --- | ---
|
7 |
PyTorch | yolov5s.pt | -
|
8 |
+
TorchScript | yolov5s.torchscript | 'torchscript'
|
9 |
ONNX | yolov5s.onnx | 'onnx'
|
10 |
CoreML | yolov5s.mlmodel | 'coreml'
|
11 |
TensorFlow SavedModel | yolov5s_saved_model/ | 'saved_model'
|
|
|
19 |
|
20 |
Inference:
|
21 |
$ python path/to/detect.py --weights yolov5s.pt
|
22 |
+
yolov5s.torchscript
|
23 |
yolov5s.onnx
|
24 |
yolov5s.mlmodel (under development)
|
25 |
yolov5s_saved_model
|
|
|
66 |
# YOLOv5 TorchScript model export
|
67 |
try:
|
68 |
LOGGER.info(f'\n{prefix} starting export with torch {torch.__version__}...')
|
69 |
+
f = file.with_suffix('.torchscript')
|
70 |
|
71 |
ts = torch.jit.trace(model, im, strict=False)
|
72 |
d = {"shape": im.shape, "stride": int(max(model.stride)), "names": model.names}
|
models/common.py
CHANGED
@@ -279,7 +279,7 @@ class DetectMultiBackend(nn.Module):
|
|
279 |
def __init__(self, weights='yolov5s.pt', device=None, dnn=True):
|
280 |
# Usage:
|
281 |
# PyTorch: weights = *.pt
|
282 |
-
# TorchScript: *.torchscript
|
283 |
# CoreML: *.mlmodel
|
284 |
# TensorFlow: *_saved_model
|
285 |
# TensorFlow: *.pb
|
@@ -289,10 +289,10 @@ class DetectMultiBackend(nn.Module):
|
|
289 |
# TensorRT: *.engine
|
290 |
super().__init__()
|
291 |
w = str(weights[0] if isinstance(weights, list) else weights)
|
292 |
-
suffix
|
|
|
293 |
check_suffix(w, suffixes) # check weights have acceptable suffix
|
294 |
-
pt, onnx, engine, tflite, pb, saved_model, coreml = (suffix == x for x in suffixes) # backend booleans
|
295 |
-
jit = pt and 'torchscript' in w.lower()
|
296 |
stride, names = 64, [f'class{i}' for i in range(1000)] # assign defaults
|
297 |
|
298 |
if jit: # TorchScript
|
@@ -304,10 +304,10 @@ class DetectMultiBackend(nn.Module):
|
|
304 |
stride, names = int(d['stride']), d['names']
|
305 |
elif pt: # PyTorch
|
306 |
from models.experimental import attempt_load # scoped to avoid circular import
|
307 |
-
model =
|
308 |
stride = int(model.stride.max()) # model stride
|
309 |
names = model.module.names if hasattr(model, 'module') else model.names # get class names
|
310 |
-
elif coreml: # CoreML
|
311 |
import coremltools as ct
|
312 |
model = ct.models.MLModel(w)
|
313 |
elif dnn: # ONNX OpenCV DNN
|
|
|
279 |
def __init__(self, weights='yolov5s.pt', device=None, dnn=True):
|
280 |
# Usage:
|
281 |
# PyTorch: weights = *.pt
|
282 |
+
# TorchScript: *.torchscript
|
283 |
# CoreML: *.mlmodel
|
284 |
# TensorFlow: *_saved_model
|
285 |
# TensorFlow: *.pb
|
|
|
289 |
# TensorRT: *.engine
|
290 |
super().__init__()
|
291 |
w = str(weights[0] if isinstance(weights, list) else weights)
|
292 |
+
suffix = Path(w).suffix.lower()
|
293 |
+
suffixes = ['.pt', '.torchscript', '.onnx', '.engine', '.tflite', '.pb', '', '.mlmodel']
|
294 |
check_suffix(w, suffixes) # check weights have acceptable suffix
|
295 |
+
pt, jit, onnx, engine, tflite, pb, saved_model, coreml = (suffix == x for x in suffixes) # backend booleans
|
|
|
296 |
stride, names = 64, [f'class{i}' for i in range(1000)] # assign defaults
|
297 |
|
298 |
if jit: # TorchScript
|
|
|
304 |
stride, names = int(d['stride']), d['names']
|
305 |
elif pt: # PyTorch
|
306 |
from models.experimental import attempt_load # scoped to avoid circular import
|
307 |
+
model = attempt_load(weights, map_location=device)
|
308 |
stride = int(model.stride.max()) # model stride
|
309 |
names = model.module.names if hasattr(model, 'module') else model.names # get class names
|
310 |
+
elif coreml: # CoreML
|
311 |
import coremltools as ct
|
312 |
model = ct.models.MLModel(w)
|
313 |
elif dnn: # ONNX OpenCV DNN
|
utils/activations.py
CHANGED
@@ -18,8 +18,8 @@ class SiLU(nn.Module): # export-friendly version of nn.SiLU()
|
|
18 |
class Hardswish(nn.Module): # export-friendly version of nn.Hardswish()
|
19 |
@staticmethod
|
20 |
def forward(x):
|
21 |
-
# return x * F.hardsigmoid(x) # for
|
22 |
-
return x * F.hardtanh(x + 3, 0.0, 6.0) / 6.0 # for
|
23 |
|
24 |
|
25 |
# Mish https://github.com/digantamisra98/Mish --------------------------------------------------------------------------
|
|
|
18 |
class Hardswish(nn.Module): # export-friendly version of nn.Hardswish()
|
19 |
@staticmethod
|
20 |
def forward(x):
|
21 |
+
# return x * F.hardsigmoid(x) # for TorchScript and CoreML
|
22 |
+
return x * F.hardtanh(x + 3, 0.0, 6.0) / 6.0 # for TorchScript, CoreML and ONNX
|
23 |
|
24 |
|
25 |
# Mish https://github.com/digantamisra98/Mish --------------------------------------------------------------------------
|
val.py
CHANGED
@@ -111,7 +111,7 @@ def run(data,
|
|
111 |
# Initialize/load model and set device
|
112 |
training = model is not None
|
113 |
if training: # called by train.py
|
114 |
-
device, pt, engine = next(model.parameters()).device, True, False # get model device, PyTorch model
|
115 |
|
116 |
half &= device.type != 'cpu' # half precision only supported on CUDA
|
117 |
model.half() if half else model.float()
|
@@ -124,10 +124,10 @@ def run(data,
|
|
124 |
|
125 |
# Load model
|
126 |
model = DetectMultiBackend(weights, device=device, dnn=dnn)
|
127 |
-
stride, pt, engine = model.stride, model.pt, model.engine
|
128 |
imgsz = check_img_size(imgsz, s=stride) # check image size
|
129 |
-
half &= (pt or engine) and device.type != 'cpu' # half precision only supported by PyTorch on CUDA
|
130 |
-
if pt:
|
131 |
model.model.half() if half else model.model.float()
|
132 |
elif engine:
|
133 |
batch_size = model.batch_size
|
@@ -166,7 +166,7 @@ def run(data,
|
|
166 |
pbar = tqdm(dataloader, desc=s, bar_format='{l_bar}{bar:10}{r_bar}{bar:-10b}') # progress bar
|
167 |
for batch_i, (im, targets, paths, shapes) in enumerate(pbar):
|
168 |
t1 = time_sync()
|
169 |
-
if pt or engine:
|
170 |
im = im.to(device, non_blocking=True)
|
171 |
targets = targets.to(device)
|
172 |
im = im.half() if half else im.float() # uint8 to fp16/32
|
|
|
111 |
# Initialize/load model and set device
|
112 |
training = model is not None
|
113 |
if training: # called by train.py
|
114 |
+
device, pt, jit, engine = next(model.parameters()).device, True, False, False # get model device, PyTorch model
|
115 |
|
116 |
half &= device.type != 'cpu' # half precision only supported on CUDA
|
117 |
model.half() if half else model.float()
|
|
|
124 |
|
125 |
# Load model
|
126 |
model = DetectMultiBackend(weights, device=device, dnn=dnn)
|
127 |
+
stride, pt, jit, engine = model.stride, model.pt, model.jit, model.engine
|
128 |
imgsz = check_img_size(imgsz, s=stride) # check image size
|
129 |
+
half &= (pt or jit or engine) and device.type != 'cpu' # half precision only supported by PyTorch on CUDA
|
130 |
+
if pt or jit:
|
131 |
model.model.half() if half else model.model.float()
|
132 |
elif engine:
|
133 |
batch_size = model.batch_size
|
|
|
166 |
pbar = tqdm(dataloader, desc=s, bar_format='{l_bar}{bar:10}{r_bar}{bar:-10b}') # progress bar
|
167 |
for batch_i, (im, targets, paths, shapes) in enumerate(pbar):
|
168 |
t1 = time_sync()
|
169 |
+
if pt or jit or engine:
|
170 |
im = im.to(device, non_blocking=True)
|
171 |
targets = targets.to(device)
|
172 |
im = im.half() if half else im.float() # uint8 to fp16/32
|