1. Architecture Overview
YOLOv5 is a single-stage, anchor-based object detection framework implemented in PyTorch. The codebase is organized as a script-driven repository — each task (detection, segmentation, classification) has its own set of train.py, predict.py/detect.py, and val.py entry points. A shared utilities layer provides model definitions, data loading, augmentation, loss computation, logging, and multi-backend inference.
The architecture follows a three-stage pipeline common to YOLO-family detectors:
Input Image
│
▼
┌─────────────────┐
│ Backbone │ Feature extraction (CSP-based)
│ Focus/Conv → C3 │ Downsamples through P1-P5 (or P6)
│ → SPPF │
└────────┬────────┘
│
▼
┌─────────────────┐
│ Neck │ Feature aggregation (PANet / FPN)
│ Upsample + C3 │ Top-down + bottom-up paths
│ Concat layers │
└────────┬────────┘
│
▼
┌─────────────────┐
│ Head │ Task-specific prediction
│ Detect / Segment│ Anchors × (5 + nc) per grid cell
│ / Classify │ (or + nm for segmentation)
└─────────────────┘
Component Inventory
| Component | Path | Role |
|---|---|---|
| Detection inference | detect.py |
Run detection on images/videos/streams |
| Detection training | train.py |
Train detection models |
| Detection validation | val.py |
Evaluate detection models |
| Segmentation inference | segment/predict.py |
Run segmentation inference |
| Segmentation training | segment/train.py |
Train segmentation models |
| Segmentation validation | segment/val.py |
Evaluate segmentation models |
| Classification inference | classify/predict.py |
Run classification inference |
| Classification training | classify/train.py |
Train classification models |
| Classification validation | classify/val.py |
Evaluate classification models |
| Model export | export.py |
Export to 11+ deployment formats |
| Benchmarks | benchmarks.py |
Compare export format performance |
| PyTorch Hub | hubconf.py |
torch.hub.load() entry points |
| Model definitions | models/common.py |
Layers, DetectMultiBackend, AutoShape |
| Model architecture | models/yolo.py |
DetectionModel, SegmentationModel, ClassificationModel |
| Experimental | models/experimental.py |
Ensemble, model loading |
| TF export | models/tf.py |
TensorFlow/Keras model equivalents |
| Model configs | models/*.yaml |
Architecture definitions (n/s/m/l/x) |
| Hub variants | models/hub/*.yaml |
Experimental architectures (BiFPN, ghost, P6, etc.) |
| Segmentation configs | models/segment/*.yaml |
Segmentation model architectures |
| Utility core | utils/general.py |
NMS, box ops, downloads, config, checks |
| Data loading | utils/dataloaders.py |
Image/video/stream loaders, dataset classes |
| Augmentations | utils/augmentations.py |
HSV, mosaic, mixup, letterbox, copy-paste |
| Loss functions | utils/loss.py |
ComputeLoss, FocalLoss, BCE |
| Metrics | utils/metrics.py |
mAP, confusion matrix, IoU |
| Plotting | utils/plots.py |
Training curves, label plots, predictions |
| Torch utilities | utils/torch_utils.py |
Device selection, EarlyStopping, EMA, optimizer |
| Anchor utilities | utils/autoanchor.py |
AutoAnchor, k-means anchor generation |
| Auto batch | utils/autobatch.py |
Automatic batch size selection |
| Downloads | utils/downloads.py |
Weight download, GitHub asset fetching |
| Callbacks | utils/callbacks.py |
Hook-based event system |
| Activations | utils/activations.py |
SiLU, Hardswish, Mish, FReLU, AconC |
| Triton client | utils/triton.py |
NVIDIA Triton Inference Server client |
| Loggers | utils/loggers/__init__.py |
Unified logging (TB, W&B, ClearML, Comet) |
| ClearML integration | utils/loggers/clearml/ |
ClearML experiment tracking + HPO |
| Comet integration | utils/loggers/comet/ |
Comet experiment tracking + HPO |
| W&B integration | utils/loggers/wandb/ |
Weights & Biases logging |
| Segmentation utils | utils/segment/ |
Seg-specific augmentations, loss, metrics, dataloaders |
| Flask REST API | utils/flask_rest_api/ |
Simple REST serving endpoint |
| Docker configs | utils/docker/ |
GPU, CPU, ARM64 Dockerfiles |
| AWS utilities | utils/aws/ |
EC2 resume, MIME, userdata scripts |
| GAE config | utils/google_app_engine/ |
Google App Engine deployment |
| Dataset configs | data/*.yaml |
14 dataset definitions |
| Hyperparameters | data/hyps/*.yaml |
6 hyperparameter presets |
Note
The
models/hub/directory contains experimental architecture variants (BiFPN, ghost convolutions, transformer, P6/P7 detection layers) that are not included in the standard pretrained weights but can be used with--cfgfor custom training.
Technology Stack
| Layer | Technology |
|---|---|
| Framework | PyTorch (>=1.8.0) |
| Model definition | torch.nn.Module subclasses |
| Training | Native PyTorch DDP, AMP |
| Data loading | torch.utils.data.DataLoader |
| Computer vision | OpenCV, PIL |
| Logging | TensorBoard, W&B, ClearML, Comet |
| Export | TorchScript, ONNX, TensorFlow/Keras, TensorRT, CoreML, OpenVINO, PaddlePaddle |
| Serving | Flask (bundled), Triton (client) |
| Containerization | Docker (GPU, CPU, ARM64) |
| License | AGPL-3.0 |
Data Flow — Training
Dataset YAML → check_dataset() → LoadImagesAndLabels
│ │
│ Augmentations
│ (mosaic, HSV, mixup,
│ copy-paste, letterbox)
│ │
▼ ▼
Model(cfg/weights) InfiniteDataLoader
│ │
├── Forward pass ◄──────────────────┘
│ │
│ ComputeLoss (box + obj + cls)
│ │
│ Backward + optimizer.step()
│ │
│ EMA update
│ │
│ EarlyStopping check
│ │
└── Loggers (TB/W&B/ClearML/Comet)
│
Save best.pt / last.pt
Data Flow — Inference
Source (image/video/stream/URL/screen/webcam)
│
▼
LoadImages / LoadStreams / LoadScreenshots
│
▼
DetectMultiBackend(weights) ←── auto-detects format
│ (PT/ONNX/TRT/OV/TF/CoreML/Paddle/Triton)
│
▼
Preprocessing (letterbox, normalize, to tensor)
│
▼
Model forward pass
│
▼
Non-Max Suppression (conf_thres + iou_thres)
│
▼
Post-processing (scale boxes, annotate, save)
2. System Requirements
Hardware Requirements
| Deployment | CPU | GPU | RAM | Disk |
|---|---|---|---|---|
| Training (small dataset) | Any modern x86_64 | 1x GPU, 4+ GB VRAM | 16 GB | 20 GB |
| Training (COCO) | 8+ cores | 1-8x GPU, 8+ GB VRAM | 32+ GB | 200+ GB |
| Inference (real-time) | 4+ cores | 1x GPU, 2+ GB VRAM | 8 GB | 5 GB |
| Inference (CPU-only) | 4+ cores | None | 8 GB | 5 GB |
| Edge (TFLite/ONNX) | ARM/x86 | Optional (Coral TPU) | 2+ GB | 1 GB |
Software Requirements
| Component | Version | Notes |
|---|---|---|
| Python | 3.8 - 3.12 | 3.10+ recommended |
| PyTorch | >=1.8.0 | 2.0+ for torch.compile |
| torchvision | >=0.9.0 | Must match PyTorch version |
| CUDA | 10.2+ | For GPU training/inference |
| cuDNN | 8.0+ | For GPU acceleration |
| GCC/G++ | 7+ | For compiling extensions |
| Docker | 20.10+ | For containerized deployment |
3. Configuration Reference
Model Architecture YAML
Model architectures are defined in YAML files under models/. Structure:
# yolov5s.yaml
nc: 80 # number of classes
depth_multiple: 0.33 # model depth multiple
width_multiple: 0.50 # layer channel multiple
anchors:
- [10,13, 16,30, 33,23] # P3/8
- [30,61, 62,45, 59,119] # P4/16
- [116,90, 156,198, 373,326] # P5/32
backbone:
# [from, number, module, args]
[[-1, 1, Conv, [64, 6, 2, 2]], # 0-P1/2
[-1, 1, Conv, [128, 3, 2]], # 1-P2/4
[-1, 3, C3, [128]], # 2
[-1, 1, Conv, [256, 3, 2]], # 3-P3/8
[-1, 6, C3, [256]], # 4
[-1, 1, Conv, [512, 3, 2]], # 5-P4/16
[-1, 9, C3, [512]], # 6
[-1, 1, Conv, [1024, 3, 2]], # 7-P5/32
[-1, 3, C3, [1024]], # 8
[-1, 1, SPPF, [1024, 5]], # 9
]
head:
[[-1, 1, Conv, [512, 1, 1]],
[-1, 1, nn.Upsample, [None, 2, 'nearest']],
[[-1, 6], 1, Concat, [1]],
[-1, 3, C3, [512, False]],
# ... (neck and detect layers)
[[17, 20, 23], 1, Detect, [nc, anchors]],
]
Important
The
depth_multipleandwidth_multipleparameters control model scaling. These are what differentiateyolov5n(0.33/0.25) fromyolov5x(1.33/1.25). Changing these values creates custom-sized models.
Available Architecture Variants
| Config | depth_multiple | width_multiple | Use |
|---|---|---|---|
yolov5n.yaml |
0.33 | 0.25 | Nano — edge/mobile |
yolov5s.yaml |
0.33 | 0.50 | Small — balanced |
yolov5m.yaml |
0.67 | 0.75 | Medium |
yolov5l.yaml |
1.00 | 1.00 | Large |
yolov5x.yaml |
1.33 | 1.25 | Extra-large |
Hub Variants (experimental)
| Config | Description |
|---|---|
yolov5-bifpn.yaml |
Bi-directional FPN neck |
yolov5-fpn.yaml |
Classic FPN (no PANet bottom-up) |
yolov5-p2.yaml |
Extra P2 small-object layer |
yolov5-p34.yaml |
P3+P4 only (2-layer head) |
yolov5-p6.yaml |
Extra P6 layer (1280px) |
yolov5-p7.yaml |
Extra P6+P7 layers |
yolov5-panet.yaml |
PANet variant |
yolov5s-ghost.yaml |
GhostNet-style convolutions |
yolov5s-transformer.yaml |
Transformer block in C3 |
yolov5s-LeakyReLU.yaml |
LeakyReLU activation |
yolov5n6.yaml - yolov5x6.yaml |
P6 models for each size |
yolov3.yaml, yolov3-spp.yaml, yolov3-tiny.yaml |
YOLOv3 legacy |
anchors.yaml |
Anchor presets |
Hyperparameter Configuration
Hyperparameters are stored in YAML files under data/hyps/. Complete reference from hyp.scratch-low.yaml:
| Parameter | Default | Description |
|---|---|---|
lr0 |
0.01 |
Initial learning rate (SGD=1e-2, Adam=1e-3) |
lrf |
0.01 |
Final OneCycleLR rate = lr0 * lrf |
momentum |
0.937 |
SGD momentum / Adam beta1 |
weight_decay |
0.0005 |
Optimizer weight decay |
warmup_epochs |
3.0 |
Warmup epochs (fractions allowed) |
warmup_momentum |
0.8 |
Warmup initial momentum |
warmup_bias_lr |
0.1 |
Warmup initial bias learning rate |
box |
0.05 |
Box loss gain |
cls |
0.5 |
Classification loss gain |
cls_pw |
1.0 |
Classification BCELoss positive_weight |
obj |
1.0 |
Objectness loss gain (scales with pixels) |
obj_pw |
1.0 |
Objectness BCELoss positive_weight |
iou_t |
0.20 |
IoU training threshold |
anchor_t |
4.0 |
Anchor-multiple threshold |
fl_gamma |
0.0 |
Focal loss gamma (0 = disabled) |
hsv_h |
0.015 |
HSV-Hue augmentation fraction |
hsv_s |
0.7 |
HSV-Saturation augmentation fraction |
hsv_v |
0.4 |
HSV-Value augmentation fraction |
degrees |
0.0 |
Rotation augmentation (+/- degrees) |
translate |
0.1 |
Translation augmentation (+/- fraction) |
scale |
0.5 |
Scale augmentation (+/- gain) |
shear |
0.0 |
Shear augmentation (+/- degrees) |
perspective |
0.0 |
Perspective augmentation (fraction, 0-0.001) |
flipud |
0.0 |
Vertical flip probability |
fliplr |
0.5 |
Horizontal flip probability |
mosaic |
1.0 |
Mosaic augmentation probability |
mixup |
0.0 |
MixUp augmentation probability |
copy_paste |
0.0 |
Copy-paste augmentation probability |
Preset Profiles
| Profile | Target | Key Differences |
|---|---|---|
hyp.scratch-low.yaml |
Default, low augmentation | lrf=0.01, scale=0.5, mixup=0.0 |
hyp.scratch-med.yaml |
Medium augmentation | scale=0.9, mixup=0.1 |
hyp.scratch-high.yaml |
High augmentation | scale=0.9, mixup=0.1, copy_paste=0.1 |
hyp.no-augmentation.yaml |
For Albumentations | All spatial/color augmentations set to 0 |
hyp.Objects365.yaml |
Evolved for Objects365 | lr0=0.00258, momentum=0.779 |
hyp.VOC.yaml |
Evolved for VOC | lr0=0.00334, momentum=0.74832 |
Environment Variables — Complete Reference
| Variable | Default | Scope | Description |
|---|---|---|---|
YOLOv5_DATASETS_DIR |
(none) | utils/general.py |
Override dataset root directory |
YOLOv5_AUTOINSTALL |
True |
utils/general.py |
Auto-install missing pip packages |
YOLOv5_VERBOSE |
True |
utils/general.py |
Enable verbose logging |
PIN_MEMORY |
True |
utils/dataloaders.py |
Pin dataloader memory |
LOCAL_RANK |
(auto) | train.py |
DDP local rank (set by torch.distributed) |
RANK |
(auto) | train.py |
DDP global rank |
WORLD_SIZE |
1 |
train.py |
DDP world size |
COMET_API_KEY |
(none) | Comet logger | Comet authentication |
COMET_PROJECT_NAME |
yolov5 |
Comet logger | Comet project name |
COMET_MODE |
online |
Comet logger | online or offline |
COMET_MODEL_NAME |
yolov5 |
Comet logger | Model name in Comet |
COMET_LOG_CONFUSION_MATRIX |
true |
Comet logger | Log confusion matrix |
COMET_MAX_IMAGE_UPLOADS |
100 |
Comet logger | Max prediction images logged |
COMET_LOG_PER_CLASS_METRICS |
false |
Comet logger | Per-class metrics |
COMET_DEFAULT_CHECKPOINT_FILENAME |
last.pt |
Comet logger | Resume checkpoint name |
COMET_LOG_BATCH_LEVEL_METRICS |
false |
Comet logger | Per-batch metrics |
COMET_LOG_PREDICTIONS |
true |
Comet logger | Log model predictions |
Note
DDP environment variables (
LOCAL_RANK,RANK,WORLD_SIZE) are managed bytorch.distributed.runand should not be set manually.
4. Deployment & Scaling
Docker Deployment
Three Dockerfiles are provided in utils/docker/:
| Dockerfile | Base | Target | Size |
|---|---|---|---|
Dockerfile |
nvidia/cuda |
GPU systems | Large (~6 GB) |
Dockerfile-cpu |
python:3.x-slim |
CPU-only systems | Small (~1 GB) |
Dockerfile-arm64 |
arm64v8/python |
ARM64 (Jetson, RPi) | Medium (~2 GB) |
Build and run:
# GPU
docker build -t yolov5 -f utils/docker/Dockerfile .
docker run --ipc=host --gpus all -v /data:/data yolov5 python detect.py --source /data/images
# CPU
docker build -t yolov5-cpu -f utils/docker/Dockerfile-cpu .
docker run --ipc=host -v /data:/data yolov5-cpu python detect.py --source /data/images --device cpu
Important
Always use
--ipc=hostor--shm-size=2gwhen running Docker containers with PyTorch dataloaders. Insufficient shared memory causesBus errorcrashes during training.
Google App Engine Deployment
A minimal GAE configuration is provided in utils/google_app_engine/:
app.yaml— GAE application configDockerfile— Container for GAEadditional_requirements.txt— Extra dependencies
AWS Deployment
AWS utilities in utils/aws/:
resume.py— Resume training on EC2 spot instances after interruptionuserdata.sh— EC2 instance bootstrap scriptmime.sh— MIME multipart user data assembly
Flask REST API Deployment
The built-in REST API (utils/flask_rest_api/restapi.py) serves models via HTTP:
python utils/flask_rest_api/restapi.py --port 5000 --model yolov5s yolov5m
API endpoint: POST /v1/object-detection/<model>
| Parameter | Type | Description |
|---|---|---|
image |
file (multipart) | Image to process |
<model> |
URL path | Model name (must match --model args) |
Response: JSON array of detections with xmin, ymin, xmax, ymax, confidence, class, name fields.
Warning
The Flask REST API is intended for development and testing. For production serving, use NVIDIA Triton Inference Server (supported via
utils/triton.pyandDetectMultiBackend) or export to TensorRT/ONNX and serve with a production-grade framework.
NVIDIA Triton Inference Server
YOLOv5 includes a Triton client (utils/triton.py) and DetectMultiBackend natively supports Triton URLs:
python detect.py --weights http://triton-server:8000/yolov5s --source image.jpg
The TritonRemoteModel class handles gRPC/HTTP communication with Triton.
Multi-GPU Scaling
Data Parallel (single-node):
python train.py --device 0,1,2,3 --batch-size 64
Distributed Data Parallel (recommended):
python -m torch.distributed.run --nproc_per_node 4 train.py \
--device 0,1,2,3 --batch-size 64 --sync-bn
Tip
DDP is significantly faster than DataParallel for multi-GPU training. Use
--sync-bnwith DDP for synchronized batch normalization across GPUs — particularly important when per-GPU batch size is small.
5. Operations
Training Output Structure
runs/train/exp/
├── weights/
│ ├── best.pt # Best mAP weights
│ └── last.pt # Latest checkpoint
├── confusion_matrix.png
├── F1_curve.png
├── P_curve.png
├── PR_curve.png
├── R_curve.png
├── results.png # Training metrics plot
├── results.csv # Epoch-level metrics
├── labels.jpg # Label distribution
├── labels_correlogram.jpg
├── train_batch0.jpg # Sample training batch
├── val_batch0_labels.jpg
├── val_batch0_pred.jpg
├── hyp.yaml # Used hyperparameters
├── opt.yaml # Command-line options
└── events.out.tfevents* # TensorBoard events
Checkpoint Format
Checkpoints (.pt files) contain:
{
'epoch': int, # Completed epochs
'best_fitness': float, # Best fitness score
'model': nn.Module, # Model state dict
'ema': nn.Module, # EMA state dict
'updates': int, # EMA updates count
'optimizer': dict, # Optimizer state
'opt': Namespace, # Training options
'git': str, # Git info
'date': str, # Timestamp
}
Monitoring with TensorBoard
tensorboard --logdir runs/train
Logged metrics include: train/box_loss, train/obj_loss, train/cls_loss, metrics/precision, metrics/recall, metrics/mAP_0.5, metrics/mAP_0.5:0.95, val/box_loss, val/obj_loss, val/cls_loss, learning rates.
Experiment Tracking Integrations
| Platform | Activation | What Gets Logged |
|---|---|---|
| TensorBoard | Auto (if installed) | Metrics, model graph, images |
| Weights & Biases | pip install wandb + wandb login |
Metrics, hyperparams, artifacts, images |
| ClearML | pip install clearml + clearml-init |
Metrics, code, packages, checkpoints, images |
| Comet | COMET_API_KEY env var |
Metrics, hyperparams, confusion matrix, predictions |
Note
Loggers are auto-detected. If a package is installed, its logger activates automatically. Use the
GenericLoggerclass inutils/loggers/__init__.pyfor console/CSV/TensorBoard fallback.
Hyperparameter Evolution
The built-in genetic algorithm evolves hyperparameters:
python train.py --weights yolov5s.pt --data coco128.yaml --evolve 300
Process:
- Initial training with current hyperparameters
- Mutation of hyperparameters using
np.random.randnwith configurable sigma - Fitness evaluation (0.1 * mAP@0.5 + 0.9 * mAP@0.5:0.95)
- Selection of best-performing configurations
- Results saved to
runs/train/evolve/hyp_evolve.yaml
6. Performance
AutoBatch
Automatic batch size detection based on GPU memory:
python train.py --batch-size -1
The autobatch() function in utils/autobatch.py profiles GPU memory usage with increasing batch sizes and selects the largest that fits within a configurable fraction (default: 0.60) of available VRAM.
Image Caching
# RAM caching (fastest, ~2x speedup)
python train.py --cache ram
# Disk caching (.npy files)
python train.py --cache disk
The check_cache_ram() method verifies sufficient system RAM before caching with a safety margin.
FP16 / Mixed Precision
Training uses PyTorch AMP (Automatic Mixed Precision) by default when a CUDA device is detected. This halves memory usage and accelerates training on Volta+ GPUs.
For inference, use --half:
python detect.py --weights yolov5s.pt --half --device 0
Model Optimization Post-Export
- ONNX simplification:
--simplifyreduces graph complexity - TensorRT: Automatically optimizes for target GPU at export time. Use
--workspaceto control memory allocation (default: 4 GB) - OpenVINO INT8:
--int8enables post-training quantization via calibration dataset - CoreML INT8:
--int8with--mlmodelfor mobile deployment - TFLite INT8:
--int8for edge deployment on Coral/mobile
Inference Speed Optimization
| Technique | Speedup | Trade-off |
|---|---|---|
Smaller model (yolov5n vs yolov5x) |
4-5x | Lower accuracy |
FP16 (--half) |
1.5-2x | Negligible accuracy loss |
| TensorRT export | 2-3x vs PyTorch | Requires NVIDIA GPU |
| OpenVINO export | 2-3x vs PyTorch (CPU) | Intel CPU only |
Reduce --imgsz |
~Linear | Lower small-object detection |
--vid-stride |
Linear | Skip frames in video |
7. Security
Input Validation
- Source paths validated against filesystem existence before processing (raises
FileNotFoundError) - Image format validation via extension whitelist (
IMG_FORMATS,VID_FORMATS) - URL scheme validation (HTTP/HTTPS/RTSP/RTMP)
- YAML files loaded with
yaml.safe_load()viautils/general.py:yaml_load()
Dependency Security
The requirements.txt includes version floors with security patches:
pillow>=10.3.0— Addresses image parsing vulnerabilitiesrequests>=2.32.2— HTTP security fixesurllib3>=2.6.0(Python >3.8) — Pinned per Snyk advisorysetuptools>=70.0.0— Snyk vulnerability fixtqdm>=4.66.3— Security patch
Warning
The codebase executes
os.system("pip install -U ultralytics")inmodels/common.pyif theultralyticspackage is not found. In locked-down environments, pre-install all dependencies to prevent runtime package installation.
Auto-Install Behavior
When YOLOv5_AUTOINSTALL=True (default), missing packages are automatically installed via pip. Disable this in production:
export YOLOv5_AUTOINSTALL=False
Model Weight Security
- Weights are downloaded from GitHub releases via HTTPS
torch.load()is used for model deserialization — only load weights from trusted sources- The
attempt_download()function inutils/downloads.pyfetches fromgithub.com/ultralytics/yolov5/releases
Caution
PyTorch model files (
.pt) can execute arbitrary code when loaded withtorch.load(). Never load model weights from untrusted sources. In production environments, consider exporting to ONNX or TorchScript, which have safer deserialization.
Network Exposure
- The Flask REST API binds to
0.0.0.0by default — restrict in production - No authentication is built into the REST API
- Triton client communicates via gRPC/HTTP without TLS by default
8. Developer Internals
Repository Structure
yolov5/
├── benchmarks.py # Export format benchmarking
├── detect.py # Detection inference entry point
├── export.py # Model export (11+ formats)
├── hubconf.py # PyTorch Hub entry points
├── train.py # Detection training entry point
├── val.py # Detection validation entry point
├── classify/
│ ├── predict.py # Classification inference
│ ├── train.py # Classification training
│ ├── val.py # Classification validation
│ └── tutorial.ipynb # Jupyter tutorial
├── segment/
│ ├── predict.py # Segmentation inference
│ ├── train.py # Segmentation training
│ ├── val.py # Segmentation validation
│ └── tutorial.ipynb # Jupyter tutorial
├── models/
│ ├── __init__.py
│ ├── common.py # Core layers + DetectMultiBackend
│ ├── experimental.py # Ensemble, model loading
│ ├── tf.py # TensorFlow/Keras equivalents
│ ├── yolo.py # Model classes + parse_model
│ ├── yolov5[n|s|m|l|x].yaml # Architecture configs
│ ├── hub/ # Experimental variants
│ └── segment/ # Segmentation model configs
├── utils/
│ ├── __init__.py # TryExcept, threading, notebook
│ ├── activations.py # Custom activation functions
│ ├── augmentations.py # Data augmentation pipeline
│ ├── autoanchor.py # Automatic anchor calculation
│ ├── autobatch.py # Automatic batch sizing
│ ├── callbacks.py # Event hook system
│ ├── dataloaders.py # All data loading classes
│ ├── downloads.py # Weight/asset downloading
│ ├── general.py # Core utilities (NMS, box ops, etc.)
│ ├── loss.py # Loss functions
│ ├── metrics.py # Evaluation metrics
│ ├── plots.py # Visualization
│ ├── torch_utils.py # PyTorch helpers
│ ├── triton.py # NVIDIA Triton client
│ ├── loggers/ # Experiment tracking
│ ├── segment/ # Segmentation utilities
│ ├── flask_rest_api/ # REST API serving
│ ├── docker/ # Container configs
│ ├── aws/ # AWS deployment helpers
│ └── google_app_engine/ # GAE config
├── data/
│ ├── *.yaml # Dataset configurations
│ ├── hyps/ # Hyperparameter presets
│ ├── images/ # Sample test images
│ └── scripts/ # Dataset download scripts
├── pyproject.toml # Project metadata
├── requirements.txt # Dependencies
├── CITATION.cff # Citation info
├── CONTRIBUTING.md # Contributing guidelines
├── LICENSE # AGPL-3.0
└── tutorial.ipynb # Getting started notebook
Architecture Deep Dive
DetectMultiBackend (models/common.py)
The central inference class that transparently handles 13 model formats. At initialization, it:
- Determines format from file extension/path pattern via
_model_type() - Loads the appropriate runtime (PyTorch, ONNX Runtime, TensorRT, OpenVINO, CoreML, TensorFlow, PaddlePaddle, Triton)
- Resolves class names, stride, and input format (NCHW vs NHWC)
- Sets FP16 flag based on format support
The forward() method normalizes input tensors and routes to the active backend.
Model Construction (models/yolo.py)
DetectionModel.__init__() uses parse_model() to build the network from a YAML config:
- Reads backbone and head layer definitions
- Resolves module classes (
Conv,C3,SPPF,Detect, etc.) - Applies
depth_multipleto repeat counts andwidth_multipleto channel counts - Builds the
nn.Sequentialmodel with proper skip connections - Initializes anchors and stride for the Detect head
SegmentationModel extends DetectionModel by replacing the Detect head with Segment, which adds prototype prediction.
ClassificationModel extends BaseModel with a classification head instead of detection anchors.
Loss Computation (utils/loss.py)
ComputeLoss.__init__() configures:
BCEWithLogitsLossfor objectness and classification (optionally wrapped inFocalLoss)- Label smoothing support
- AutoBalance for multi-scale loss weighting
ComputeLoss.__call__():
build_targets()matches ground truth to anchors across all detection layers- Computes box regression loss (CIoU)
- Computes objectness loss (BCE)
- Computes classification loss (BCE)
- Returns weighted sum scaled by batch size
For segmentation, utils/segment/loss.py:ComputeLoss adds mask loss via prototype-based instance segmentation.
Data Loading (utils/dataloaders.py)
Key classes:
LoadImagesAndLabels(Dataset)— Main dataset class for training/validation. Handles mosaic augmentation, image caching, label verification, and rectangular batching.InfiniteDataLoader— WrapsDataLoaderwith_RepeatSamplerfor endless iteration.LoadImages— Simple image/video file loader for inference.LoadStreams— Multi-threaded stream loader for webcam/RTSP/RTMP.LoadScreenshots— Screen capture loader usingmss.ClassificationDataset— Extendstorchvision.datasets.ImageFolder.
The create_dataloader() factory function configures all loader parameters including distributed sampling.
Augmentation Pipeline (utils/augmentations.py)
Training augmentations applied in order:
- Mosaic (4 or 9 images composited) —
load_mosaic()/load_mosaic9() - Copy-paste — pastes instances from other images
- Random perspective — rotation, translation, scale, shear, perspective
- MixUp — weighted average of two images
- HSV augmentation — hue, saturation, value jitter
- Flip — horizontal and vertical
- Albumentations (optional) — third-party augmentation pipeline
Optional Albumentations integration wraps the albumentations library for Blur, MedianBlur, ToGray, CLAHE, RandomBrightnessContrast, RandomGamma, and ImageCompression transforms.
Callback System (utils/callbacks.py)
The Callbacks class provides a hook-based event system:
callbacks = Callbacks()
callbacks.register_action('on_train_start', name='my_hook', callback=my_fn)
callbacks.run('on_train_start')
Registered hooks:
on_pretrain_routine_start/endon_train_start/epoch_end/batch_end/endon_val_start/image_end/batch_end/endon_fit_epoch_endon_model_saveon_params_update
EarlyStopping (utils/torch_utils.py)
The EarlyStopping class monitors fitness (mAP combination):
stopper = EarlyStopping(patience=100)
if stopper(epoch, fitness):
break # Stop training
Fitness is computed as 0.1 * mAP@0.5 + 0.9 * mAP@0.5:0.95 by default.
Exponential Moving Average (utils/torch_utils.py)
ModelEMA maintains a shadow copy of model weights:
- Decay starts near 0 and increases to target (default: 0.9999)
- Updated after each training step
- Used for validation and checkpoint saving
update_attr()copies non-weight attributes (names, stride, etc.)
Module Reference
models/common.py — Neural Network Layers
| Class | Purpose | Key Detail |
|---|---|---|
Conv |
Standard Conv+BN+SiLU | forward_fuse() for inference (merged BN) |
DWConv |
Depthwise convolution | Groups = GCD(c1, c2) |
DWConvTranspose2d |
Depthwise transposed conv | For upsampling |
TransformerLayer |
Multi-head attention | No LayerNorm (performance) |
TransformerBlock |
Conv + Transformer layers | Used in C3TR variant |
Bottleneck |
Standard bottleneck | Optional shortcut |
BottleneckCSP |
CSP bottleneck | Cross-stage partial connections |
CrossConv |
Cross convolution | (1,k)→(k,1) factored convs |
C3 |
CSP with 3 convolutions | Primary building block |
C3x |
C3 with CrossConv | Extended C3 |
C3TR |
C3 with TransformerBlock | Attention-enhanced C3 |
C3SPP |
C3 with SPP | Spatial pyramid pooling |
C3Ghost |
C3 with GhostBottleneck | Lightweight variant |
SPP |
Spatial Pyramid Pooling | Multi-kernel max pooling |
SPPF |
Fast SPP | Sequential max pooling (equiv to SPP) |
Focus |
Space-to-depth | 4x channel expansion |
GhostConv |
Ghost convolution | Half channels + depthwise |
GhostBottleneck |
Ghost bottleneck | Efficient feature extraction |
Contract |
Spatial→Channel | Downscale spatial dims |
Expand |
Channel→Spatial | Upscale spatial dims |
Concat |
Tensor concatenation | Configurable dimension |
DetectMultiBackend |
Multi-format inference | 13 backend support |
AutoShape |
Automatic preprocessing | Handles diverse inputs |
Detections |
Detection results | .show(), .save(), .pandas() |
Proto |
Prototype generation | For segmentation masks |
Classify |
Classification head | Pooling + dropout + linear |
models/yolo.py — Model Classes
| Class | Purpose |
|---|---|
Detect(nn.Module) |
Detection head — generates anchor-based predictions |
Segment(Detect) |
Segmentation head — adds prototype coefficients |
BaseModel(nn.Module) |
Abstract base with fuse() and info() |
DetectionModel(BaseModel) |
Full detection model from YAML |
SegmentationModel(DetectionModel) |
Detection model with Segment head |
ClassificationModel(BaseModel) |
Classification model from detection backbone |
utils/general.py — Core Utilities
Key functions:
| Function | Purpose |
|---|---|
non_max_suppression() |
NMS with class-agnostic option, multi-label support |
scale_boxes() |
Rescale boxes from inference to original image size |
xyxy2xywh() / xywh2xyxy() |
Box format conversion |
check_dataset() |
Validate and auto-download datasets |
check_img_size() |
Round image size to stride multiple |
check_version() |
Version string comparison |
increment_path() |
Auto-increment experiment directory |
strip_optimizer() |
Remove optimizer from checkpoint (shrink file) |
colorstr() |
ANSI-colored terminal strings |
yaml_load() / yaml_save() |
Safe YAML I/O |
download() |
Multi-threaded file downloads |
utils/metrics.py — Evaluation
| Class/Function | Purpose |
|---|---|
ConfusionMatrix |
Builds and plots confusion matrix |
ap_per_class() |
Compute AP for each class |
compute_ap() |
Average precision from P-R curve |
bbox_iou() |
Box IoU (supports GIoU/DIoU/CIoU) |
fitness() |
Weighted mAP combination (0.1mAP50 + 0.9mAP50-95) |
plot_pr_curve() |
Precision-recall curve plot |
plot_mc_curve() |
Metric-confidence curve plot |
utils/torch_utils.py — PyTorch Helpers
| Function/Class | Purpose |
|---|---|
select_device() |
Auto-select CUDA/CPU device |
smart_inference_mode() |
torch.no_grad() / torch.inference_mode() decorator |
smart_optimizer() |
Create optimizer with proper weight decay groups |
smart_resume() |
Resume from checkpoint with validation |
fuse_conv_and_bn() |
Merge Conv2d + BatchNorm2d for inference |
model_info() |
Print model summary (layers, params, gradients) |
EarlyStopping |
Stop training when fitness plateaus |
ModelEMA |
Exponential moving average of model weights |
profile() |
Benchmark operations (FLOPs, time) |
Code Patterns and Conventions
Path handling: All scripts use Path(__file__).resolve() to establish ROOT. Paths are resolved relative to the repo root.
Device management: select_device() centralizes CUDA/CPU selection. Returns a torch.device and validates GPU availability.
Experiment naming: increment_path() auto-increments experiment directories: exp, exp2, exp3, etc. Use --exist-ok to reuse.
Logging: The LOGGER object (from utils/general.py) is a standard Python logger configured in set_logging(). All modules use it.
Model fusing: model.fuse() merges Conv2d+BatchNorm2d layers for faster inference. Called automatically before export and inference.
NMS implementation: non_max_suppression() in utils/general.py handles multi-class and class-agnostic NMS, with optional multi-label output and mask coefficient support for segmentation.
Testing
The repository includes a tutorial.ipynb notebook (and classify/tutorial.ipynb, segment/tutorial.ipynb) for interactive testing. Automated tests are primarily through:
benchmarks.py --test— Validates all export formatsbenchmarks.py— Full inference benchmarks across formats- GitHub Actions CI (referenced in the repository badges)
Tip
Use
python benchmarks.py --test --pt-onlyfor a quick smoke test that only validates the PyTorch path.
Potential Improvements
Based on the codebase analysis:
- MPS backend support: Incomplete —
aten::nonzeroand other ops not supported (#8508) - Parallel inference with DDP: Not supported (#9687) — DDP is training-only
- ConvTranspose2d upsampling: Feature request to replace nearest-neighbor upsample (#7520)
- Quantization-aware training: Not built-in — export-time quantization only
- Model pruning:
utils/torch_utils.py:prune()exists but is not integrated into the training pipeline - Attention mechanisms: C3TR variant exists but is not default — could improve small-object detection
Contributing
The project uses the AGPL-3.0 license. Contributions follow standard GitHub flow:
- Fork the repository
- Create a feature branch
- Make changes following existing code patterns
- Test with
benchmarks.py - Submit a pull request
Code style conventions:
- Type hints used in function signatures (not exhaustive)
- Docstrings on all public functions (Google style)
LOGGERfor all output (no bareprint()in library code)- YAML for all configuration
- Argparse for all CLI interfaces
pathlib.Pathpreferred overos.path
9. Integration Points
PyTorch Hub API
hubconf.py exposes the following entry points for torch.hub.load():
| Function | Description |
|---|---|
custom(path) |
Load a custom model from local path |
yolov5n(pretrained) |
YOLOv5 Nano |
yolov5s(pretrained) |
YOLOv5 Small |
yolov5m(pretrained) |
YOLOv5 Medium |
yolov5l(pretrained) |
YOLOv5 Large |
yolov5x(pretrained) |
YOLOv5 Extra-Large |
yolov5n6(pretrained) |
YOLOv5 Nano P6 |
yolov5s6(pretrained) |
YOLOv5 Small P6 |
yolov5m6(pretrained) |
YOLOv5 Medium P6 |
yolov5l6(pretrained) |
YOLOv5 Large P6 |
yolov5x6(pretrained) |
YOLOv5 Extra-Large P6 |
All Hub functions accept channels, classes, autoshape, _verbose, and device parameters.
The AutoShape wrapper enables inference on diverse inputs (PIL, numpy, torch, filename, URL) with automatic preprocessing, NMS, and result formatting.
Logger Integration Architecture
utils/loggers/__init__.py defines the Loggers class that aggregates all active loggers:
Loggers
├── TensorBoard (via SummaryWriter)
├── Weights & Biases (if wandb installed)
├── ClearML (if clearml installed)
└── Comet (if comet_ml installed + COMET_API_KEY)
Each logger receives callbacks at training lifecycle events. The GenericLogger class provides a simpler interface for standalone use.
ClearML Integration
Files: utils/loggers/clearml/
clearml_utils.py— Task creation, dataset handlinghpo.py— Hyperparameter optimization via ClearML + Optuna- Auto-logs: source code, packages, hyperparameters, checkpoints, metrics, images
Comet Integration
Files: utils/loggers/comet/
comet_utils.py— Experiment logging, artifact managementhpo.py— Hyperparameter sweeps via Comet Optimizeroptimizer_config.json— Default HPO sweep configuration- Supports dataset versioning via Comet Artifacts
- Resume training via
comet://WORKSPACE/PROJECT/EXPERIMENT_ID
W&B Integration
Files: utils/loggers/wandb/
wandb_utils.py— Experiment tracking, artifact logging
Dataset Download Scripts
Shell scripts in data/scripts/:
| Script | Purpose |
|---|---|
get_coco.sh |
Download full COCO 2017 dataset |
get_coco128.sh |
Download COCO128 mini dataset |
get_imagenet.sh |
Download ImageNet (requires manual auth) |
get_imagenet10.sh |
Download ImageNet-10 subset |
get_imagenet100.sh |
Download ImageNet-100 subset |
get_imagenet1000.sh |
Download ImageNet-1000 |
download_weights.sh |
Download all pretrained weights |