使用 Ultralytics YOLO 进行图像分类#
图像分类 是受支持任务中最简单的一种,它将整张图像分类为预定义类别中的某一类。
图像分类器的输出是单个类别标签和置信度分数。当你只需要知道图像属于哪个类别,而不需要知道该类别对象位于何处或其确切形状时,图像分类就非常有用。
Watch: Explore Ultralytics YOLO Tasks: Image Classification using Ultralytics Platform
YOLO26 Classify 模型使用 -cls 后缀,即 yolo26n-cls.pt,并在 ImageNet 上进行预训练。
模型#
此处展示了经过预训练的 YOLO26 Classify 模型。Detect、Segment 和 Pose 模型在 COCO 数据集上进行预训练,Semantic 模型在 Cityscapes 上进行预训练,而 Classify 模型在 ImageNet 数据集上进行预训练。
模型首次使用时会自动从最新的 Ultralytics 发行版下载。
| 模型 | 尺寸 (像素) | acc top1 | acc top5 | 速度 CPU ONNX (毫秒) | 速度 T4 TensorRT10 (毫秒) | 参数量 (M) | FLOPs (B),输入尺寸 224 |
|---|---|---|---|---|---|---|---|
| YOLO26n-cls | 224 | 71.4 | 90.1 | 5.0 ± 0.3 | 1.1 ± 0.0 | 2.8 | 0.4 |
| YOLO26s-cls | 224 | 76.0 | 92.9 | 7.9 ± 0.2 | 1.3 ± 0.0 | 6.7 | 1.5 |
| YOLO26m-cls | 224 | 78.1 | 94.2 | 17.2 ± 0.4 | 2.0 ± 0.0 | 11.6 | 4.8 |
| YOLO26l-cls | 224 | 79.0 | 94.6 | 23.2 ± 0.3 | 2.8 ± 0.0 | 14.1 | 6.0 |
| YOLO26x-cls | 224 | 79.9 | 95.0 | 41.4 ± 0.9 | 3.8 ± 0.0 | 29.6 | 13.5 |
- acc 值表示模型在 ImageNet 数据集验证集上的准确率。
通过yolo val classify data=path/to/ImageNet device=0复现 - Speed 是使用 Amazon EC2 P4d 实例在 ImageNet 验证图像上测得的平均速度。
通过yolo val classify data=path/to/ImageNet batch=1 device=0|cpu复现 - Params 和 FLOPs 值对应于执行
model.fuse()后的融合模型,该操作会合并 Conv 和 BatchNorm 层。预训练检查点保留完整的训练架构,因此显示的数量可能更高。
请参阅未发布的 YOLO27 预览版以了解初步的分类速度和模型大小。
训练#
在 MNIST160 数据集上以图像尺寸 64 训练 YOLO26n-cls,训练 100 个 epoch。有关可用参数的完整列表,请参阅 Configuration 页面。
from ultralytics import YOLO
# Load a model
model = YOLO("yolo26n-cls.yaml") # build a new model from YAML
model = YOLO("yolo26n-cls.pt") # load a pretrained model (recommended for training)
model = YOLO("yolo26n-cls.yaml").load("yolo26n-cls.pt") # build from YAML and transfer weights
# Train the model
results = model.train(data="mnist160", epochs=100, imgsz=64)Ultralytics YOLO 分类使用 torchvision.transforms.RandomResizedCrop 进行训练,并使用 torchvision.transforms.CenterCrop 进行验证和推理。
这些基于裁剪的变换假设输入为正方形,对于宽高比极端的图像,可能会意外裁掉重要区域,从而导致训练期间丢失关键信息。
若要在保持图像比例的同时保留完整图像,可以考虑使用 torchvision.transforms.Resize 代替裁剪变换。
你可以通过自定义 ClassificationDataset 和 ClassificationTrainer 来定制数据增强流程,从而实现这一点。
import torch
import torchvision.transforms as T
from ultralytics import YOLO
from ultralytics.data.dataset import ClassificationDataset
from ultralytics.models.yolo.classify import ClassificationTrainer, ClassificationValidator
class CustomizedDataset(ClassificationDataset):
"""A customized dataset class for image classification with enhanced data augmentation transforms."""
def __init__(self, root: str, args, augment: bool = False, prefix: str = ""):
"""Initialize a customized classification dataset with enhanced data augmentation transforms."""
super().__init__(root, args, augment, prefix)
# Add your custom training transforms here
train_transforms = T.Compose(
[
T.Resize((args.imgsz, args.imgsz)),
T.RandomHorizontalFlip(p=args.fliplr),
T.RandomVerticalFlip(p=args.flipud),
T.RandAugment(interpolation=T.InterpolationMode.BILINEAR),
T.ColorJitter(brightness=args.hsv_v, contrast=args.hsv_v, saturation=args.hsv_s, hue=args.hsv_h),
T.ToTensor(),
T.Normalize(mean=torch.tensor(0), std=torch.tensor(1)),
T.RandomErasing(p=args.erasing, inplace=True),
]
)
# Add your custom validation transforms here
val_transforms = T.Compose(
[
T.Resize((args.imgsz, args.imgsz)),
T.ToTensor(),
T.Normalize(mean=torch.tensor(0), std=torch.tensor(1)),
]
)
self.torch_transforms = train_transforms if augment else val_transforms
class CustomizedTrainer(ClassificationTrainer):
"""A customized trainer class for YOLO classification models with enhanced dataset handling."""
def build_dataset(self, img_path: str, mode: str = "train", batch=None):
"""Build a customized dataset for classification training and the validation during training."""
return CustomizedDataset(root=img_path, args=self.args, augment=mode == "train", prefix=mode)
class CustomizedValidator(ClassificationValidator):
"""A customized validator class for YOLO classification models with enhanced dataset handling."""
def build_dataset(self, img_path: str):
"""Build a customized dataset for classification standalone validation (no augmentation)."""
return CustomizedDataset(root=img_path, args=self.args, augment=False, prefix=self.args.split)
model = YOLO("yolo26n-cls.pt")
model.train(data="imagenet", trainer=CustomizedTrainer, epochs=10, imgsz=224, batch=64)
model.val(data="imagenet", validator=CustomizedValidator, imgsz=224, batch=64)数据集格式#
Dataset Guide 中详细介绍了 YOLO 分类数据集格式。你还可以使用 Ultralytics Platform annotation tools 管理和标注分类数据集。
验证#
在 MNIST160 数据集上验证经过训练的 YOLO26n-cls 模型的准确率。无需提供参数,因为 model 会将其训练用的 data 和参数保留为模型属性。
from ultralytics import YOLO
# Load a model
model = YOLO("yolo26n-cls.pt") # load an official model
model = YOLO("path/to/best.pt") # load a custom model
# Validate the model
metrics = model.val() # no arguments needed, dataset and settings remembered
metrics.top1 # top1 accuracy
metrics.top5 # top5 accuracyPredict#
使用经过训练的 YOLO26n-cls 模型对图像运行预测。
from ultralytics import YOLO
# Load a model
model = YOLO("yolo26n-cls.pt") # load an official model
model = YOLO("path/to/best.pt") # load a custom model
# Predict with the model
results = model("https://ultralytics.com/images/bus.jpg") # predict on an image
# Access the results
for result in results:
top1 = result.probs.top1 # top predicted class ID
top1_conf = result.probs.top1conf # top prediction confidence
top1_name = result.names[top1] # top predicted class name请在预测页面查看完整的 predict 模式详情。
结果输出#
图像分类会为每张图像返回一个 Results 对象。主要预测字段是 result.probs,其中
包含类别概率向量以及用于获取最高预测结果的辅助工具。
| 属性 | 类型 | 形状 | 描述 |
|---|---|---|---|
result.probs | Probs | (C,) | 类别概率。 |
result.probs.data | torch.float32 | (C,) | 每个类别的概率。 |
result.probs.top1 | int | () | 最高类别 ID。 |
result.probs.top1conf | torch.float32 | () | 最高置信度。 |
result.probs.top5 | list[int] | (<=5) | 排名前 5 的类别 ID。 |
有关每项任务中任务专用的 Results 字段,请参阅按任务划分的预测结果部分。
导出#
将 YOLO26n-cls 模型导出为其他格式,例如 ONNX、CoreML 等。
from ultralytics import YOLO
# Load a model
model = YOLO("yolo26n-cls.pt") # load an official model
model = YOLO("path/to/best.pt") # load a custom-trained model
# Export the model
model.export(format="onnx")下表列出了可用的 YOLO26-cls 导出格式。你可以使用 format 参数导出为任意格式,即 format='onnx' 或 format='engine'。你可以直接在导出的模型上进行预测或验证,即 yolo predict model=yolo26n-cls.onnx。导出完成后,系统会为你的模型显示用法示例。
| 格式 | format 参数 | 模型 | 元数据 | 参数 |
|---|---|---|---|---|
| PyTorch | - | yolo26n-cls.pt | ✅ | - |
| TorchScript | torchscript | yolo26n-cls.torchscript | ✅ | imgsz、quantize、dynamic、nms、batch、device |
| ONNX | onnx | yolo26n-cls.onnx | ✅ | imgsz、quantize、dynamic、simplify、opset、nms、batch、data、fraction、device |
| OpenVINO | openvino | yolo26n-cls_openvino_model/ | ✅ | imgsz、quantize、dynamic、nms、batch、data、fraction、device |
| TensorRT | engine | yolo26n-cls.engine | ✅ | imgsz、quantize、dynamic、simplify、opset、workspace、nms、batch、data、fraction、device |
| CoreML | coreml | yolo26n-cls.mlpackage | ✅ | imgsz、dynamic、quantize、nms、batch、device |
| TF SavedModel | saved_model | yolo26n-cls_saved_model/ | ✅ | imgsz、keras、quantize、opset、nms、batch、data、fraction、device |
| TF GraphDef | pb | yolo26n-cls.pb | ❌ | imgsz、opset、batch、device |
| TF Edge TPU | edgetpu | yolo26n-cls_edgetpu.tflite | ✅ | imgsz、quantize、opset、data、fraction、device |
| PaddlePaddle | paddle | yolo26n-cls_paddle_model/ | ✅ | imgsz、batch、device |
| MNN | mnn | yolo26n-cls.mnn | ✅ | imgsz、batch、dynamic、quantize、simplify、opset、nms、device |
| NCNN | ncnn | yolo26n-cls_ncnn_model/ | ✅ | imgsz、quantize、batch、device |
| IMX500 | imx | yolo26n-cls_imx_model/ | ✅ | imgsz、quantize、data、fraction、nms、device |
| RKNN | rknn | yolo26n-cls_rknn_model/ | ✅ | imgsz、batch、name、quantize、simplify、opset、data、fraction、device |
| ExecuTorch | executorch | yolo26n-cls_executorch_model/ | ✅ | imgsz、batch、device |
| Axelera | axelera | yolo26n-cls_axelera_model/ | ✅ | imgsz、batch、quantize、data、fraction、device |
| DEEPX | deepx | yolo26n-cls_deepx_model/ | ✅ | imgsz、quantize、simplify、opset、data、optimize、device |
| Qualcomm QNN | qnn | yolo26n-cls_qnn.onnx | ✅ | imgsz、batch、name、quantize、simplify、opset、data、fraction、device |
| LiteRT | litert | yolo26n-cls.tflite | ✅ | imgsz、quantize、batch、data、fraction、device |
| Hailo | hailo | yolo26n-cls_hailo_model/ | ✅ | imgsz、name、quantize、data、fraction、simplify、conf、iou |
| Huawei Ascend | ascend | yolo26n-cls_ascend_model/ | ✅ | imgsz、batch、name、quantize、opset、simplify、nms |
| Apple Core AI | coreai | yolo26n-cls.aimodel | ✅ | imgsz、batch、quantize |
nms=None 默认对外部 NMS 使用原始输出。设置 nms=False 以选择可用的无 NMS 检测头;不支持的格式将回退到其本机输出路径。上面的 nms 条目标识了可以通过 nms=True 嵌入 NMS 的格式。
请在导出页面查看完整的 export 详情。
常见问题#
YOLO26 模型(例如
yolo26n-cls.pt)专为高效图像分类而设计。它会为整张图像分配一个类别标签和置信度分数。当只需知道图像的具体类别,而不必识别图像中对象的位置或形状时,这一点尤其有用。要训练 YOLO26 模型,你可以使用 Python 或 CLI 命令。例如,在 MNIST160 数据集上以图像尺寸 64 训练
yolo26n-cls模型 100 个 epoch:示例from ultralytics import YOLO # Load a model model = YOLO("yolo26n-cls.pt") # load a pretrained model (recommended for training) # Train the model results = model.train(data="mnist160", epochs=100, imgsz=64)有关更多配置选项,请访问 Configuration 页面。
你可以使用 Python 或 CLI 命令将经过训练的 YOLO26 模型导出为多种格式。例如,将模型导出为 ONNX 格式:
示例from ultralytics import YOLO # Load a model model = YOLO("yolo26n-cls.pt") # load the trained model # Export the model to ONNX model.export(format="onnx")有关详细的导出选项,请参阅 Export 页面。
要在 MNIST160 等数据集上验证经过训练的模型准确率,你可以使用以下 Python 或 CLI 命令:
示例from ultralytics import YOLO # Load a model model = YOLO("yolo26n-cls.pt") # load the trained model # Validate the model metrics = model.val() # no arguments needed, uses the dataset and settings from training metrics.top1 # top1 accuracy metrics.top5 # top5 accuracy有关更多信息,请访问 Validate 部分。