From e61756910758f59406255269921e55992ca0b64b Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Wed, 17 Aug 2022 15:33:37 +0200 Subject: [PATCH] Improve classification comments (#8997) --- .github/README_cn.md | 10 +++++----- README.md | 10 +++++----- classify/predict.py | 2 +- classify/train.py | 4 +++- classify/val.py | 3 ++- 5 files changed, 16 insertions(+), 13 deletions(-) diff --git a/.github/README_cn.md b/.github/README_cn.md index 86b502df61f7..816adf6b0449 100644 --- a/.github/README_cn.md +++ b/.github/README_cn.md @@ -269,7 +269,7 @@ We trained YOLOv5-cls classification models on ImageNet for 90 epochs using a 4x
Table Notes (click to expand) -- 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. +- All checkpoints are trained to 90 epochs with SGD optimizer with `lr0=0.001` and `weight_decay=5e-5` 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](https://www.image-net.org/index.php) dataset.
Reproduce by `python classify/val.py --data ../datasets/imagenet --img 224` - **Speed** averaged over 100 inference images using a Google [Colab Pro](https://colab.research.google.com/signup) 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` @@ -291,14 +291,14 @@ python -m torch.distributed.run --nproc_per_node 4 --master_port 1 classify/trai ``` ### Val -Validate accuracy on a pretrained model. To validate YOLOv5s-cls accuracy on ImageNet. +Validate YOLOv5m-cls accuracy on ImageNet-1k dataset: ```bash 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 +python classify/val.py --weights yolov5m-cls.pt --data ../datasets/imagenet --img 224 # validate ``` ### Predict -Run a classification prediction on an image. +Use pretrained YOLOv5s-cls.pt to predict bus.jpg: ```bash python classify/predict.py --weights yolov5s-cls.pt --data data/images/bus.jpg ``` @@ -307,7 +307,7 @@ model = torch.hub.load('ultralytics/yolov5', 'custom', 'yolov5s-cls.pt') # load ``` ### Export -Export a group of trained YOLOv5-cls, ResNet and EfficientNet models to ONNX and TensorRT. +Export a group of trained YOLOv5s-cls, ResNet and EfficientNet models to ONNX and TensorRT: ```bash python export.py --weights yolov5s-cls.pt resnet50.pt efficientnet_b0.pt --include onnx engine --img 224 ``` diff --git a/README.md b/README.md index b368d1d6e264..7335394402da 100644 --- a/README.md +++ b/README.md @@ -278,7 +278,7 @@ We trained YOLOv5-cls classification models on ImageNet for 90 epochs using a 4x
Table Notes (click to expand) -- 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. +- All checkpoints are trained to 90 epochs with SGD optimizer with `lr0=0.001` and `weight_decay=5e-5` 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](https://www.image-net.org/index.php) dataset.
Reproduce by `python classify/val.py --data ../datasets/imagenet --img 224` - **Speed** averaged over 100 inference images using a Google [Colab Pro](https://colab.research.google.com/signup) 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` @@ -300,14 +300,14 @@ python -m torch.distributed.run --nproc_per_node 4 --master_port 1 classify/trai ``` ### Val -Validate accuracy on a pretrained model. To validate YOLOv5s-cls accuracy on ImageNet. +Validate YOLOv5m-cls accuracy on ImageNet-1k dataset: ```bash 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 +python classify/val.py --weights yolov5m-cls.pt --data ../datasets/imagenet --img 224 # validate ``` ### Predict -Run a classification prediction on an image. +Use pretrained YOLOv5s-cls.pt to predict bus.jpg: ```bash python classify/predict.py --weights yolov5s-cls.pt --data data/images/bus.jpg ``` @@ -316,7 +316,7 @@ model = torch.hub.load('ultralytics/yolov5', 'custom', 'yolov5s-cls.pt') # load ``` ### Export -Export a group of trained YOLOv5-cls, ResNet and EfficientNet models to ONNX and TensorRT. +Export a group of trained YOLOv5s-cls, ResNet and EfficientNet models to ONNX and TensorRT: ```bash python export.py --weights yolov5s-cls.pt resnet50.pt efficientnet_b0.pt --include onnx engine --img 224 ``` diff --git a/classify/predict.py b/classify/predict.py index 419830d43952..4247e3c8e7fa 100644 --- a/classify/predict.py +++ b/classify/predict.py @@ -3,7 +3,7 @@ Run classification inference on images Usage: - $ python classify/predict.py --weights yolov5s-cls.pt --source im.jpg + $ python classify/predict.py --weights yolov5s-cls.pt --source data/images/bus.jpg """ import argparse diff --git a/classify/train.py b/classify/train.py index f2b465567446..b85f14236039 100644 --- a/classify/train.py +++ b/classify/train.py @@ -2,8 +2,10 @@ """ Train a YOLOv5 classifier model on a classification dataset Datasets: --data mnist, fashion-mnist, cifar10, cifar100, imagenette, imagewoof, imagenet, or 'path/to/custom/dataset' +YOLOv5-cls models: --model yolov5n-cls.pt, yolov5s-cls.pt, yolov5m-cls.pt, yolov5l-cls.pt, yolov5x-cls.pt +Torchvision models: --model resnet50, efficientnet_b0, etc. See https://pytorch.org/vision/stable/models.html -Usage: +Usage - Single-GPU and Multi-GPU DDP $ python classify/train.py --model yolov5s-cls.pt --data cifar100 --epochs 5 --img 128 $ 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 """ diff --git a/classify/val.py b/classify/val.py index 0930ba8c9c51..9d965d9f1fdc 100644 --- a/classify/val.py +++ b/classify/val.py @@ -3,7 +3,8 @@ Validate a classification model on a dataset Usage: - $ python classify/val.py --weights yolov5s-cls.pt --data ../datasets/imagenet + $ bash data/scripts/get_imagenet.sh --val # download ImageNet val split (6.3G, 50000 images) + $ python classify/val.py --weights yolov5m-cls.pt --data ../datasets/imagenet --img 224 # validate """ import argparse