EfficientNet: Rethinking Model Scaling for Convolutional Neural Networks

Quickstart

Install with pip install efficientnet_pytorch and load a pretrained EfficientNet with:

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

Overview

This repository contains an op-for-op PyTorch reimplementation of EfficientNet, along with pre-trained models and examples.
The goal of this implementation is to be simple, highly extensible, and easy to integrate into your own projects. This implementation is a work in progress -- new features are currently being implemented.

Details about the models are below:

Name	# Params	Top-1 Acc.	Pretrained?
`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	✓

Example: Feature Extraction

You can easily extract features with model.extract_features:

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

# ... image preprocessing as in the classification example ...
print(img.shape) # torch.Size([1, 3, 224, 224])

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