EfficientNet implementation in Julia.
]add https://github.com/pxl-th/EfficientNet.jl.git
model = EffNet("efficientnet-b0"; n_classes=10, in_channels=3)model = EfficientNet.from_pretrained("efficientnet-b3")By default, weights are stored in ~/.cache/EfficientNet.jl/ directory.
If there are no weights, it will attempt to download them.
Additionally, you can specify cache directory as a second parameter to from_pretrained function.
Available pretrained models are B0-B7. They are loaded from EfficientNet-PyTorch and converted to Julia's format.
using Images
using EfficientNet
device = gpu
model = EfficientNet.from_pretrained("efficientnet-b3")
model = model |> testmode! |> device
image = "./spaceshuttle.png"
x = Images.load(image) |> channelview .|> Float32
x = Flux.unsqueeze(permutedims(x, (3, 2, 1)), 4)
o = x |> device |> model |> softmax |> cpu
o = sortperm(o[:, 1])
@info "Top 5 classes: $(o[end:-1:end - 5] .- 1)"To extract list of features, pass Val(:stages) as the second parameter:
features = model(x, Val(:stages))It will contain features, extracted from different resolution levels.
This can be used in something like UNet architecture as an encoder.
To get those resolution levels, use model.stages and model.stages_channels
to get ids of specific layers and output channels.