def __init__(self, compound_coef):
super(EfficientNetFeatures, self).__init__()
model = EfficientNet.from_pretrained(f'efficientnet-b{compound_coef}',
False)
del model._conv_head
del model._bn1
del model._avg_pooling
del model._dropout
del model._fc
self.model = model
def model_from_dataset(dataset, **kwargs):
if dataset == 'adult':
return FullyConnected(**kwargs)
elif dataset == 'credit':
return FullyConnected(**kwargs)
elif dataset == 'compass':
return FullyConnected(**kwargs)
elif dataset == 'multi_mnist' or dataset == 'multi_fashion_mnist' or dataset == 'multi_fashion':
return MultiLeNet(**kwargs)
elif dataset == 'celeba':
if 'efficientnet' in kwargs['model_name']:
return EfficientNet.from_pretrained(**kwargs)
elif kwargs['model_name'] == 'resnet18':
return ResNet.from_name(**kwargs)
else:
raise ValueError("Unknown model name {}".format(dataset))
def initialize_model(model_name,
num_classes,
learning_rate,
use_pretrained=True):
# Initialize these variables which will be set in this if statement. Each of these
# variables is model specific.
model = None
input_size = 0
optimizer = None
criterion = None
if model_name == "resnet101":
""" Resnet101
"""
model = models.resnet101(pretrained=use_pretrained)
num_ftrs = model.fc.in_features
model.fc = nn.Linear(num_ftrs, num_classes)
input_size = 224
params_to_update = model.parameters()
optimizer = optim.Adam(params_to_update, lr=learning_rate)
criterion = nn.CrossEntropyLoss()
elif model_name == "resnet152":
""" Resnet101
"""
model = models.resnet152(pretrained=use_pretrained)
num_ftrs = model.fc.in_features
model.fc = nn.Linear(num_ftrs, num_classes)
input_size = 224
params_to_update = model.parameters()
optimizer = optim.Adam(params_to_update, lr=learning_rate)
criterion = nn.CrossEntropyLoss()
elif model_name == "resnet101-next":
""" Resnet101
"""
model = models.resnext101_32x8d(pretrained=use_pretrained)
num_ftrs = model.fc.in_features
model.fc = nn.Linear(num_ftrs, num_classes)
input_size = 224
params_to_update = model.parameters()
optimizer = optim.Adam(params_to_update, lr=learning_rate)
criterion = nn.CrossEntropyLoss()
elif model_name == "resnet50-next":
""" Resnet101
"""
model = models.resnext50_32x4d(pretrained=use_pretrained)
num_ftrs = model.fc.in_features
model.fc = nn.Linear(num_ftrs, num_classes)
input_size = 224
params_to_update = model.parameters()
optimizer = optim.Adam(params_to_update, lr=learning_rate)
criterion = nn.CrossEntropyLoss()
elif model_name == "wide-resnet101-2":
""" Resnet101
"""
model = models.wide_resnet101_2(pretrained=use_pretrained)
num_ftrs = model.fc.in_features
model.fc = nn.Linear(num_ftrs, num_classes)
input_size = 224
params_to_update = model.parameters()
optimizer = optim.Adam(params_to_update, lr=learning_rate)
criterion = nn.CrossEntropyLoss()
elif model_name == "densenet161":
""" Densenet
"""
model = models.densenet161(pretrained=use_pretrained)
num_ftrs = model.classifier.in_features
model.classifier = nn.Linear(num_ftrs, num_classes)
input_size = 224
params_to_update = model.parameters()
optimizer = optim.Adam(params_to_update, lr=learning_rate)
criterion = nn.CrossEntropyLoss()
elif model_name == "densenet201":
""" Densenet
"""
model = models.densenet201(pretrained=use_pretrained)
num_ftrs = model.classifier.in_features
model.classifier = nn.Linear(num_ftrs, num_classes)
input_size = 224
params_to_update = model.parameters()
optimizer = optim.Adam(params_to_update, lr=learning_rate)
criterion = nn.CrossEntropyLoss()
elif model_name == "mobilenet":
""" MobileNet
"""
model = models.mobilenet_v2(pretrained=use_pretrained)
# num_ftrs = model.classifier.in_features
# model.classifier = nn.Linear(num_ftrs, num_classes)
model.classifier[1] = nn.Linear(model.last_channel, num_classes)
input_size = 224
params_to_update = model.parameters()
optimizer = optim.Adam(params_to_update, lr=learning_rate)
criterion = nn.CrossEntropyLoss()
elif model_name == "inception":
""" Inception v3
Be careful, expects (299,299) sized images and has auxiliary output
"""
model = models.inception_v3(pretrained=use_pretrained)
# Handle the auxilary net
num_ftrs = model.AuxLogits.fc.in_features
model.AuxLogits.fc = nn.Linear(num_ftrs, num_classes)
# Handle the primary net
num_ftrs = model.fc.in_features
model.fc = nn.Linear(num_ftrs, num_classes)
input_size = 299
params_to_update = model.parameters()
optimizer = optim.Adam(params_to_update, lr=learning_rate)
criterion = nn.CrossEntropyLoss()
elif model_name == "efficientnet-b3":
model = EfficientNet.from_pretrained(model_name)
in_ftrs = model._fc.in_features
model._fc = nn.Linear(in_ftrs, num_classes)
input_size = 224
params_to_update = model.parameters()
optimizer = optim.RMSprop(params_to_update, lr=learning_rate)
criterion = nn.CrossEntropyLoss()
else:
print("Invalid model name, exiting...")
exit()
return model, optimizer, criterion, input_size
