def efficientnet(variant, pretrained=False):
if variant is None:
print(
"please specify an EfficientNet variant (b0-b7). Using b0 as default."
)
model = torch.hub.load('rwightman/gen-efficientnet-pytorch',
'efficientnet_b0',
pretrained=pretrained)
return model
# return geffnet.efficientnet_b0(pretrained=pretrained, drop_rate=0.25, drop_connect_rate=0.2, as_sequential=True)
if variant == 0:
return geffnet.efficientnet_b0(pretrained=pretrained,
as_sequential=True)
if variant == 1:
return geffnet.efficientnet_b1(pretrained=pretrained,
as_sequential=True)
elif variant == 2:
return geffnet.efficientnet_b2(pretrained=pretrained,
as_sequential=True)
elif variant == 3:
return geffnet.efficientnet_b3(pretrained=pretrained,
as_sequential=True)
elif variant == 4:
return geffnet.efficientnet_b4(pretrained=pretrained,
as_sequential=True)
elif variant == 5:
return geffnet.efficientnet_b5(pretrained=pretrained,
as_sequential=True)
elif variant == 6:
return geffnet.efficientnet_b6(pretrained=pretrained,
as_sequential=True)
elif variant == 7:
return geffnet.efficientnet_b7(pretrained=pretrained,
as_sequential=True)
def define_model(self):
return geffnet.efficientnet_b0(pretrained=False, drop_rate=0.25, drop_connect_rate=0.2).cuda()
def efficientnet_b0(pretrained=True): return geffnet.efficientnet_b0(as_sequential=True, pretrained=pretrained) def efficientnet_b1(pretrained=True): return geffnet.efficientnet_b1(as_sequential=True, pretrained=pretrained)
import geffnet
import torch
if __name__ == '__main__':
model = geffnet.efficientnet_b0(pretrained=True)
feature_extractor = torch.nn.Sequential(*list(model.children())[:-1])
torch.save(feature_extractor, 'efficientnet_b0_fe.pth.tar')
