* Refactor code for speed and clarity

* Update dataloaders.py

Signed-off-by: Glenn Jocher <glenn.jocher@ultralytics.com>

---------

Signed-off-by: Glenn Jocher <glenn.jocher@ultralytics.com>
This commit is contained in:
Glenn Jocher authored and GitHub committed 2024-07-01 00:17:24 +02:00
1 parent b9019671f4
commit 100a423b66
14 files changed
+30

No files matched your search

+3
View File
@@ -1066,6 +1066,9 @@ class Classify(nn.Module):
def __init__(
self, c1, c2, k=1, s=1, p=None, g=1, dropout_p=0.0
): # ch_in, ch_out, kernel, stride, padding, groups, dropout probability
"""Initializes YOLOv5 classification head with convolution, pooling, and dropout layers for input to output
channel transformation.
"""
super().__init__()
c_ = 1280 # efficientnet_b0 size
self.conv = Conv(c1, c_, k, s, autopad(k, p), g)
+2
View File
@@ -612,6 +612,7 @@ class TFModel:
iou_thres=0.45,
conf_thres=0.25,
):
"""Runs inference on input data, with an option for TensorFlow NMS."""
y = [] # outputs
x = inputs
for m in self.model.layers:
@@ -730,6 +731,7 @@ def run(
dynamic=False, # dynamic batch size
):
# PyTorch model
"""Exports YOLOv5 model from PyTorch to TensorFlow and Keras formats, performing inference for validation."""
im = torch.zeros((batch_size, 3, *imgsz)) # BCHW image
model = attempt_load(weights, device=torch.device("cpu"), inplace=True, fuse=False)
_ = model(im) # inference