keras_cv_attention_models:强大而灵活的计算机视觉模型库
keras_cv_attention_models(简称kecam)是一个基于Keras的计算机视觉模型库,包含了大量最新的注意力机制模型,支持图像分类、目标检测、语言模型等多种任务。该库由GitHub用户leondgarse开发和维护,旨在为深度学习研究者和实践者提供一个强大而灵活的工具。
主要特性
-
丰富的模型支持:包括EfficientNet、VOLO、CoAtNet、ConvNeXt等多种SOTA模型。
-
多任务支持:除了图像分类,还支持目标检测(YOLO系列)、语言模型(GPT2、LLaMA2)等任务。
-
预训练模型:提供大量在ImageNet等数据集上预训练的模型权重。
-
灵活的训练和评估:提供了训练和评估脚本,支持自定义数据集。
-
模型手术:支持模型结构修改,如替换激活函数、融合卷积和BN层等。
-
多backend支持:除了TensorFlow,还支持PyTorch作为backend。
-
ONNX导出:支持将模型导出为ONNX格式,便于跨平台部署。
安装和基本使用
可以通过pip安装keras_cv_attention_models:
pip install -U kecam
基本使用示例:
from keras_cv_attention_models import volo
model = volo.VOLO_d1(pretrained="imagenet")
# 预测
from keras_cv_attention_models.test_images import cat
preds = model(model.preprocess_input(cat()))
print(model.decode_predictions(preds))
这将加载预训练的VOLO_d1模型,并对一张猫的图片进行预测。
模型概览
keras_cv_attention_models提供了大量最新的计算机视觉模型,主要分为以下几类:
-
图像分类模型:
- EfficientNet系列
- VOLO
- CoAtNet
- ConvNeXt
- Swin Transformer
- MaxViT 等
-
目标检测模型:
- EfficientDet
- YOLO系列(YOLOR, YOLOv7, YOLOv8等)
-
语言模型:
- GPT2
- LLaMA2
-
分割模型:
- YOLOv8 Segmentation
- Segment Anything
每个模型都提供了预训练权重,可以直接用于推理或微调。
模型训练
keras_cv_attention_models提供了训练脚本train_script.py,可以方便地训练各种模型。以下是一个训练AotNet50的示例:
CUDA_VISIBLE_DEVICES='0' TF_XLA_FLAGS="--tf_xla_auto_jit=2" python3 train_script.py --seed 0 -s aotnet50
这将使用默认参数训练AotNet50模型。train_script.py支持多种参数来控制训练过程,如学习率、优化器、数据增强等。
对于目标检测任务,可以使用coco_train_script.py进行COCO数据集的训练。
模型评估
eval_script.py用于评估模型的准确率:
CUDA_VISIBLE_DEVICES='1' python3 eval_script.py -m regnet.RegNetZD8
这将评估预训练的RegNetZD8模型在ImageNet验证集上的准确率。
模型导出
可以将模型导出为ONNX格式,便于在不同平台上部署:
from keras_cv_attention_models import volo, model_surgery
mm = volo.VOLO_d1(pretrained=True)
model_surgery.export_onnx(mm, fuse_conv_bn=True, simplify=True)
这将导出一个优化后的ONNX模型。
高级特性
-
模型手术:model_surgery模块提供了一些修改模型结构的函数,如替换激活函数、融合卷积和BN层等。
-
渐进式训练:支持在训练过程中逐步增加输入图像尺寸,可以加快训练速度并提高精度。
-
迁移学习:支持冻结backbone或norm层进行迁移学习。
-
自定义数据集:提供了脚本用于创建自定义数据集的JSON文件,可用于训练。
性能对比
keras_cv_attention_models提供了model_summary.csv文件,包含了各个模型的参数量、FLOPs、推理速度等信息,方便用户进行模型选择和对比。
从图中可以看出,不同模型在准确率和效率之间有不同的权衡。用户可以根据自己的需求选择合适的模型。
总结
keras_cv_attention_models是一个功能丰富、易用性强的计算机视觉模型库。它不仅提供了大量最新的模型实现,还提供了完整的训练、评估和部署流程。无论是研究新的模型架构,还是将模型应用到实际问题中,keras_cv_attention_models都是一个强大的工具。
对于深度学习研究者来说,keras_cv_attention_models提供了一个方便的平台来实验和比较不同的模型架构。对于实践者来说,它提供了即用即取的预训练模型和完整的训练流程,可以快速应用到实际问题中。
未来,keras_cv_attention_models还将继续跟进最新的研究进展,添加更多的模型和功能。研究者和开发者也可以通过GitHub为这个项目做出贡献,共同推动计算机视觉领域的发展。
如果您正在从事计算机视觉相关的研究或应用,不妨尝试使用keras_cv_attention_models,相信它会为您的工作带来便利和启发。
