New YOLOv5 Classification Models

[glenn-jocher]

We trained YOLOv5-cls classification models on ImageNet for 90 epochs using a 4xA100 instance, and we trained ResNet and EfficientNet models alongside with the same default training settings to compare. We exported all models to ONNX FP32 for CPU speed tests and to TensorRT FP16 for GPU speed tests. We ran all speed tests on Google Colab Pro for easy reproducibility.

New Classification Checkpoints

Model	size (pixels)	acc top1	acc top5	Training 90 epochs 4xA100 (hours)	Speed ONNX CPU (ms)	Speed TensorRT V100 (ms)	params (M)	FLOPs @224 (B)
YOLOv5n-cls	224	64.6	85.4	7:59	3.3	0.5	2.5	0.5
YOLOv5s-cls	224	71.5	90.2	8:09	6.6	0.6	5.4	1.4
YOLOv5m-cls	224	75.9	92.9	10:06	15.5	0.9	12.9	3.9
YOLOv5l-cls	224	78.0	94.0	11:56	26.9	1.4	26.5	8.5
YOLOv5x-cls	224	79.0	94.4	15:04	54.3	1.8	48.1	15.9
ResNet18	224	70.3	89.5	6:47	11.2	0.5	11.7	3.7
ResNet34	224	73.9	91.8	8:33	20.6	0.9	21.8	7.4
ResNet50	224	76.8	93.4	11:10	23.4	1.0	25.6	8.5
ResNet101	224	78.5	94.3	17:10	42.1	1.9	44.5	15.9
EfficientNet_b0	224	75.1	92.4	13:03	12.5	1.3	5.3	1.0
EfficientNet_b1	224	76.4	93.2	17:04	14.9	1.6	7.8	1.5
EfficientNet_b2	224	76.6	93.4	17:10	15.9	1.6	9.1	1.7
EfficientNet_b3	224	77.7	94.0	19:19	18.9	1.9	12.2	2.4

• All checkpoints are trained to 90 epochs with SGD optimizer with lr0=0.001 at image size 224 and all default settings. Runs logged to https://wandb.ai/glenn-jocher/YOLOv5-Classifier-v6-2.
• Accuracy values are for single-model single-scale on ImageNet-1k dataset.
  Reproduce by python classify/val.py --data ../datasets/imagenet --img 224
• Speed averaged over 100 inference images using a Google Colab Pro V100 High-RAM instance.
  Reproduce by python classify/val.py --data ../datasets/imagenet --img 224 --batch 1
• Export to ONNX at FP32 and TensorRT at FP16 done with export.py.
  Reproduce by python export.py --weights yolov5s-cls.pt --include engine onnx --imgsz 224

New Classification Usage Examples

Train

YOLOv5 classification training supports auto-download of MNIST, Fashion-MNIST, CIFAR10, CIFAR100, Imagenette, Imagewoof, and ImageNet datasets with the --data argument. To start training on MNIST for example use --data mnist.

# Single-GPU
python classify/train.py --model yolov5s-cls.pt --data cifar100 --epochs 5 --img 224 --batch 128

# Multi-GPU DDP
python -m torch.distributed.run --nproc_per_node 4 --master_port 1 classify/train.py --model yolov5s-cls.pt --data imagenet --epochs 5 --img 224 --device 0,1,2,3

Val

Validate accuracy on a pretrained model. To validate YOLOv5s-cls accuracy on ImageNet.

bash data/scripts/get_imagenet.sh --val  # download ImageNet val split (6.3G, 50000 images)
python classify/val.py --weights yolov5s-cls.pt --data ../datasets/imagenet --img 224

Predict

Run a classification prediction on an image.

python classify/predict.py --weights yolov5s-cls.pt --data data/images/bus.jpg

model = torch.hub.load('ultralytics/yolov5', 'custom', 'yolov5s-cls.pt')  # load from PyTorch Hub

Export

Export a group of trained YOLOv5-cls, ResNet and EfficientNet models to ONNX and TensorRT.

python export.py --weights yolov5s-cls.pt resnet50.pt efficientnet_b0.pt --include onnx engine --img 224

🛠️ PR Summary

_{Made with ❤️ by Ultralytics Actions}

🌟 Summary

Enhance YOLOv5 classification models, introduce training features, and improve GitHub workflows.

📊 Key Changes

• Provided support for training classification models
• Updated GitHub action workflows for better CI testing
• Added new classification usage examples
• Removed outdated image competition banner
• Implemented image classification checkpoint comparisons
• Incorporated additional logging capabilities for classification

🎯 Purpose & Impact

• Purpose:

  • Create a simplified pipeline for training image classification models, enabling better performance and ease of use.
  • Streamline CI workflows to ensure codebase stability and reliability.
  • Update documentation to give users clear examples of how to work with image classification features.

• Impact to Users:

  • Users can now train classification models using YOLOv5 with simple commands, making it more accessible to a wider audience.
  • Improved testing ensures that users receive more reliable updates and enhancements.
  • New usage examples and documentation offer users better guidance on leveraging new features, such as checkpoints for model speed tests and comparisons.

[zhiqwang]

Hi @glenn-jocher , First of all, congratulations to YOLOv5 for finally having classification models, do you plan to give a detection model pre-trained on this classification model? I guess this will also help to promote the application of this classification models.

[glenn-jocher]

@zhiqwang good question. I haven't tried to build detection models from the classification models, I actually went to the other way around and build the classification models by just attaching a Classify() module to the detection backbones.

But yes we should be able to populate just the detection backbones from the ImageNet pretrained cls models now...

[glenn-jocher]

@zhiqwang oh wow it actually works out of the box. Amazing. This is the command to transfer ImageNet backbone for detection training (head remains randomly initialised).

python train.py --cfg yolov5s.yaml --weights yolov5s-cls.pt
...
# Transferred 186/349 items from yolov5s-cls.pt

@AyushExel

[zhiqwang]

oh wow it actually works out of the box. Amazing. This is the command to transfer ImageNet backbone for detection training (head remains randomly initialised).

Hi @glenn-jocher , The flexibility of the YOLOv5 framework guarantees ease of use here!

Another question is that I see that the classification model uses a different mean and std than the detection model, not sure if this would cause the pre-training mechanism to be disabled here?

[glenn-jocher]

@zhiqwang yes the cls models use ImageNet normalization whereas the detection models do not normalize, that is a good point.

[Azzam-Alhussain]

Could you please help me to know I can train this classifier from scratch rather than using knowledge transfer from pretrained?