EfficientNet PyTorch

快速入门

使用 pip install efficientnet_pytorch 安装，并通过以下代码加载预训练的 EfficientNet 模型：

from efficientnet_pytorch import EfficientNet
model = EfficientNet.from_pretrained('efficientnet-b0')

更新日志

更新（2021年4月2日）

EfficientNetV2 论文 已发布！在您阅读本文时，我正在着手实现它 :)

关于 EfficientNetV2：

EfficientNetV2 是一个新的卷积神经网络系列，其训练速度更快，参数效率也优于之前的模型。为了开发这一系列模型，我们结合了训练感知的神经架构搜索和缩放技术，以共同优化训练速度和参数效率。这些模型是从包含 Fused-MBConv 等新操作符的搜索空间中搜索出来的。

以下是对比图：

更新（2020年8月25日）

本次更新新增：

• 一个新的 include_top 参数（默认为 True）(#208)
• 使用 sotabench 进行持续测试
• 代码质量改进及修复 (#215 和 #223)

更新（2020年5月14日）

本次更新增加了全面的注释和文档说明（感谢 @workingcoder）。

更新（2020年1月23日）

本次更新新增了一类基于对抗训练的预训练模型，称为 advprop。需要注意的是，advprop 预训练模型所需的预处理与常规 ImageNet 预处理略有不同。因此，默认情况下不会使用 advprop 模型。若要加载 advprop 模型，请使用以下代码：

model = EfficientNet.from_pretrained("efficientnet-b0", advprop=True)

此外，还有一个新的大型 efficientnet-b8 预训练模型，仅以 advprop 形式提供。使用这些模型时，请按如下方式替换 ImageNet 预处理代码：

if advprop:  # 对于使用 advprop 预训练权重的模型
    normalize = transforms.Lambda(lambda img: img * 2.0 - 1.0)
else:
    normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                     std=[0.229, 0.224, 0.225])

此次更新还解决了多个其他问题 (#115, #128)。

更新（2019年10月15日）

本次更新允许您选择是否使用内存高效的 Swish 激活函数。默认选择内存高效的版本，但在使用 PyTorch JIT 导出时无法使用该版本。为此，我们也提供了标准的（适合导出的）Swish 激活函数。要切换到适合导出的版本，只需在加载所需模型后调用 model.set_swish(memory_efficient=False)。此次更新解决了 #88 和 #89 中的问题。

更新（2019年10月12日）

本次更新使 Swish 激活函数更加节省内存。同时，也解决了 #72, #73, #85 和 #86 中的请求。感谢所有提交请求的作者！

更新（2019年7月31日）

请使用 pip install --upgrade efficientnet-pytorch 升级 pip 包。

现在 B6 和 B7 模型已上线。此外，所有预训练模型均已更新为使用 AutoAugment 预处理，这将带来整体性能的提升。使用方法与之前相同：

from efficientnet_pytorch import EfficientNet
model = EfficientNet.from_pretrained('efficientnet-b7')

更新（2019年6月29日）

本次更新增加了便捷的模型导出功能 (#20) 和特征提取功能 (#38)。

• 示例：导出为 ONNX 格式
• 示例：特征提取
• 此外：修复了一个 CUDA/CPU 错误 (#32)

现在，加载一个具有新类别数目的预训练模型进行迁移学习也变得非常简单：

model = EfficientNet.from_pretrained('efficientnet-b1', num_classes=23)

更新（2019年6月23日）

B4 和 B5 模型现已上线。它们的使用方式与其他模型完全相同：

from efficientnet_pytorch import EfficientNet
model = EfficientNet.from_pretrained('efficientnet-b4')

概述

本仓库包含了对 EfficientNet 的逐层 PyTorch 重实现，以及预训练模型和示例。

该实现的目标是简单易用、高度可扩展，并且易于集成到您自己的项目中。目前该实现仍在开发中——新功能正在不断添加。

现阶段，您可以轻松地：

• 加载预训练的 EfficientNet 模型
• 将 EfficientNet 模型用于分类或特征提取
• 在 ImageNet 或您自己的数据集上评估 EfficientNet 模型

即将推出的功能：在未来几天内，您将能够：

• 使用一条简单的命令从头开始在 ImageNet 上训练新模型
• 快速地在您自己的数据集上微调 EfficientNet 模型
• 将 EfficientNet 模型导出用于生产环境

目录

1. 关于 EfficientNet
2. 关于 EfficientNet-PyTorch
3. 安装
4. 使用
   • 加载预训练模型
   • 示例：分类
   • 示例：特征提取
   • 示例：导出为 ONNX 格式
5. 贡献

关于 EfficientNet

如果你是第一次接触 EfficientNet，以下是来自 TensorFlow 官方实现的直接解释：

EfficientNet 是一系列图像分类模型，它们不仅达到了最先进的准确率，而且在规模和速度上都比之前的模型小了一个数量级。我们基于 AutoML 和复合缩放方法开发了 EfficientNet。具体来说，我们首先使用 AutoML Mobile 框架 来构建一个移动端大小的基准网络，命名为 EfficientNet-B0；然后，我们采用复合缩放方法将该基准网络逐步放大，从而得到 EfficientNet-B1 到 B7。

EfficientNet 在 ImageNet 数据集上实现了最先进的准确率，同时效率提升了整整一个数量级：

• 在高准确率场景下，我们的 EfficientNet-B7 在 ImageNet 上达到了 84.4% 的 top-1 准确率和 97.1% 的 top-5 准确率，参数量仅为 6600 万，浮点运算次数为 370 亿，相比此前的最佳模型 Gpipe，其规模缩小了 8.4 倍，CPU 推理速度提升了 6.1 倍。

• 在中等准确率场景下，我们的 EfficientNet-B1 在 CPU 推理速度上比 ResNet-152 快 5.7 倍，而参数量仅为其 7.6 分之一，且在 ImageNet 上的准确率相近。

• 与广泛使用的 ResNet-50 相比，在相似的浮点运算次数约束下，我们的 EfficientNet-B4 将 top-1 准确率从 ResNet-50 的 76.3% 提升至 82.6%（提高了 6.3%）。

关于 EfficientNet PyTorch

EfficientNet PyTorch 是 EfficientNet 的 PyTorch 重新实现版本。它与 原始的 TensorFlow 实现 保持一致，因此可以轻松地从 TensorFlow 检查点加载权重。与此同时，我们致力于使我们的 PyTorch 实现尽可能简单、灵活且易于扩展。

如果你有任何功能请求或问题，请随时在 GitHub 上提交 Issue！

安装

通过 pip 安装：

pip install efficientnet_pytorch

或者从源码安装：

git clone https://github.com/lukemelas/EfficientNet-PyTorch
cd EfficientNet-PyTorch
pip install -e .

使用

加载预训练模型

加载一个 EfficientNet 模型：

from efficientnet_pytorch import EfficientNet
model = EfficientNet.from_name('efficientnet-b0')

加载一个预训练的 EfficientNet 模型：

from efficientnet_pytorch import EfficientNet
model = EfficientNet.from_pretrained('efficientnet-b0')

以下是各模型的详细信息：

名称	参数量	Top-1 准确率	是否预训练？
efficientnet-b0	5.3M	76.3	✓
efficientnet-b1	7.8M	78.8	✓
efficientnet-b2	9.2M	79.8	✓
efficientnet-b3	12M	81.1	✓
efficientnet-b4	19M	82.6	✓
efficientnet-b5	30M	83.3	✓
efficientnet-b6	43M	84.0	✓
efficientnet-b7	66M	84.4	✓

示例：分类任务

以下是一个简单完整的示例。你也可以在 examples/simple 中找到对应的 Jupyter Notebook，或者在 Colab Notebook 中查看。

假设你的当前目录下有一个 img.jpg 文件和一个 labels_map.txt 文件（包含 ImageNet 类别名称）。这两个文件都包含在 examples/simple 中。

import json
from PIL import Image
import torch
from torchvision import transforms

from efficientnet_pytorch import EfficientNet
model = EfficientNet.from_pretrained('efficientnet-b0')

# 图像预处理
tfms = transforms.Compose([transforms.Resize(224), transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),])
img = tfms(Image.open('img.jpg')).unsqueeze(0)
print(img.shape) # torch.Size([1, 3, 224, 224])

# 加载 ImageNet 类别名称
labels_map = json.load(open('labels_map.txt'))
labels_map = [labels_map[str(i)] for i in range(1000)]

# 分类
model.eval()
with torch.no_grad():
    outputs = model(img)

# 打印预测结果
print('-----')
for idx in torch.topk(outputs, k=5).indices.squeeze(0).tolist():
    prob = torch.softmax(outputs, dim=1)[0, idx].item()
    print('{label:<75} ({p:.2f}%)'.format(label=labels_map[idx], p=prob*100))

示例：特征提取

你可以使用 model.extract_features 轻松提取特征：

from efficientnet_pytorch import EfficientNet
model = EfficientNet.from_pretrained('efficientnet-b0')

# ... 图像预处理步骤同分类示例 ...
print(img.shape) # torch.Size([1, 3, 224, 224])

features = model.extract_features(img)
print(features.shape) # torch.Size([1, 1280, 7, 7])

示例：导出为 ONNX

现在，将模型导出为 ONNX 以用于生产部署变得非常简单：

import torch
from efficientnet_pytorch import EfficientNet

model = EfficientNet.from_pretrained('efficientnet-b1')
dummy_input = torch.randn(10, 3, 240, 240)

model.set_swish(memory_efficient=False)
torch.onnx.export(model, dummy_input, "test-b1.onnx", verbose=True)

这里 是一个 Colab 示例。

ImageNet

有关在 ImageNet 上评估的详细信息，请参阅 examples/imagenet。

贡献

如果你发现任何错误，请创建一个 GitHub 问题，或者更好的是，提交一个拉取请求。同样地，如果你有任何疑问，也请直接在 GitHub 上提出。

我非常期待看到社区如何利用这些模型！

EfficientNet-PyTorch 快速上手指南

EfficientNet-PyTorch 是 Google EfficientNet 系列模型的 PyTorch 复现版本，提供了从 B0 到 B7 多种预训练模型。该库以高精度、低参数量和高推理速度著称，非常适合图像分类和特征提取任务。

环境准备

在开始之前，请确保您的开发环境满足以下要求：

• 操作系统：Linux, macOS, 或 Windows
• Python 版本：推荐 Python 3.6 及以上
• 核心依赖：
  • torch (PyTorch)
  • torchvision
  • Pillow (用于图像处理)
  • numpy

您可以使用以下命令检查 PyTorch 是否已安装：

python -c "import torch; print(torch.__version__)"

安装步骤

推荐使用 pip 进行安装。为了提高下载速度，国内用户建议使用清华或阿里镜像源。

方式一：通过 pip 安装（推荐）

使用国内镜像源加速安装：

pip install efficientnet_pytorch -i https://pypi.tuna.tsinghua.edu.cn/simple

或使用官方源：

pip install efficientnet_pytorch

方式二：从源码安装

如果您需要最新的功能或贡献代码，可以从 GitHub 克隆源码安装：

git clone https://github.com/lukemelas/EfficientNet-PyTorch
cd EfficientNet-PyTorch
pip install -e .

基本使用

以下是加载预训练模型并进行图像分类的最简示例。

1. 加载预训练模型

只需一行代码即可加载指定的 EfficientNet 模型（如 efficientnet-b0）。模型权重会自动下载并缓存。

from efficientnet_pytorch import EfficientNet

# 加载预训练的 efficientnet-b0 模型
model = EfficientNet.from_pretrained('efficientnet-b0')

# 将模型设置为评估模式
model.eval()

支持的模型名称包括：efficientnet-b0 至 efficientnet-b7。

2. 图像预处理与推理

EfficientNet 需要特定的图像预处理（缩放、转 Tensor、标准化）。以下是一个完整的分类示例：

import torch
from PIL import Image
from torchvision import transforms
from efficientnet_pytorch import EfficientNet

# 1. 加载模型
model = EfficientNet.from_pretrained('efficientnet-b0')
model.eval()

# 2. 定义预处理流程
# 注意：EfficientNet 通常需要将图像缩放到对应模型的分辨率 (B0 为 224x224)
tfms = transforms.Compose([
    transforms.Resize(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])

# 3. 加载并处理图像
img = Image.open('img.jpg') # 替换为您的图片路径
img_tensor = tfms(img).unsqueeze(0) # 增加 batch 维度 [1, 3, 224, 224]

# 4. 执行推理
with torch.no_grad():
    outputs = model(img_tensor)

# 5. 获取预测结果 (Top-1)
predicted_class_idx = outputs.argmax(dim=1).item()
print(f"预测类别索引：{predicted_class_idx}")

# 如果需要获取概率分布
probabilities = torch.softmax(outputs, dim=1)[0]
print(f"最高置信度：{probabilities[predicted_class_idx].item():.4f}")

3. 提取特征

如果您需要将 EfficientNet 作为骨干网络用于其他任务（如检测、分割），可以使用 extract_features 方法直接获取中间层特征图，而不经过最后的分类层：

from efficientnet_pytorch import EfficientNet
import torch

model = EfficientNet.from_pretrained('efficientnet-b0')
model.eval()

# 假设 img_tensor 已经过预处理，形状为 [1, 3, 224, 224]
features = model.extract_features(img_tensor)

# 输出特征图形状，例如: torch.Size([1, 1280, 7, 7])
print(features.shape)