def build_model(model_name): device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") # load pretrained model model_name = model_name # could be fbresnet152 or inceptionresnetv2 if (model_name == 'senet154'): model = pretrainedmodels.senet154(pretrained='imagenet') elif (model_name == 'se_resnet152'): model = pretrainedmodels.se_resnet152(pretrained='imagenet') elif (model_name == 'se_resnext101_32x4d'): model = pretrainedmodels.se_resnext101_32x4d(pretrained='imagenet') elif (model_name == 'resnet152'): model = pretrainedmodels.resnet152(pretrained='imagenet') elif (model_name == 'resnet101'): model = pretrainedmodels.resnet101(pretrained='imagenet') elif (model_name == 'densenet201'): model = pretrainedmodels.densenet201(pretrained='imagenet') model.to(device) for param in model.parameters(): param.requires_grad = False num_ftrs = model.last_linear.in_features class CustomModel(nn.Module): def __init__(self, model): super(CustomModel, self).__init__() self.features = nn.Sequential(*list(model.children())[:-1]) self.classifier = nn.Sequential( torch.nn.Linear(num_ftrs, 128), torch.nn.Dropout(0.3), # drop 50% of the neuron torch.nn.Linear(128, 7)) def forward(self, x): x = self.features(x) x = x.view(x.size(0), -1) x = self.classifier(x) return x model = CustomModel(model) freeze_layer(model.features) num_ftrs = list(model.classifier.children())[-1].out_features model.to(device) model.name = model_name PATH = os.path.abspath(os.path.dirname(__file__)) PATH_par = os.path.abspath(os.path.join(PATH, os.pardir)) path_to_model = os.path.join(PATH_par, 'pretrained_model', '128_7') model.load_state_dict( torch.load(os.path.join(path_to_model, '%s.pth' % (model_name)))) model.to(device) for param in model.parameters(): param.requires_grad = False return model, num_ftrs
def densenet201(input_size=(3, 224, 224), num_classes=1000, pretrained=None): model = models.densenet201(num_classes=num_classes, pretrained=pretrained) model = add_instances_to_torchvisionmodel(model) # Change the First Convol2D layer into new input shape if input_size != (3, 224, 224): model.features[0] = nn.Conv2d(input_size[0], 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False) model.input_size = input_size # calc kernel_size on new_avgpool2d layer test_tensor = torch.randn((1, input_size[0], input_size[1], input_size[2])) features = model.features(test_tensor) #print(features, features.shape[2], features.shape[3]) avg_pool2d_kernel_size = (features.shape[2], features.shape[3]) # calc last linear size x = F.avg_pool2d(features, kernel_size=avg_pool2d_kernel_size, stride=1) x = x.view(x.size(0), -1).shape[1] model.last_linear = nn.Linear(in_features=x, out_features=num_classes) #del model.logits #del model.forward def logits(self, features): x = F.relu(features, inplace=True) x = F.avg_pool2d(x, kernel_size=avg_pool2d_kernel_size, stride=1) x = x.view(x.size(0), -1) x = self.last_linear(x) return x def forward(self, input): x = self.features(input) x = self.logits(x) return x # Modify methods model.logits = types.MethodType(logits, model) model.forward = types.MethodType(forward, model) return model
def __init__(self, class_names): super(DenseNet201, self).__init__() self.base = pretrainedmodels.densenet201() self.backbone = nn.Sequential(*list(self.base.children())[:-1]) self.head = nn.Linear(self.base.last_linear.in_features, len(class_names))
def build_model(model_name): device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") # load pretrained model model_name = model_name # could be fbresnet152 or inceptionresnetv2 if(model_name == 'senet154'): model = pretrainedmodels.senet154(pretrained='imagenet') elif(model_name == 'se_resnet152'): model = pretrainedmodels.se_resnet152(pretrained='imagenet') elif(model_name == 'se_resnext101_32x4d'): model = pretrainedmodels.se_resnext101_32x4d(pretrained='imagenet') elif(model_name == 'resnet152'): model = pretrainedmodels.resnet152(pretrained='imagenet') elif(model_name == 'resnet101'): model = pretrainedmodels.resnet101(pretrained='imagenet') elif(model_name == 'densenet201'): model = pretrainedmodels.densenet201(pretrained='imagenet') model.to(device) for param in model.parameters(): param.requires_grad = False num_ftrs = model.last_linear.in_features class CustomModel(nn.Module): def __init__(self, model): super(CustomModel, self).__init__() self.features = nn.Sequential(*list(model.children())[:-1] ) self.classifier = nn.Sequential( torch.nn.Linear(num_ftrs, 128), torch.nn.Dropout(0.3), # drop 50% of the neuron torch.nn.Linear(128, 7) ) def forward(self, x): x = self.features(x) x = x.view(x.size(0), -1) x = self.classifier(x) return x model = CustomModel(model) freeze_layer(model.features) model.to(device) for param in model.parameters(): param.requires_grad = False class CustomModel1(nn.Module): def __init__(self, model): super(CustomModel1, self).__init__() self.features = nn.Sequential(*list(model.children())[:-1]) self.classifier = nn.Sequential( *[list(model.classifier.children())[i] for i in [0]] ) def forward(self, x): x = self.features(x) x = x.view(x.size(0), -1) x = self.classifier(x) return x CustomModel = CustomModel1(model) num_ftrs = list(CustomModel.classifier.children())[-1].out_features CustomModel.to(device) return CustomModel, num_ftrs