[Feature Request] Add EfficientNet

[timonbimon]

https://ai.googleblog.com/2019/05/efficientnet-improving-accuracy-and.html

It seems like EfficientNet is faster and better than most the other nets out there currently. Would be awesome to see it in torchvision soon! :)

[soumith]

fyi it's available in https://github.com/rwightman/pytorch-image-models/blob/master/models/gen_efficientnet.py already

[cdluminate]

Would be awesome if it's available in torchvision.

[fmassa]

It would also be great to have a training script that reproduces EfficientNet, from what I saw in the paper they also use some different training hyperparameters / tricks to improve the trained models, so it might take some effort to have reproducible training scripts for those.

[bhack]

https://github.com/lukemelas/EfficientNet-PyTorch/blob/master/README.md /cc @lukemelas

[lukemelas]

Thanks for looping me in, @bhack. Yes, it would be great to merge EfficientNet into torchvision!

I don't think we're quite ready yet, though -- as @fmassa mentioned, we have to first have reliable training scripts. I would also like to rewrite EfficientNet in a more PyTorchic style; my implementation and the others I have seen generally mirror the TensorFlow implementation, which is great for quickly reproducing results, but could be improved for merging into torchvision.

[bhack]

/cc @rwightman @zsef123

[rwightman]

@bhack thanks for the vote, definitely think a variant of it should be here, but it should be 100% PyTorch trained. The TF ported weights with the necessary padding are not suitable as it's an unacceptable overhead.

Given the nature of these models, it takes a long time to train, and is fiddly to find the right hyper-params. I finally just got some MobileNet-V3 weights (very close to an EfficientNet-B0) that match/slightly exceed the paper after two weeks of attempts. I can do it with my training scripts, but it takes much more time than training something like a ResNet to great results.

[fmassa]

@rwightman do you have the scripts for training MobileNet-V3 from scratch in PyTorch?

The MobileNetV2 pre-trained weights that I made available also had to be trained on a lot of epochs, and it indeed seems to take longer to train than a ResNet.

[rwightman]

@fmassa TLDR 'sort of' :)

My train script can handle most of the training right now, basically using RMSprop with Google's crazy high epsilons for opt and batch norm and different BN momentum for 350-450+ epochs. That gets you most of the way, but the final necessary bit is EMA decay of the model weights, which PyTorch doesn't have a built-in default for. I have a hack for that working right now, and found it is very effective to enable it at the late stages of training, but validating the smoothed copy of the weights doesn't play nicely with distributed GPU training, so it adds an extra CPU only validation step. Trying to see if there is a cleaner way of doing it...

[rwightman]

I also suspect that odd optimizer/bn params could be relaxed back to PyTorch defaults, and that the EMA weight decay is more significant, as it alone adds a solid 1.5-2% top-1 bump once the model is getting close to converged.

[lukemelas]

Agreed on all counts with @rwightman. It's important to have a fully PyTorchic model and training procedure before merging into torchvision. I'm also finding the training hyperparameters to be very sensitive, but we should get them right before officially releasing.

EMA weight decay ... alone adds a solid 1.5-2% top-1 bump once the model is getting close to converged.

I have not done extensive experimentation yet, but from early impressions, it seems to me as well that weight EMA could make quite a significant difference. @rwightman Do you think feel could be something worth trying to build into a core optimizer at some point in the future?

[rwightman]

It's not torchvision, but it's pretty convenient with the spiffy new hub feature...

torch.hub.list('rwightman/gen-efficientnet-pytorch', force_reload=True)  
model = torch.hub.load('rwightman/gen-efficientnet-pytorch', 'efficientnet_b0', pretrained=True)

That model is native PyTorch, no TF padding hacks, tested to be fully exportable to ONNX/Caffe2. 76.9% top1, 93.2% top5. It's going be awhile before any of the larger ones are satisfactory. Heh.

If anyone wants to take a crack at B1-B3, the model weight EMA smoothing impl I've been using to train, on a dev branch in https://github.com/rwightman/pytorch-image-models/tree/ema-cleanup