Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Updata #1

Merged
merged 50 commits into from
Mar 21, 2022
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
50 commits
Select commit Hold shift + click to select a range
4bab565
Update tutorial.ipynb (2 CPUs, 12.7 GB RAM, 42.2/166.8 GB disk) (#6767)
glenn-jocher Feb 24, 2022
47f265d
Update min warmup iterations from 1k to 100 (#6768)
glenn-jocher Feb 24, 2022
c161557
Default `OMP_NUM_THREADS=8` (#6770)
glenn-jocher Feb 24, 2022
741cd0e
Update tutorial.ipynb (#6771)
glenn-jocher Feb 24, 2022
cea994b
Update hyp.VOC.yaml (#6772)
glenn-jocher Feb 24, 2022
b2adc7c
Fix export for 1-channel images (#6780)
lcombaldieu Feb 25, 2022
0f81991
Update EMA decay `tau` (#6769)
glenn-jocher Feb 25, 2022
9ec51a6
YOLOv5s6 params FLOPs fix (#6782)
glenn-jocher Feb 25, 2022
c2403eb
Update PULL_REQUEST_TEMPLATE.md (#6783)
glenn-jocher Feb 25, 2022
63ddb6f
Update autoanchor.py (#6794)
glenn-jocher Feb 26, 2022
bcc92e2
Update sweep.yaml (#6825)
lcombaldieu Mar 4, 2022
601dbb8
AutoAnchor improved initialization robustness (#6854)
glenn-jocher Mar 4, 2022
8a66eba
Add `*.ts` to `VID_FORMATS` (#6859)
glenn-jocher Mar 4, 2022
4728840
Update `--cache disk` deprecate `*_npy/` dirs (#6876)
glenn-jocher Mar 6, 2022
7e98b48
Update yolov5s.yaml (#6865)
vnekat Mar 6, 2022
596de6d
Default FP16 TensorRT export (#6798)
DavidBaldsiefen Mar 6, 2022
c8a5899
Bump actions/setup-python from 2 to 3 (#6880)
dependabot[bot] Mar 7, 2022
a5a1760
Bump actions/checkout from 2 to 3 (#6881)
dependabot[bot] Mar 7, 2022
acc58c1
Fix TRT `max_workspace_size` deprecation notice (#6856)
glenn-jocher Mar 7, 2022
e6e36aa
Update bytes to GB with bitshift (#6886)
glenn-jocher Mar 7, 2022
6dd82c0
Move `git_describe()` to general.py (#6918)
glenn-jocher Mar 9, 2022
d3d9cbc
PyTorch 1.11.0 compatibility updates (#6932)
glenn-jocher Mar 10, 2022
055e72a
Optimize PyTorch 1.11.0 compatibility update (#6933)
glenn-jocher Mar 10, 2022
caf7ad0
Allow 3-point segments (#6938)
glenn-jocher Mar 10, 2022
84efa62
Fix PyTorch Hub export inference shapes (#6949)
glenn-jocher Mar 11, 2022
b94b59e
DetectMultiBackend() `--half` handling (#6945)
glenn-jocher Mar 11, 2022
c6b4f84
Update Dockerfile `torch==1.11.0+cu113` (#6954)
glenn-jocher Mar 11, 2022
c84dd27
New val.py `cuda` variable (#6957)
glenn-jocher Mar 12, 2022
52c1399
DetectMultiBackend() return `device` update (#6958)
glenn-jocher Mar 12, 2022
701e117
Tensor initialization on device improvements (#6959)
glenn-jocher Mar 12, 2022
c13d4ce
EdgeTPU optimizations (#6808)
paradigmn Mar 12, 2022
2d45de6
Model `ema` key backward compatibility fix (#6972)
glenn-jocher Mar 14, 2022
99de551
pt model to cpu on TF export
glenn-jocher Mar 14, 2022
932dc78
YOLOv5 Export Benchmarks for GPU (#6963)
glenn-jocher Mar 14, 2022
c09fb2a
Update TQDM bar format (#6988)
glenn-jocher Mar 15, 2022
3f634d4
Conditional `Timeout()` by OS (disable on Windows) (#7013)
glenn-jocher Mar 16, 2022
7c6a335
fix: add default PIL font as fallback (#7010)
maxstrobel Mar 17, 2022
4effd06
Consistent saved_model output format (#7032)
MrinalJain17 Mar 18, 2022
b0ba101
`ComputeLoss()` indexing/speed improvements (#7048)
glenn-jocher Mar 20, 2022
9ebec78
Update Dockerfile to `git clone` instead of `COPY` (#7053)
glenn-jocher Mar 20, 2022
6843ea5
Create SECURITY.md (#7054)
glenn-jocher Mar 20, 2022
9f4d113
Fix incomplete URL substring sanitation (#7056)
glenn-jocher Mar 20, 2022
529fbc1
Use PIL to eliminate chroma subsampling in crops (#7008)
LaserBorg Mar 20, 2022
f327eee
Fix `check_anchor_order()` in pixel-space not grid-space (#7060)
glenn-jocher Mar 20, 2022
d5e363f
Update detect.py non-inplace with `y.tensor_split()` (#7062)
glenn-jocher Mar 20, 2022
0529b77
Update common.py lists for tuples (#7063)
glenn-jocher Mar 20, 2022
e278fd6
Update W&B message to `LOGGER.info()` (#7064)
glenn-jocher Mar 20, 2022
9e75cbf
Update __init__.py (#7065)
glenn-jocher Mar 20, 2022
178c109
Add non-zero `da` `check_anchor_order()` condition (#7066)
glenn-jocher Mar 20, 2022
9cd89b7
Fix2 `check_anchor_order()` in pixel-space not grid-space (#7067)
glenn-jocher Mar 20, 2022
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions .github/PULL_REQUEST_TEMPLATE.md
Original file line number Diff line number Diff line change
@@ -1,7 +1,9 @@
<!--
Thank you for submitting a YOLOv5 🚀 Pull Request! We want to make contributing to YOLOv5 as easy and transparent as possible. A few tips to get you started:

- Search existing YOLOv5 [PRs](https://github.com/ultralytics/yolov5/pull) to see if a similar PR already exists.
- Link this PR to a YOLOv5 [issue](https://github.com/ultralytics/yolov5/issues) to help us understand what bug fix or feature is being implemented.
- Provide before and after profiling/inference/training results to help us quantify the improvement your PR provides (if applicable).

Please see our ✅ [Contributing Guide](https://github.com/ultralytics/yolov5/blob/master/CONTRIBUTING.md) for more details.
-->
7 changes: 7 additions & 0 deletions .github/SECURITY.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,7 @@
# Security Policy

We aim to make YOLOv5 🚀 as secure as possible! If you find potential vulnerabilities or have any concerns please let us know so we can investigate and take corrective action if needed.

### Reporting a Vulnerability

To report vulnerabilities please email us at [email protected] or visit https://ultralytics.com/contact. Thank you!
4 changes: 2 additions & 2 deletions .github/workflows/ci-testing.yml
Original file line number Diff line number Diff line change
Expand Up @@ -25,9 +25,9 @@ jobs:
# Timeout: https://stackoverflow.com/a/59076067/4521646
timeout-minutes: 60
steps:
- uses: actions/checkout@v2
- uses: actions/checkout@v3
- name: Set up Python ${{ matrix.python-version }}
uses: actions/setup-python@v2
uses: actions/setup-python@v3
with:
python-version: ${{ matrix.python-version }}

Expand Down
2 changes: 1 addition & 1 deletion .github/workflows/codeql-analysis.yml
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,7 @@ jobs:

steps:
- name: Checkout repository
uses: actions/checkout@v2
uses: actions/checkout@v3

# Initializes the CodeQL tools for scanning.
- name: Initialize CodeQL
Expand Down
2 changes: 1 addition & 1 deletion .github/workflows/rebase.yml
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,7 @@ jobs:
runs-on: ubuntu-latest
steps:
- name: Checkout the latest code
uses: actions/checkout@v2
uses: actions/checkout@v3
with:
token: ${{ secrets.ACTIONS_TOKEN }}
fetch-depth: 0 # otherwise, you will fail to push refs to dest repo
Expand Down
5 changes: 3 additions & 2 deletions Dockerfile
Original file line number Diff line number Diff line change
Expand Up @@ -11,15 +11,16 @@ COPY requirements.txt .
RUN python -m pip install --upgrade pip
RUN pip uninstall -y torch torchvision torchtext
RUN pip install --no-cache -r requirements.txt albumentations wandb gsutil notebook \
torch==1.10.2+cu113 torchvision==0.11.3+cu113 -f https://download.pytorch.org/whl/cu113/torch_stable.html
torch==1.11.0+cu113 torchvision==0.12.0+cu113 -f https://download.pytorch.org/whl/cu113/torch_stable.html
# RUN pip install --no-cache -U torch torchvision

# Create working directory
RUN mkdir -p /usr/src/app
WORKDIR /usr/src/app

# Copy contents
COPY . /usr/src/app
RUN git clone https://github.com/ultralytics/yolov5 /usr/src/app
# COPY . /usr/src/app

# Downloads to user config dir
ADD https://ultralytics.com/assets/Arial.ttf /root/.config/Ultralytics/
Expand Down
2 changes: 1 addition & 1 deletion README.md
Original file line number Diff line number Diff line change
Expand Up @@ -245,7 +245,7 @@ We are super excited about our first-ever Ultralytics YOLOv5 🚀 EXPORT Competi
|[YOLOv5x][assets] |640 |50.7 |68.9 |766 |12.1 |4.8 |86.7 |205.7
| | | | | | | | |
|[YOLOv5n6][assets] |1280 |36.0 |54.4 |153 |8.1 |2.1 |3.2 |4.6
|[YOLOv5s6][assets] |1280 |44.8 |63.7 |385 |8.2 |3.6 |16.8 |12.6
|[YOLOv5s6][assets] |1280 |44.8 |63.7 |385 |8.2 |3.6 |12.6 |16.8
|[YOLOv5m6][assets] |1280 |51.3 |69.3 |887 |11.1 |6.8 |35.7 |50.0
|[YOLOv5l6][assets] |1280 |53.7 |71.3 |1784 |15.8 |10.5 |76.8 |111.4
|[YOLOv5x6][assets]<br>+ [TTA][TTA]|1280<br>1536 |55.0<br>**55.8** |72.7<br>**72.7** |3136<br>- |26.2<br>- |19.4<br>- |140.7<br>- |209.8<br>-
Expand Down
42 changes: 21 additions & 21 deletions data/hyps/hyp.VOC.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -4,37 +4,37 @@
# See Hyperparameter Evolution tutorial for details https://github.com/ultralytics/yolov5#tutorials

# YOLOv5 Hyperparameter Evolution Results
# Best generation: 319
# Last generation: 434
# Best generation: 467
# Last generation: 996
# metrics/precision, metrics/recall, metrics/mAP_0.5, metrics/mAP_0.5:0.95, val/box_loss, val/obj_loss, val/cls_loss
# 0.86236, 0.86184, 0.91274, 0.72647, 0.0077056, 0.0042449, 0.0013846
# 0.87729, 0.85125, 0.91286, 0.72664, 0.0076739, 0.0042529, 0.0013865

lr0: 0.0033
lrf: 0.15184
momentum: 0.74747
lr0: 0.00334
lrf: 0.15135
momentum: 0.74832
weight_decay: 0.00025
warmup_epochs: 3.4278
warmup_momentum: 0.59032
warmup_bias_lr: 0.18742
warmup_epochs: 3.3835
warmup_momentum: 0.59462
warmup_bias_lr: 0.18657
box: 0.02
cls: 0.21563
cls: 0.21638
cls_pw: 0.5
obj: 0.50843
obj_pw: 0.6729
obj: 0.51728
obj_pw: 0.67198
iou_t: 0.2
anchor_t: 3.4172
anchor_t: 3.3744
fl_gamma: 0.0
hsv_h: 0.01032
hsv_s: 0.5562
hsv_v: 0.28255
hsv_h: 0.01041
hsv_s: 0.54703
hsv_v: 0.27739
degrees: 0.0
translate: 0.04575
scale: 0.73711
translate: 0.04591
scale: 0.75544
shear: 0.0
perspective: 0.0
flipud: 0.0
fliplr: 0.5
mosaic: 0.87158
mixup: 0.04294
mosaic: 0.85834
mixup: 0.04266
copy_paste: 0.0
anchors: 3.3556
anchors: 3.412
13 changes: 4 additions & 9 deletions detect.py
Original file line number Diff line number Diff line change
Expand Up @@ -89,15 +89,10 @@ def run(weights=ROOT / 'yolov5s.pt', # model.pt path(s)

# Load model
device = select_device(device)
model = DetectMultiBackend(weights, device=device, dnn=dnn, data=data)
stride, names, pt, jit, onnx, engine = model.stride, model.names, model.pt, model.jit, model.onnx, model.engine
model = DetectMultiBackend(weights, device=device, dnn=dnn, data=data, fp16=half)
stride, names, pt = model.stride, model.names, model.pt
imgsz = check_img_size(imgsz, s=stride) # check image size

# Half
half &= (pt or jit or onnx or engine) and device.type != 'cpu' # FP16 supported on limited backends with CUDA
if pt or jit:
model.model.half() if half else model.model.float()

# Dataloader
if webcam:
view_img = check_imshow()
Expand All @@ -110,12 +105,12 @@ def run(weights=ROOT / 'yolov5s.pt', # model.pt path(s)
vid_path, vid_writer = [None] * bs, [None] * bs

# Run inference
model.warmup(imgsz=(1 if pt else bs, 3, *imgsz), half=half) # warmup
model.warmup(imgsz=(1 if pt else bs, 3, *imgsz)) # warmup
dt, seen = [0.0, 0.0, 0.0], 0
for path, im, im0s, vid_cap, s in dataset:
t1 = time_sync()
im = torch.from_numpy(im).to(device)
im = im.half() if half else im.float() # uint8 to fp16/32
im = im.half() if model.fp16 else im.float() # uint8 to fp16/32
im /= 255 # 0 - 255 to 0.0 - 1.0
if len(im.shape) == 3:
im = im[None] # expand for batch dim
Expand Down
40 changes: 20 additions & 20 deletions export.py
Original file line number Diff line number Diff line change
Expand Up @@ -75,18 +75,18 @@

def export_formats():
# YOLOv5 export formats
x = [['PyTorch', '-', '.pt'],
['TorchScript', 'torchscript', '.torchscript'],
['ONNX', 'onnx', '.onnx'],
['OpenVINO', 'openvino', '_openvino_model'],
['TensorRT', 'engine', '.engine'],
['CoreML', 'coreml', '.mlmodel'],
['TensorFlow SavedModel', 'saved_model', '_saved_model'],
['TensorFlow GraphDef', 'pb', '.pb'],
['TensorFlow Lite', 'tflite', '.tflite'],
['TensorFlow Edge TPU', 'edgetpu', '_edgetpu.tflite'],
['TensorFlow.js', 'tfjs', '_web_model']]
return pd.DataFrame(x, columns=['Format', 'Argument', 'Suffix'])
x = [['PyTorch', '-', '.pt', True],
['TorchScript', 'torchscript', '.torchscript', True],
['ONNX', 'onnx', '.onnx', True],
['OpenVINO', 'openvino', '_openvino_model', False],
['TensorRT', 'engine', '.engine', True],
['CoreML', 'coreml', '.mlmodel', False],
['TensorFlow SavedModel', 'saved_model', '_saved_model', True],
['TensorFlow GraphDef', 'pb', '.pb', True],
['TensorFlow Lite', 'tflite', '.tflite', False],
['TensorFlow Edge TPU', 'edgetpu', '_edgetpu.tflite', False],
['TensorFlow.js', 'tfjs', '_web_model', False]]
return pd.DataFrame(x, columns=['Format', 'Argument', 'Suffix', 'GPU'])


def export_torchscript(model, im, file, optimize, prefix=colorstr('TorchScript:')):
Expand Down Expand Up @@ -218,6 +218,7 @@ def export_engine(model, im, file, train, half, simplify, workspace=4, verbose=F
builder = trt.Builder(logger)
config = builder.create_builder_config()
config.max_workspace_size = workspace * 1 << 30
# config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, workspace << 30) # fix TRT 8.4 deprecation notice

flag = (1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
network = builder.create_network(flag)
Expand All @@ -233,9 +234,8 @@ def export_engine(model, im, file, train, half, simplify, workspace=4, verbose=F
for out in outputs:
LOGGER.info(f'{prefix}\toutput "{out.name}" with shape {out.shape} and dtype {out.dtype}')

half &= builder.platform_has_fast_fp16
LOGGER.info(f'{prefix} building FP{16 if half else 32} engine in {f}')
if half:
LOGGER.info(f'{prefix} building FP{16 if builder.platform_has_fast_fp16 else 32} engine in {f}')
if builder.platform_has_fast_fp16:
config.set_flag(trt.BuilderFlag.FP16)
with builder.build_engine(network, config) as engine, open(f, 'wb') as t:
t.write(engine.serialize())
Expand All @@ -260,9 +260,9 @@ def export_saved_model(model, im, file, dynamic,
batch_size, ch, *imgsz = list(im.shape) # BCHW

tf_model = TFModel(cfg=model.yaml, model=model, nc=model.nc, imgsz=imgsz)
im = tf.zeros((batch_size, *imgsz, 3)) # BHWC order for TensorFlow
im = tf.zeros((batch_size, *imgsz, ch)) # BHWC order for TensorFlow
_ = tf_model.predict(im, tf_nms, agnostic_nms, topk_per_class, topk_all, iou_thres, conf_thres)
inputs = tf.keras.Input(shape=(*imgsz, 3), batch_size=None if dynamic else batch_size)
inputs = tf.keras.Input(shape=(*imgsz, ch), batch_size=None if dynamic else batch_size)
outputs = tf_model.predict(inputs, tf_nms, agnostic_nms, topk_per_class, topk_all, iou_thres, conf_thres)
keras_model = tf.keras.Model(inputs=inputs, outputs=outputs)
keras_model.trainable = False
Expand All @@ -275,7 +275,7 @@ def export_saved_model(model, im, file, dynamic,
m = m.get_concrete_function(spec)
frozen_func = convert_variables_to_constants_v2(m)
tfm = tf.Module()
tfm.__call__ = tf.function(lambda x: frozen_func(x), [spec])
tfm.__call__ = tf.function(lambda x: frozen_func(x)[0], [spec])
tfm.__call__(im)
tf.saved_model.save(
tfm,
Expand Down Expand Up @@ -331,7 +331,7 @@ def export_tflite(keras_model, im, file, int8, data, ncalib, prefix=colorstr('Te
converter.target_spec.supported_types = []
converter.inference_input_type = tf.uint8 # or tf.int8
converter.inference_output_type = tf.uint8 # or tf.int8
converter.experimental_new_quantizer = False
converter.experimental_new_quantizer = True
f = str(file).replace('.pt', '-int8.tflite')

tflite_model = converter.convert()
Expand Down Expand Up @@ -494,7 +494,7 @@ def run(data=ROOT / 'data/coco128.yaml', # 'dataset.yaml path'
if int8 or edgetpu: # TFLite --int8 bug https://github.com/ultralytics/yolov5/issues/5707
check_requirements(('flatbuffers==1.12',)) # required before `import tensorflow`
assert not (tflite and tfjs), 'TFLite and TF.js models must be exported separately, please pass only one type.'
model, f[5] = export_saved_model(model, im, file, dynamic, tf_nms=nms or agnostic_nms or tfjs,
model, f[5] = export_saved_model(model.cpu(), im, file, dynamic, tf_nms=nms or agnostic_nms or tfjs,
agnostic_nms=agnostic_nms or tfjs, topk_per_class=topk_per_class,
topk_all=topk_all, conf_thres=conf_thres, iou_thres=iou_thres) # keras model
if pb or tfjs: # pb prerequisite to tfjs
Expand Down
41 changes: 24 additions & 17 deletions models/common.py
Original file line number Diff line number Diff line change
Expand Up @@ -31,7 +31,7 @@
def autopad(k, p=None): # kernel, padding
# Pad to 'same'
if p is None:
p = k // 2 if isinstance(k, int) else [x // 2 for x in k] # auto-pad
p = k // 2 if isinstance(k, int) else (x // 2 for x in k) # auto-pad
return p


Expand Down Expand Up @@ -133,7 +133,7 @@ def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5): # ch_in, ch_out, nu
self.cv2 = Conv(c1, c_, 1, 1)
self.cv3 = Conv(2 * c_, c2, 1) # act=FReLU(c2)
self.m = nn.Sequential(*(Bottleneck(c_, c_, shortcut, g, e=1.0) for _ in range(n)))
# self.m = nn.Sequential(*[CrossConv(c_, c_, 3, 1, g, 1.0, shortcut) for _ in range(n)])
# self.m = nn.Sequential(*(CrossConv(c_, c_, 3, 1, g, 1.0, shortcut) for _ in range(n)))

def forward(self, x):
return self.cv3(torch.cat((self.m(self.cv1(x)), self.cv2(x)), dim=1))
Expand Down Expand Up @@ -194,7 +194,7 @@ def forward(self, x):
warnings.simplefilter('ignore') # suppress torch 1.9.0 max_pool2d() warning
y1 = self.m(x)
y2 = self.m(y1)
return self.cv2(torch.cat([x, y1, y2, self.m(y2)], 1))
return self.cv2(torch.cat((x, y1, y2, self.m(y2)), 1))


class Focus(nn.Module):
Expand All @@ -205,7 +205,7 @@ def __init__(self, c1, c2, k=1, s=1, p=None, g=1, act=True): # ch_in, ch_out, k
# self.contract = Contract(gain=2)

def forward(self, x): # x(b,c,w,h) -> y(b,4c,w/2,h/2)
return self.conv(torch.cat([x[..., ::2, ::2], x[..., 1::2, ::2], x[..., ::2, 1::2], x[..., 1::2, 1::2]], 1))
return self.conv(torch.cat((x[..., ::2, ::2], x[..., 1::2, ::2], x[..., ::2, 1::2], x[..., 1::2, 1::2]), 1))
# return self.conv(self.contract(x))


Expand All @@ -219,7 +219,7 @@ def __init__(self, c1, c2, k=1, s=1, g=1, act=True): # ch_in, ch_out, kernel, s

def forward(self, x):
y = self.cv1(x)
return torch.cat([y, self.cv2(y)], 1)
return torch.cat((y, self.cv2(y)), 1)


class GhostBottleneck(nn.Module):
Expand Down Expand Up @@ -277,7 +277,7 @@ def forward(self, x):

class DetectMultiBackend(nn.Module):
# YOLOv5 MultiBackend class for python inference on various backends
def __init__(self, weights='yolov5s.pt', device=None, dnn=False, data=None):
def __init__(self, weights='yolov5s.pt', device=torch.device('cpu'), dnn=False, data=None, fp16=False):
# Usage:
# PyTorch: weights = *.pt
# TorchScript: *.torchscript
Expand All @@ -297,6 +297,7 @@ def __init__(self, weights='yolov5s.pt', device=None, dnn=False, data=None):
pt, jit, onnx, xml, engine, coreml, saved_model, pb, tflite, edgetpu, tfjs = self.model_type(w) # get backend
stride, names = 64, [f'class{i}' for i in range(1000)] # assign defaults
w = attempt_download(w) # download if not local
fp16 &= (pt or jit or onnx or engine) and device.type != 'cpu' # FP16
if data: # data.yaml path (optional)
with open(data, errors='ignore') as f:
names = yaml.safe_load(f)['names'] # class names
Expand All @@ -305,11 +306,13 @@ def __init__(self, weights='yolov5s.pt', device=None, dnn=False, data=None):
model = attempt_load(weights if isinstance(weights, list) else w, map_location=device)
stride = max(int(model.stride.max()), 32) # model stride
names = model.module.names if hasattr(model, 'module') else model.names # get class names
model.half() if fp16 else model.float()
self.model = model # explicitly assign for to(), cpu(), cuda(), half()
elif jit: # TorchScript
LOGGER.info(f'Loading {w} for TorchScript inference...')
extra_files = {'config.txt': ''} # model metadata
model = torch.jit.load(w, _extra_files=extra_files)
model.half() if fp16 else model.float()
if extra_files['config.txt']:
d = json.loads(extra_files['config.txt']) # extra_files dict
stride, names = int(d['stride']), d['names']
Expand Down Expand Up @@ -342,12 +345,15 @@ def __init__(self, weights='yolov5s.pt', device=None, dnn=False, data=None):
with open(w, 'rb') as f, trt.Runtime(logger) as runtime:
model = runtime.deserialize_cuda_engine(f.read())
bindings = OrderedDict()
fp16 = False # default updated below
for index in range(model.num_bindings):
name = model.get_binding_name(index)
dtype = trt.nptype(model.get_binding_dtype(index))
shape = tuple(model.get_binding_shape(index))
data = torch.from_numpy(np.empty(shape, dtype=np.dtype(dtype))).to(device)
bindings[name] = Binding(name, dtype, shape, data, int(data.data_ptr()))
if model.binding_is_input(index) and dtype == np.float16:
fp16 = True
binding_addrs = OrderedDict((n, d.ptr) for n, d in bindings.items())
context = model.create_execution_context()
batch_size = bindings['images'].shape[0]
Expand Down Expand Up @@ -435,7 +441,7 @@ def forward(self, im, augment=False, visualize=False, val=False):
else: # TensorFlow (SavedModel, GraphDef, Lite, Edge TPU)
im = im.permute(0, 2, 3, 1).cpu().numpy() # torch BCHW to numpy BHWC shape(1,320,192,3)
if self.saved_model: # SavedModel
y = (self.model(im, training=False) if self.keras else self.model(im)[0]).numpy()
y = (self.model(im, training=False) if self.keras else self.model(im)).numpy()
elif self.pb: # GraphDef
y = self.frozen_func(x=self.tf.constant(im)).numpy()
else: # Lite or Edge TPU
Expand All @@ -452,15 +458,17 @@ def forward(self, im, augment=False, visualize=False, val=False):
y = (y.astype(np.float32) - zero_point) * scale # re-scale
y[..., :4] *= [w, h, w, h] # xywh normalized to pixels

y = torch.tensor(y) if isinstance(y, np.ndarray) else y
if isinstance(y, np.ndarray):
y = torch.tensor(y, device=self.device)
return (y, []) if val else y

def warmup(self, imgsz=(1, 3, 640, 640), half=False):
def warmup(self, imgsz=(1, 3, 640, 640)):
# Warmup model by running inference once
if self.pt or self.jit or self.onnx or self.engine: # warmup types
if isinstance(self.device, torch.device) and self.device.type != 'cpu': # only warmup GPU models
im = torch.zeros(*imgsz).to(self.device).type(torch.half if half else torch.float) # input image
self.forward(im) # warmup
if any((self.pt, self.jit, self.onnx, self.engine, self.saved_model, self.pb)): # warmup types
if self.device.type != 'cpu': # only warmup GPU models
im = torch.zeros(*imgsz, dtype=torch.half if self.fp16 else torch.float, device=self.device) # input
for _ in range(2 if self.jit else 1): #
self.forward(im) # warmup

@staticmethod
def model_type(p='path/to/model.pt'):
Expand Down Expand Up @@ -541,10 +549,9 @@ def forward(self, imgs, size=640, augment=False, profile=False):
g = (size / max(s)) # gain
shape1.append([y * g for y in s])
imgs[i] = im if im.data.contiguous else np.ascontiguousarray(im) # update
shape1 = [make_divisible(x, self.stride) for x in np.stack(shape1, 0).max(0)] # inference shape
x = [letterbox(im, new_shape=shape1 if self.pt else size, auto=False)[0] for im in imgs] # pad
x = np.stack(x, 0) if n > 1 else x[0][None] # stack
x = np.ascontiguousarray(x.transpose((0, 3, 1, 2))) # BHWC to BCHW
shape1 = [make_divisible(x, self.stride) if self.pt else size for x in np.array(shape1).max(0)] # inf shape
x = [letterbox(im, new_shape=shape1, auto=False)[0] for im in imgs] # pad
x = np.ascontiguousarray(np.array(x).transpose((0, 3, 1, 2))) # stack and BHWC to BCHW
x = torch.from_numpy(x).to(p.device).type_as(p) / 255 # uint8 to fp16/32
t.append(time_sync())

Expand Down
Loading