Skip to content

Commit

Permalink
Improve classification comments (ultralytics#8997)
Browse files Browse the repository at this point in the history
  • Loading branch information
glenn-jocher authored and Clay Januhowski committed Sep 8, 2022
1 parent b6f8a2d commit 734b699
Show file tree
Hide file tree
Showing 5 changed files with 16 additions and 13 deletions.
10 changes: 5 additions & 5 deletions .github/README_cn.md
Original file line number Diff line number Diff line change
Expand Up @@ -269,7 +269,7 @@ We trained YOLOv5-cls classification models on ImageNet for 90 epochs using a 4x
<details>
<summary>Table Notes (click to expand)</summary>

- 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.<br>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.<br>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.<br>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`. <br>Reproduce by `python export.py --weights yolov5s-cls.pt --include engine onnx --imgsz 224`
Expand All @@ -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
```
Expand All @@ -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
```
Expand Down
10 changes: 5 additions & 5 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -278,7 +278,7 @@ We trained YOLOv5-cls classification models on ImageNet for 90 epochs using a 4x
<details>
<summary>Table Notes (click to expand)</summary>

- 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.<br>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.<br>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.<br>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`. <br>Reproduce by `python export.py --weights yolov5s-cls.pt --include engine onnx --imgsz 224`
Expand All @@ -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
```
Expand All @@ -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
```
Expand Down
2 changes: 1 addition & 1 deletion classify/predict.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
4 changes: 3 additions & 1 deletion classify/train.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
"""
Expand Down
3 changes: 2 additions & 1 deletion classify/val.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down

0 comments on commit 734b699

Please sign in to comment.