def get_model(args):
if args.model_name == 'inception':
if args.pre_trained_checkpoint is not None:
model = models.inception_v3(pretrained=False, transform_input=False)
model.fc = nn.Linear(2048, args.n_outputs)
model.AuxLogits = InceptionAux(768, args.n_outputs)
model.aux_logits = False
new_conv = BasicConv2d(1, 32, kernel_size=3, stride=2)
model.Conv2d_1a_3x3 = new_conv
sd = torch.load(args.pre_trained_checkpoint)['model']
model.load_state_dict(sd)
else:
model = models.inception_v3(pretrained=True, transform_input=False)
model.fc = nn.Linear(2048, args.n_outputs)
model.AuxLogits = InceptionAux(768, args.n_outputs)
model.aux_logits = False
new_conv = BasicConv2d(1, 32, kernel_size=3, stride=2)
first_layer_sd = model.Conv2d_1a_3x3.state_dict()
first_layer_sd['conv.weight'] = first_layer_sd['conv.weight'].mean(dim=1, keepdim=True)
new_conv.load_state_dict(first_layer_sd)
model.Conv2d_1a_3x3 = new_conv
elif args.model_name == 'resnet':
if args.pre_trained_checkpoint is not None:
model = models.resnet50(pretrained=False)
model.fc = nn.Linear(in_features=2048, out_features=args.n_outputs, bias=True)
new_conv = nn.Conv2d(1, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)
model.conv1 = new_conv
sd = torch.load(args.pre_trained_checkpoint)['model']
model.load_state_dict(sd)
else:
model = models.resnet50(pretrained=True)
model.fc = nn.Linear(in_features=2048, out_features=args.n_outputs, bias=True)
new_conv = nn.Conv2d(1, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)
first_layer_sd = model.conv1.state_dict()
first_layer_sd['weight'] = first_layer_sd['weight'].mean(dim=1, keepdim=True)
new_conv.load_state_dict(first_layer_sd)
model.conv1 = new_conv
elif args.model_name == 'resnext':
if args.pre_trained_checkpoint is not None:
model = models.resnext50_32x4d(pretrained=False)
model.fc = nn.Linear(in_features=2048, out_features=args.n_outputs, bias=True)
new_conv = nn.Conv2d(1, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)
model.conv1 = new_conv
sd = torch.load(args.pre_trained_checkpoint)['model']
model.load_state_dict(sd)
else:
model = models.resnext50_32x4d(pretrained=True)
model.fc = nn.Linear(in_features=2048, out_features=args.n_outputs, bias=True)
new_conv = nn.Conv2d(1, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)
first_layer_sd = model.conv1.state_dict()
first_layer_sd['weight'] = first_layer_sd['weight'].mean(dim=1, keepdim=True)
new_conv.load_state_dict(first_layer_sd)
model.conv1 = new_conv
return model
def __init__(self, pretrained, output_channel):
super(ResNeXt_50, self).__init__()
if pretrained == True:
self.model = models.resnext50_32x4d(pretrained=True, progress=True)
else:
self.model = models.resnext50_32x4d(pretrained=False,
progress=True)
self.l0 = nn.Linear(1000, output_channel)
self.dropout = nn.Dropout(p=0.5)
self.fc1 = nn.Linear(2048, 1000)
def __init__(self):
super().__init__()
self.resnext = models.resnext50_32x4d(pretrained=True)
n_features = self.resnext.fc.in_features
self.resnext.fc = nn.Linear(n_features, 5)
def get_model(name, pretrained=True):
if name == 'resnet18':
model = models.resnet18(pretrained=pretrained)
out_channels = [64, 128, 256, 512]
elif name == 'resnet34':
model = models.resnet34(pretrained=pretrained)
out_channels = [64, 128, 256, 512]
elif name == 'resnet50':
model = models.resnet50(pretrained=pretrained)
out_channels = [256, 512, 1024, 2048]
elif name == 'resnet101':
model = models.resnet101(pretrained=pretrained)
out_channels = [256, 512, 1024, 2048]
# ResNeXt
elif name == 'resnext50':
model = models.resnext50_32x4d(pretrained=pretrained)
out_channels = [256, 512, 1024, 2048]
elif name == 'resnext101':
model = models.resnext101_32x8d(pretrained=pretrained)
out_channels = [256, 512, 1024, 2048]
elif name == 'wide_resnet50':
model = models.wide_resnet50_2(pretrained=pretrained)
out_channels = [256, 512, 1024, 2048]
elif name == 'wide_resnet101':
model = models.wide_resnet101_2(pretrained=pretrained)
out_channels = [256, 512, 1024, 2048]
# Error
else:
raise NotImplemented('{} backbone model is not implemented so far.'.format(name))
return model, out_channels
def pretrained_model(name):
if name == "resnet18":
return models.resnet18(pretrained=True)
elif name == "alexnet":
return models.alexnet(pretrained=True)
elif name == "squeezenet":
return models.squeezenet1_0(pretrained=True)
elif name == "vgg16":
return models.vgg16(pretrained=True)
elif name == "densenet":
return models.densenet161(pretrained=True)
elif name == "inception":
return models.inception_v3(pretrained=True)
elif name == "googlenet":
return models.googlenet(pretrained=True)
elif name == "shufflenet":
return models.shufflenet_v2_x1_0(pretrained=True)
elif name == "mobilenet":
return models.mobilenet_v2(pretrained=True)
elif name == "resnext50_32x4d":
return models.resnext50_32x4d(pretrained=True)
elif name == "wide_resnet50_2":
return models.wide_resnet50_2(pretrained=True)
elif name == "mnasnet":
return models.mnasnet1_0(pretrained=True)
else:
print("model {} don't know".format(name))
def __init__(self, classes=2, pt=True):
super(resnext_50_32x4d, self).__init__()
self.base_model = models.resnext50_32x4d(pretrained=pt)
in_features = self.base_model.fc.out_features
#self.nb_features = self.base_model.fc.in_features
self.l0 = nn.Linear(in_features, classes)
