gen-efficientnet-pytorch:高效的EfficientNet PyTorch实现

gen-efficientnet-pytorch:高效的EfficientNet PyTorch实现

gen-efficientnet-pytorch是一个通用的EfficientNet PyTorch实现,支持多种高效CNN模型,包括EfficientNet、MixNet、MobileNetV3等。该项目由Ross Wightman开发,旨在提供一个灵活、高效的EfficientNet系列模型实现。

主要特性

1. 支持多种高效CNN模型

gen-efficientnet-pytorch实现了以下模型:

• EfficientNet (B0-B8, L2)
• EfficientNet-Lite
• MixNet (Small, Medium, Large)
• MobileNetV3 (Large, Small)
• MNASNet
• FBNet-C
• Single-Path NAS

这些模型都是基于MobileNetV1/V2的block序列派生出来的计算/参数高效的架构,包括一些通过神经架构搜索(NAS)得到的模型。

2. 提供预训练模型

项目提供了多种预训练模型,包括:

• 作者自行训练的模型,如EfficientNet-B3达到82.1% top-1准确率
• 从TensorFlow官方实现移植的预训练权重

这些预训练模型可以直接用于迁移学习或评估。

3. 便捷的使用接口

gen-efficientnet-pytorch提供了简单易用的接口:

import geffnet

# 创建预训练模型
model = geffnet.create_model('efficientnet_b0', pretrained=True)

# 使用entrypoint直接创建
model = geffnet.efficientnet_b2(pretrained=True, drop_rate=0.25)

# 创建为nn.Sequential,方便与fast.ai等框架集成
model = geffnet.mixnet_l(pretrained=True, as_sequential=True)

4. 支持模型导出

项目提供了导出模型到ONNX、Caffe2等格式的脚本,方便部署:

• 导出到ONNX
• 优化ONNX图
• 验证ONNX模型
• 转换ONNX到Caffe2
• 在Caffe2中验证和benchmark

性能表现

gen-efficientnet-pytorch中的模型在ImageNet上展现了出色的性能,以下是部分模型的top-1准确率:

• EfficientNet-B3: 82.1%
• MixNet-XL: 80.5%
• EfficientNet-B2: 80.3%
• MobileNetV3 Large: 75.8%
• MNASNet-A1: 75.4%

项目还提供了详细的benchmarks,包括参数量、MACs等指标。

使用方法

1. 安装

pip install geffnet

2. 创建模型

import geffnet

# 创建预训练模型
model = geffnet.create_model('efficientnet_b0', pretrained=True)

# 评估模式
model.eval()

# 前向传播
output = model(input) 

3. 导出模型

# 导出到ONNX
python onnx_export.py --model mobilenetv3_large_100 ./mobilenetv3_100.onnx

# 验证ONNX模型
python onnx_validate.py /imagenet/validation/ --onnx-input ./mobilenetv3_100.onnx

总结

gen-efficientnet-pytorch为EfficientNet系列模型提供了一个高效、灵活的PyTorch实现。它支持多种模型架构,提供预训练权重,使用简单,并支持模型导出,是一个很好的EfficientNet相关研究和应用的工具。

对于想要在PyTorch中使用EfficientNet等高效CNN模型的研究者和开发者来说,gen-efficientnet-pytorch是一个值得尝试的选择。它可以帮助你快速地将这些先进的模型集成到自己的项目中,进行迁移学习或进一步的研究。

希望这个项目能够帮助更多人了解和使用EfficientNet等高效CNN模型,推动计算机视觉领域的发展。如果你对这个项目感兴趣,不妨star一下GitHub仓库,并尝试使用它来提升你的深度学习项目的效率。