added include_top option to model

[nwschurink]

Hi there,

I've added extra functionality such that a user can now specify whether to include the last layers of the network (avgpool + dropout + fc + swish) or to omit these last layers instead. This is especially handy if one wants to use the efficientnet just for feature extraction and wants to build a different network on top of it.

A simple minimum working example is given below:

from efficientnet_pytorch_3d import EfficientNet3D
from efficientnet_pytorch_3d.utils import MemoryEfficientSwish
from torchsummary import summary
from torch import nn

# A model with the final layers
model = EfficientNet3D.from_name("efficientnet-b0", override_params={'num_classes':2, 'include_top':True}, in_channels=1)
model.to('cuda')
summary(model,input_size=(1,100,100,100))

# A model without the final layers
model = EfficientNet3D.from_name("efficientnet-b0", override_params={'num_classes':2, 'include_top':False}, in_channels=1)
model.to('cuda')
summary(model,input_size=(1,100,100,100))

# A custom model build on top of the feature extraction part of EfficientNet
model = EfficientNet3D.from_name("efficientnet-b0", override_params={'num_classes':2, 'include_top':False}, in_channels=1)
new_model = nn.Sequential(model,nn.AdaptiveAvgPool3d(1),nn.Dropout(0.2),nn.Flatten(),nn.Linear(1280,100),nn.Linear(100,2),MemoryEfficientSwish())

new_model.to('cuda')
summary(new_model,input_size=(1,100,100,100))

[nwschurink]

@shijianjian is this a feature you'd like to add?