A PyTorch implementation of EfficientNet

Overview

EfficientNet PyTorch

Quickstart

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

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

Updates

Update (Aug 25, 2020)

This update adds:

  • A new include_top (default: True) option (#208)
  • Continuous testing with sotabench
  • Code quality improvements and fixes (#215 #223)

Update (May 14, 2020)

This update adds comprehensive comments and documentation (thanks to @workingcoder).

Update (January 23, 2020)

This update adds a new category of pre-trained model based on adversarial training, called advprop. It is important to note that the preprocessing required for the advprop pretrained models is slightly different from normal ImageNet preprocessing. As a result, by default, advprop models are not used. To load a model with advprop, use:

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

There is also a new, large efficientnet-b8 pretrained model that is only available in advprop form. When using these models, replace ImageNet preprocessing code as follows:

if advprop:  # for models using advprop pretrained weights
    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])

This update also addresses multiple other issues (#115, #128).

Update (October 15, 2019)

This update allows you to choose whether to use a memory-efficient Swish activation. The memory-efficient version is chosen by default, but it cannot be used when exporting using PyTorch JIT. For this purpose, we have also included a standard (export-friendly) swish activation function. To switch to the export-friendly version, simply call model.set_swish(memory_efficient=False) after loading your desired model. This update addresses issues #88 and #89.

Update (October 12, 2019)

This update makes the Swish activation function more memory-efficient. It also addresses pull requests #72, #73, #85, and #86. Thanks to the authors of all the pull requests!

Update (July 31, 2019)

Upgrade the pip package with pip install --upgrade efficientnet-pytorch

The B6 and B7 models are now available. Additionally, all pretrained models have been updated to use AutoAugment preprocessing, which translates to better performance across the board. Usage is the same as before:

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

Update (June 29, 2019)

This update adds easy model exporting (#20) and feature extraction (#38).

It is also now incredibly simple to load a pretrained model with a new number of classes for transfer learning:

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

Update (June 23, 2019)

The B4 and B5 models are now available. Their usage is identical to the other models:

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

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.

At the moment, you can easily:

  • Load pretrained EfficientNet models
  • Use EfficientNet models for classification or feature extraction
  • Evaluate EfficientNet models on ImageNet or your own images

Upcoming features: In the next few days, you will be able to:

  • Train new models from scratch on ImageNet with a simple command
  • Quickly finetune an EfficientNet on your own dataset
  • Export EfficientNet models for production

Table of contents

  1. About EfficientNet
  2. About EfficientNet-PyTorch
  3. Installation
  4. Usage
  5. Contributing

About EfficientNet

If you're new to EfficientNets, here is an explanation straight from the official TensorFlow implementation:

EfficientNets are a family of image classification models, which achieve state-of-the-art accuracy, yet being an order-of-magnitude smaller and faster than previous models. We develop EfficientNets based on AutoML and Compound Scaling. In particular, we first use AutoML Mobile framework to develop a mobile-size baseline network, named as EfficientNet-B0; Then, we use the compound scaling method to scale up this baseline to obtain EfficientNet-B1 to B7.

EfficientNets achieve state-of-the-art accuracy on ImageNet with an order of magnitude better efficiency:

  • In high-accuracy regime, our EfficientNet-B7 achieves state-of-the-art 84.4% top-1 / 97.1% top-5 accuracy on ImageNet with 66M parameters and 37B FLOPS, being 8.4x smaller and 6.1x faster on CPU inference than previous best Gpipe.

  • In middle-accuracy regime, our EfficientNet-B1 is 7.6x smaller and 5.7x faster on CPU inference than ResNet-152, with similar ImageNet accuracy.

  • Compared with the widely used ResNet-50, our EfficientNet-B4 improves the top-1 accuracy from 76.3% of ResNet-50 to 82.6% (+6.3%), under similar FLOPS constraint.

About EfficientNet PyTorch

EfficientNet PyTorch is a PyTorch re-implementation of EfficientNet. It is consistent with the original TensorFlow implementation, such that it is easy to load weights from a TensorFlow checkpoint. At the same time, we aim to make our PyTorch implementation as simple, flexible, and extensible as possible.

If you have any feature requests or questions, feel free to leave them as GitHub issues!

Installation

Install via pip:

pip install efficientnet_pytorch

Or install from source:

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

Usage

Loading pretrained models

Load an EfficientNet:

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

Load a pretrained EfficientNet:

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

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: Classification

Below is a simple, complete example. It may also be found as a jupyter notebook in examples/simple or as a Colab Notebook.

We assume that in your current directory, there is a img.jpg file and a labels_map.txt file (ImageNet class names). These are both included in 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')

# Preprocess image
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])

# Load ImageNet class names
labels_map = json.load(open('labels_map.txt'))
labels_map = [labels_map[str(i)] for i in range(1000)]

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

# Print predictions
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))

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])

Example: Export to ONNX

Exporting to ONNX for deploying to production is now simple:

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)

Here is a Colab example.

ImageNet

See examples/imagenet for details about evaluating on ImageNet.

Contributing

If you find a bug, create a GitHub issue, or even better, submit a pull request. Similarly, if you have questions, simply post them as GitHub issues.

I look forward to seeing what the community does with these models!

Comments
  • efficientnet-b8 and AdvProp

    efficientnet-b8 and AdvProp

    With advprop, efficientnet got greater score in ImageNet. Would you update to the new ckpt? the paper: https://arxiv.org/pdf/1911.09665.pdf https://github.com/tensorflow/tpu/tree/master/models/official/efficientnet

    enhancement 
    opened by seefun 19
  • FLOPs count seems to be off

    FLOPs count seems to be off

    Hi Luke, Thanks for your great work !

    I am interested in the FLOPs of the models implemented.

    I have always been using this to count FLOPs.

    And for most cases, models from vision:

    1. Resnet50
    2. VGG19
    3. Densenet121
    4. Densenet169
    5. Shufflenetv2_2_0

    they all seems to match the FLOPs count from paper.

    However in this case it is not. I can not think of a reason why, when i traverse the code down to each module, even added the FLOPs count for both padding inside the conv2dblock and Swish activation .

    Do you have any idea?

    Thanks in advance.

    opened by matthewygf 17
  • NotImplementedError when use forward function.

    NotImplementedError when use forward function.

    I try define a model based on pretrained EfficientNet as below. But I get a NotImplementedError: when use 'forward' function. However, when I use other pretrained CNN e.g., resnet18 from torchvision, there is no such problem. Can anyone help me? Thanks a lot

    'Model definition' class EfficientNet_scene(nn.Module):

    def __init__(self,model_name='efficientnet-b0',class_num=45,initfc_type='normal',gain=0.2):
        super(EfficientNet_scene, self).__init__()
        model = EfficientNet.from_pretrained(model_name)
        aul = [*model.children()]
        self.features = nn.Sequential(*aul[:-1])
        self.fc = nn.Linear(aul[-1].in_features,class_num)
    
        if hasattr(self.fc, 'bias') and self.fc.bias is not None:
            nn.init.constant_(self.fc.bias.data, 0.0)
        if initfc_type == 'normal':
            nn.init.normal_(self.fc.weight.data, 0.0, gain)
        elif initfc_type == 'xavier':
            nn.init.xavier_normal_(self.fc.weight.data, gain=gain)
        elif initfc_type == 'kaiming':
            nn.init.kaiming_normal_(self.fc.weight.data, a=0, mode='fan_in')
        elif initfc_type == 'orthogonal':
            nn.init.orthogonal_(self.fc.weight.data, gain=gain)
    
    def forward(self,x):
        x = self.features(x)
        x = self.fc(x)
        return x
    

    net = EfficientNet_scene() image = torch.randn(1,3,224,224) b = net(image)

    Error information: NotImplementedError Traceback (most recent call last) in () 37 print(net) 38 image = torch.randn(1,3,224,224) ---> 39 b = net(image) 40 print(b)

    ~/anaconda3/lib/python3.6/site-packages/torch/nn/modules/module.py in call(self, *input, **kwargs) 491 result = self._slow_forward(*input, **kwargs) 492 else: --> 493 result = self.forward(*input, **kwargs) 494 for hook in self._forward_hooks.values(): 495 hook_result = hook(self, input, result)

    in forward(self, x) 30 31 def forward(self,x): ---> 32 x = self.features(x) 33 x = x.reshape(x.size(0), -1) 34 x = self.fc(x)

    ~/anaconda3/lib/python3.6/site-packages/torch/nn/modules/module.py in call(self, *input, **kwargs) 491 result = self._slow_forward(*input, **kwargs) 492 else: --> 493 result = self.forward(*input, **kwargs) 494 for hook in self._forward_hooks.values(): 495 hook_result = hook(self, input, result)

    ~/anaconda3/lib/python3.6/site-packages/torch/nn/modules/container.py in forward(self, input) 90 def forward(self, input): 91 for module in self._modules.values(): ---> 92 input = module(input) 93 return input 94

    ~/anaconda3/lib/python3.6/site-packages/torch/nn/modules/module.py in call(self, *input, **kwargs) 491 result = self._slow_forward(*input, **kwargs) 492 else: --> 493 result = self.forward(*input, **kwargs) 494 for hook in self._forward_hooks.values(): 495 hook_result = hook(self, input, result)

    ~/anaconda3/lib/python3.6/site-packages/torch/nn/modules/module.py in forward(self, *input) 86 registered hooks while the latter silently ignores them. 87 """ ---> 88 raise NotImplementedError 89 90 def register_buffer(self, name, tensor):

    NotImplementedError:

    opened by henanjun 15
  • How to tranform efficient-pytorch to efficient-onnx

    How to tranform efficient-pytorch to efficient-onnx

    I have tried to convert efficient-pytorch to efficient-onnx with api (torch.onnx.export), but I meet a problem showing below info Failed to export an ONNX attribute, since it's not constant, please try to make things (e.g., kernel size) static if possible How could I fix it ?

    opened by peyer 14
  • Memory Issues

    Memory Issues

    Hi Luke,

    Thank you for the awesome work. I tried running EfficientNet-B0 on my GTX 1070 (8GB RAM) with an input batch of dimension [44x1x256x256] (single channel image) and I am running into 'CUDA out of memory' (with the model in 'training' mode).

    I tried running another implementation and wasn't getting this issue, and after digging in the code, it seems as if the implementation for MBConv (or the re-iteration of MBConv) was too memory hungry.

    I really like your implementation of EfficientNet and if I did have more time, I would definitely have a deeper dive into your code. At the mean time, if possible, could you help me check this issue out (maybe it'll speed up training in the future?) ? Thank you!

    opened by mxtsai 14
  • urllib.error.HTTPError: HTTP Error 403: Forbidden when downloading the efficientnet-b4-e116e8b3.pth

    urllib.error.HTTPError: HTTP Error 403: Forbidden when downloading the efficientnet-b4-e116e8b3.pth

    Hi,

    the error "Anonymous caller does not have storage.objects.get access to public-models/efficientnet-b4-e116e8b3.pth." appears when trying to download the pretrained efficientnet-b4.

    opened by qiminchen 13
  • Pretrained models have terrible performance

    Pretrained models have terrible performance

    I evaluate all the efficient architecture with the pre-trained weights available on this repo and the performances are different than the ones shared in this repo.
    What is going on?

    |                 |    top1 |    top5 |   time (s) |
    |:----------------|--------:|--------:|-----------:|
    | efficientnet-b0 | 0.7465  | 0.91932 |    78.2806 |
    | efficientnet-b1 | 0.7461  | 0.91616 |   120.922  |
    | efficientnet-b2 | 0.79394 | 0.9459  |   159.07   |
    | efficientnet-b3 | 0.81212 | 0.95526 |   253.05   |
    | efficientnet-b4 | 0.82788 | 0.96234 |   542.88   |
    | efficientnet-b5 | 0.8345  | 0.9663  |  1049.4    |
    | efficientnet-b6 | 0.84024 | 0.96908 |  1853.14   |
    | efficientnet-b7 | 0.84106 | 0.9692  |  3245.83   |
    

    Images are resized using the following code

    resize_size = {
    
        'efficientnet-b0': 224,
        'efficientnet-b1': 240,
        'efficientnet-b2': 260,
        'efficientnet-b3': 300,
        'efficientnet-b4': 380,
        'efficientnet-b5': 456,
        'efficientnet-b6': 528,
        'efficientnet-b7': 600,
    
    }
    
    transform = Compose([
                Resize(size, Image.BICUBIC),
                CenterCrop(size),
                ToTensor(),
                Normalize(
                mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
       ])
    
    valid_dataset = ImageNet(root='/home/zuppif/Downloads/ImageNet', split='val', transform=transform)
    
    
    valid_loader = torch.utils.data.DataLoader(valid_dataset, batch_size=batch_size, shuffle=False,
                                                    num_workers=12, pin_memory=True)
    

    My code to compute (I used sotabencheval) top-1 and top-5 is correct.

    opened by FrancescoSaverioZuppichini 11
  • How to remove the last layer?

    How to remove the last layer?

    Hi, Great repo! I'm doing an image retrieval task. I've already trained the model on my dataset and would like to use the backbone for feature extraction. In other words, I would like the output to be a feature vector. Should I simply remove the last _fc layer? Thank you in advance! Best, Zach

    opened by Zacchaeus14 7
  • EfficientNET.onnx does not run in TensorRT

    EfficientNET.onnx does not run in TensorRT

    Hi, I've got this error for running the converted EfficientNet from PyTorch to Onnx in TensorRT:

    Traceback (most recent call last):
      File "tensorrt_python.py", line 59, in <module>
        context = engine.create_execution_context()
    AttributeError: 'NoneType' object has no attribute 'create_execution_context'
    

    Can anybody help me?

    TensorRT version: 6.1.05 Pytorch: 1.1.0 Onnx: 1.5.0

    opened by Soroorsh 7
  • Why do we need `Conv2dStaticSamePadding`

    Why do we need `Conv2dStaticSamePadding`

    If I do not need to transform the tensorflow pre-trained weight into this model, instead I train it from scratch. Can I just delete the Conv2dStaticSamePadding?

    opened by yifanjiang19 7
  • Input type (torch.FloatTensor) and weight type (torch.cuda.FloatTensor) should be the same

    Input type (torch.FloatTensor) and weight type (torch.cuda.FloatTensor) should be the same

    Hi there! Thanks for your great repo, but i faced with some difficulties while trying to inference on device ('cuda:0'). My code: device = torch.device('cuda:0') 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.png')).unsqueeze(0) img.to(device) model.to(device) model.eval() with torch.no_grad(): outputs = model(img) And thrown error: RuntimeError: Input type (torch.FloatTensor) and weight type (torch.cuda.FloatTensor) should be the same

    opened by litvinich 7
  • multilabel multiclass image classification

    multilabel multiclass image classification

    is it possible to do multilabel and multiclass image classification with efficientnet b5? i googled around and there's a page saying in order to do multilabeling, we need to change the classifier activation to sigmoid and the criterion to BinaryCrossEntropy. just wondering if anyone has any clue?

    opened by crabmon 0
  • how to use  the include_top parameter in pretrained model?

    how to use the include_top parameter in pretrained model?

    I can use the include_top parameter in EfficientNet.from_name, but in EfficientNet.from_pretrained, I cannot use that. So, if I want to change the last layer of the pretrained model by self.base_model = nn.Sequential(*list(base_model.children())[:-1]), it has the error in forward function.

    opened by lck666666 0
  • Can not script the model with torchscript.

    Can not script the model with torchscript.

    I'm trying to script the EfficientNet module with torchscript but there's an error while MemoryEfficientSwish seems to be not supported with torchscript. Have anyone successfully scripted that module yet?

    opened by PhanThanhTrung 0
  • Changing to num class to train to smaller than default

    Changing to num class to train to smaller than default

    Using the pretrained model, encounter the error when i change the num of class to 3 from default 1000. Or is it not meant to change class if using the pretrained model. Pls advise. thanks

    referring to issue #152

    model = EfficientNet.from_pretrained(args.arch, advprop=args.advprop, num_classes=3)

    File "train.py", line 664, in accuracy _, pred = output.topk(maxk, 1, True, True) RuntimeError: selected index k out of range

    opened by lchunleo 0
Owner
Luke Melas-Kyriazi
I'm student at Harvard University studying mathematics and computer science, always open to collaborate on interesting projects!
Luke Melas-Kyriazi
Unofficial PyTorch implementation of DeepMind's Perceiver IO with PyTorch Lightning scripts for distributed training

Unofficial PyTorch implementation of DeepMind's Perceiver IO with PyTorch Lightning scripts for distributed training

Martin Krasser 251 Dec 25, 2022
Tez is a super-simple and lightweight Trainer for PyTorch. It also comes with many utils that you can use to tackle over 90% of deep learning projects in PyTorch.

Tez: a simple pytorch trainer NOTE: Currently, we are not accepting any pull requests! All PRs will be closed. If you want a feature or something does

abhishek thakur 1.1k Jan 4, 2023
null 270 Dec 24, 2022
A lightweight wrapper for PyTorch that provides a simple declarative API for context switching between devices, distributed modes, mixed-precision, and PyTorch extensions.

A lightweight wrapper for PyTorch that provides a simple declarative API for context switching between devices, distributed modes, mixed-precision, and PyTorch extensions.

Fidelity Investments 56 Sep 13, 2022
A PyTorch repo for data loading and utilities to be shared by the PyTorch domain libraries.

A PyTorch repo for data loading and utilities to be shared by the PyTorch domain libraries.

null 878 Dec 30, 2022
PyTorch framework A simple and complete framework for PyTorch, providing a variety of data loading and simple task solutions that are easy to extend and migrate

PyTorch framework A simple and complete framework for PyTorch, providing a variety of data loading and simple task solutions that are easy to extend and migrate

Cong Cai 12 Dec 19, 2021
PyTorch implementation of TabNet paper : https://arxiv.org/pdf/1908.07442.pdf

README TabNet : Attentive Interpretable Tabular Learning This is a pyTorch implementation of Tabnet (Arik, S. O., & Pfister, T. (2019). TabNet: Attent

DreamQuark 2k Dec 27, 2022
An implementation of Performer, a linear attention-based transformer, in Pytorch

Performer - Pytorch An implementation of Performer, a linear attention-based transformer variant with a Fast Attention Via positive Orthogonal Random

Phil Wang 900 Dec 22, 2022
PyTorch implementation of Glow, Generative Flow with Invertible 1x1 Convolutions

glow-pytorch PyTorch implementation of Glow, Generative Flow with Invertible 1x1 Convolutions

Kim Seonghyeon 433 Dec 27, 2022
This is an differentiable pytorch implementation of SIFT patch descriptor.

This is an differentiable pytorch implementation of SIFT patch descriptor. It is very slow for describing one patch, but quite fast for batch. It can

Dmytro Mishkin 150 Dec 24, 2022
GPU-accelerated PyTorch implementation of Zero-shot User Intent Detection via Capsule Neural Networks

GPU-accelerated PyTorch implementation of Zero-shot User Intent Detection via Capsule Neural Networks This repository implements a capsule model Inten

Joel Huang 15 Dec 24, 2022
Tacotron 2 - PyTorch implementation with faster-than-realtime inference

Tacotron 2 (without wavenet) PyTorch implementation of Natural TTS Synthesis By Conditioning Wavenet On Mel Spectrogram Predictions. This implementati

NVIDIA Corporation 4.1k Jan 3, 2023
A Pytorch Implementation for Compact Bilinear Pooling.

CompactBilinearPooling-Pytorch A Pytorch Implementation for Compact Bilinear Pooling. Adapted from tensorflow_compact_bilinear_pooling Prerequisites I

null 169 Dec 23, 2022
A pure Python implementation of Compact Bilinear Pooling and Count Sketch for PyTorch.

Compact Bilinear Pooling for PyTorch. This repository has a pure Python implementation of Compact Bilinear Pooling and Count Sketch for PyTorch. This

Grégoire Payen de La Garanderie 234 Dec 7, 2022
Pytorch implementation of Distributed Proximal Policy Optimization

Pytorch-DPPO Pytorch implementation of Distributed Proximal Policy Optimization: https://arxiv.org/abs/1707.02286 Using PPO with clip loss (from https

Alexis David Jacq 164 Jan 5, 2023
A PyTorch implementation of L-BFGS.

PyTorch-LBFGS: A PyTorch Implementation of L-BFGS Authors: Hao-Jun Michael Shi (Northwestern University) and Dheevatsa Mudigere (Facebook) What is it?

Hao-Jun Michael Shi 478 Dec 27, 2022
A PyTorch implementation of Learning to learn by gradient descent by gradient descent

Intro PyTorch implementation of Learning to learn by gradient descent by gradient descent. Run python main.py TODO Initial implementation Toy data LST

Ilya Kostrikov 300 Dec 11, 2022
PyTorch Implementation of [1611.06440] Pruning Convolutional Neural Networks for Resource Efficient Inference

PyTorch implementation of [1611.06440 Pruning Convolutional Neural Networks for Resource Efficient Inference] This demonstrates pruning a VGG16 based

Jacob Gildenblat 836 Dec 26, 2022
Pretrained ConvNets for pytorch: NASNet, ResNeXt, ResNet, InceptionV4, InceptionResnetV2, Xception, DPN, etc.

Pretrained models for Pytorch (Work in progress) The goal of this repo is: to help to reproduce research papers results (transfer learning setups for

Remi 8.7k Dec 31, 2022