Keras CV Attention Models:强大而灵活的计算机视觉模型库

keras_cv_attention_models:强大而灵活的计算机视觉模型库

keras_cv_attention_models(简称kecam)是一个基于Keras的计算机视觉模型库,包含了大量最新的注意力机制模型,支持图像分类、目标检测、语言模型等多种任务。该库由GitHub用户leondgarse开发和维护,旨在为深度学习研究者和实践者提供一个强大而灵活的工具。

主要特性

1. 丰富的模型支持:包括EfficientNet、VOLO、CoAtNet、ConvNeXt等多种SOTA模型。

2. 多任务支持:除了图像分类,还支持目标检测(YOLO系列)、语言模型(GPT2、LLaMA2)等任务。

3. 预训练模型:提供大量在ImageNet等数据集上预训练的模型权重。

4. 灵活的训练和评估:提供了训练和评估脚本,支持自定义数据集。

5. 模型手术:支持模型结构修改,如替换激活函数、融合卷积和BN层等。

6. 多backend支持:除了TensorFlow,还支持PyTorch作为backend。

7. 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提供了大量最新的计算机视觉模型,主要分为以下几类:

1. 图像分类模型:

   • EfficientNet系列
   • VOLO
   • CoAtNet
   • ConvNeXt
   • Swin Transformer
   • MaxViT 等

2. 目标检测模型:

   • EfficientDet
   • YOLO系列(YOLOR, YOLOv7, YOLOv8等)

3. 语言模型:

   • GPT2
   • LLaMA2

4. 分割模型:

   • 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模型。

高级特性

1. 模型手术:model_surgery模块提供了一些修改模型结构的函数,如替换激活函数、融合卷积和BN层等。

2. 渐进式训练:支持在训练过程中逐步增加输入图像尺寸,可以加快训练速度并提高精度。

3. 迁移学习:支持冻结backbone或norm层进行迁移学习。

4. 自定义数据集:提供了脚本用于创建自定义数据集的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,相信它会为您的工作带来便利和启发。