A PyTorch 1.0 Implementation of Unet with EfficientNet as encoder
Install efficientunet-pytorch
:
pip install efficientunet-pytorch
e.g. say you want a pretrained efficientnet-b5 model with 5 classes:
from efficientunet import *
model = EfficientNet.from_name('efficientnet-b5', n_classes=5, pretrained=True)
If you prefer to use a model with a custom head rather than just a simple change of the
output_channels
of the last fully-connected layer, use:
from efficientunet import *
model = EfficientNet.custom_head('efficientnet-b5', n_classes=5, pretrained=True)
The structure of model with custom head:
encoder
->concatenation of [AvgPool2d, MaxPool2d]
->Flatten
->Dropout
->Linear(512)
->ReLU
->Dropout
->Linear(n_classes)
e.g. say you want a pretrained efficientunet-b0 model with 2 output channels:
from efficientunet import *
b0unet = get_efficientunet_b0(out_channels=2, concat_input=True, pretrained=True)
The pretrained weights are directly borrowed from this repo.
Version | Tag | Published |
---|---|---|
0.0.6 | 2yrs ago | |
0.0.5 | 3yrs ago | |
0.0.4 | 3yrs ago | |
0.0.3 | 3yrs ago |