def get_model(arch, nlabels, hidden_units, learning_rate):
if arch == 'vgg16':
model = tmodels.vgg16(pretrained=True)
elif arch == 'vgg13':
model = tmodels.vgg13(pretrained=True)
elif arch == 'alexnet':
model = tmodels.alexnet(pretrained=True)
elif arch == 'resnet18':
model = tmodels.resnet18(pretrained=True)
elif arch == 'squeezenet1_0':
model = tmodels.squeezenet1_0(pretrained=True)
elif arch == 'densenet161':
model = tmodels.densenet161(pretrained=True)
else:
raise ValueError('Unspected network architecture ', arch)
#Freeze Parameters
for param in model.parameters():
param.requires_grad = False
model.classifier = nn.Sequential(nn.Linear(25088, hidden_units), nn.ReLU(),
nn.Dropout(0.2),
nn.Linear(hidden_units, hidden_units),
nn.ReLU(), nn.Dropout(0.3),
nn.Linear(hidden_units, nlabels),
nn.LogSoftmax(dim=1))
criterion = nn.NLLLoss()
optimizer = optim.Adam(model.classifier.parameters(), lr=learning_rate)
return model, model.classifier, criterion, optimizer
def define_C(output_nc, ndf, init_type='normal', gpu_ids=[]):
#if output_nc == 3:
# netC = get_model('DTN', num_cls=10)
#else:
# Exception('classifier only implemented for 32x32x3 images')
#=================================================================================================
num_classes = 31
netC = models.squeezenet1_0(pretrained=True)
set_parameter_requires_grad(netC, False)
netC.classifier[1] = nn.Conv2d(512,
num_classes,
kernel_size=(1, 1),
stride=(1, 1))
netC.num_classes = num_classes
input_size = 224
print('Learnable parameters')
params_to_update = []
for name, param in netC.named_parameters():
if param.requires_grad == True:
params_to_update.append(param)
print("\t", name)
if len(gpu_ids) > 0:
assert (torch.cuda.is_available())
netC.to(gpu_ids[0])
netC = torch.nn.DataParallel(netC, gpu_ids)
return netC
def __init__(self, c_in, model_name):
super(PretrainedNet, self).__init__()
if model_name == 'resnet101':
model_ft = models.resnet101(pretrained=True)
elif model_name == 'densenet121':
model_ft = models.densenet121(pretrained=True)
elif model_name == 'densenet161':
model_ft = models.densenet161(pretrained=True)
elif model_name == 'squeezenet':
model_ft = models.squeezenet1_0(pretrained=True)
elif model_name == 'resnext101_32x8d':
model_ft = models.wide_resnet101_2(pretrained=True)
elif model_name == 'inception_v3':
model_ft = models.inception_v3(pretrained=True)
self._set_parameter_requires_grad(model_ft)
if model_name.startswith('res'):
layers = [i for j, i in enumerate(model_ft.children()) if j < 7]
self.model = nn.Sequential(*layers, nn.MaxPool2d((4, 1)))
elif model_name.startswith('dense'):
layers = model_ft.features[:-3]
self.model = nn.Sequential(*layers, nn.MaxPool2d((4, 1)))
elif model_name.startswith('squeeze'):
layers = model_ft.features[:-2]
self.model = nn.Sequential(*layers, nn.MaxPool2d((4, 1)))
elif model_name.startswith('incep'):
layers = [i for j, i in enumerate(model_ft.children()) if j < 13]
self.model = nn.Sequential(*layers, nn.MaxPool2d((8, 4)))
def net(self, net):
if net == "ResNet50":
Net = models.resnet50(pretrained=self.pretrained)
Net.fc = nn.Linear(2048, self.class_num)
elif net == 'DenseNet121':
Net = models.densenet121(pretrained=True)
Net.fc = nn.Linear(1024, self.class_num)
elif net == 'SENet':
Net = models.squeezenet1_0(pretrained=True)
Net.classifier = nn.Sequential(
nn.Dropout(p=0.5), nn.Conv2d(512,
self.class_num,
kernel_size=1),
nn.ReLU(inplace=True), nn.AdaptiveAvgPool2d((1, 1)))
elif net == 'pnasnet':
model_name = 'pnasnet5large'
if self.pretrained:
Net = pretrainedmodels.__dict__[model_name](
num_classes=1000, pretrained='imagenet')
else:
Net = pretrainedmodels.__dict__[model_name](num_classes=1000)
Net.last_linear = nn.Linear(4320, 2)
Net.eval()
elif net == 'efficientNet':
if self.pretrained:
Net = EfficientNet.from_pretrained('efficientnet-b0')
else:
Net = EfficientNet.from_name('efficientnet-b0')
Net._fc = nn.Linear(1280, self.class_num, bias=True)
Net.eval()
return Net
def initialize_model(model_name,
num_classes,
feature_extract,
use_pretrained=True):
# Initialize these variables which will be set in this if statement. Each of these
# variables is model specific.
model_ft = None
input_size = 0
if model_name == "resnet":
""" Resnet34
"""
model_ft = models.resnet34(pretrained=use_pretrained)
set_parameter_requires_grad(model_ft, feature_extract)
num_ftrs = model_ft.fc.in_features
model_ft.fc = nn.Linear(num_ftrs, num_classes)
input_size = 224
elif model_name == "alexnet":
""" Alexnet
"""
model_ft = models.alexnet(pretrained=use_pretrained)
set_parameter_requires_grad(model_ft, feature_extract)
num_ftrs = model_ft.classifier[6].in_features
model_ft.classifier[6] = nn.Linear(num_ftrs, num_classes)
input_size = 224
elif model_name == "vgg":
""" VGG11_bn
"""
model_ft = models.vgg11_bn(pretrained=use_pretrained)
set_parameter_requires_grad(model_ft, feature_extract)
num_ftrs = model_ft.classifier[6].in_features
model_ft.classifier[6] = nn.Linear(num_ftrs, num_classes)
input_size = 224
elif model_name == "vgg16":
""" VGG16
"""
model_ft = models.vgg16(pretrained=use_pretrained)
set_parameter_requires_grad(model_ft, feature_extract)
num_ftrs = model_ft.classifier[6].in_features
model_ft.classifier[6] = nn.Linear(num_ftrs, num_classes)
input_size = 224
elif model_name == "squeezenet":
""" Squeezenet
"""
model_ft = models.squeezenet1_0(pretrained=use_pretrained)
set_parameter_requires_grad(model_ft, feature_extract)
model_ft.classifier[1] = nn.Conv2d(512,
num_classes,
kernel_size=(1, 1),
stride=(1, 1))
model_ft.num_classes = num_classes
input_size = 224
else:
print("Invalid model name, exiting...")
exit()
return model_ft, input_size
def initialize_model(model_name,
num_classes,
feature_extract,
use_pretrained=True):
model_ft = None
input_size = 0
if model_name == "resnet":
""" Resnet18
"""
model_ft = models.resnet18(pretrained=use_pretrained)
set_parameter_requires_grad(model_ft, feature_extract)
num_ftrs = model_ft.fc.in_features
model_ft.fc = nn.Linear(num_ftrs, num_classes)
input_size = 224
elif model_name == "squeezenet":
""" Squeezenet
"""
model_ft = models.squeezenet1_0(pretrained=use_pretrained)
#set_parameter_requires_grad(model_ft, feature_extract)
model_ft.classifier[1] = nn.Conv2d(512,
num_classes,
kernel_size=(1, 1),
stride=(1, 1))
model_ft.num_classes = num_classes
input_size = 224
return model_ft, input_size
def sn00_predict(fc_img_tensor, fc_label_list):
bs, ncrops, c, h, w = fc_img_tensor.size()
fc_img_tensor = fc_img_tensor.view(-1, c, h, w)
img = Variable(fc_img_tensor)
squeeze = models.squeezenet1_0(pretrained=True)
squeeze._modules['features'][2] = torch.nn.modules.AdaptiveAvgPool2d(
(54, 54))
print('SqueezeNet00 created. Start predicting......')
prediction = squeeze(img)
print('Calculating accuracy......')
prediction = prediction.view(bs, ncrops, -1).mean(1)
pred_np = prediction.detach().numpy()
opt = []
for i in range(pred_np.shape[0]):
temp = pred_np[i].argmax()
opt.append(temp)
correct = 0
for i in range(len(opt)):
if opt[i] == fc_label_list[i]:
correct += 1
acc = float(correct) / len(opt)
return opt, acc
def get_dx(self, datasetdir,model):
# ImageFolder 一个通用的数据加载器
train_dataset = datasets.ImageFolder(
os.path.join(datasetdir, 'getP'),
# 对数据进行预处理
transforms.Compose([
# transforms.RandomResizedCrop(16),
transforms.ToTensor(), # 把一个取值范围是[0,255]或者shape为(H,W,C)的numpy.ndarray，
# 转换成形状为[C,H,W]，取值范围是[0,1.0]的torch.FloadTensor
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])
)
train_loader = torch.utils.data.DataLoader(train_dataset,batch_size=5)
if model=='resnet':
nnmodel = models.resnet18(pretrained=False)
nnmodel.load_state_dict(torch.load('H:/resnet18-epoch100-acc70.316.pth'))
elif model=='alexnet':
nnmodel = models.alexnet(pretrained=True)
elif model=='vgg':
nnmodel = models.vgg11(pretrained=True)
elif model== 'squeezenet':
nnmodel = models.squeezenet1_0(pretrained=True)
nnmodel = torch.nn.DataParallel(nnmodel).cuda()
lossfunc = torch.nn.CrossEntropyLoss()
for i, (imgs,targets ) in enumerate(train_loader):
input_var = torch.autograd.Variable(imgs,requires_grad=True)
target_var = torch.autograd.Variable(targets).cuda()
output = nnmodel(input_var)
loss = lossfunc(output, target_var).cuda()
loss.backward()
dx = input_var.grad
return dx,imgs
def __init__(self):
"""Select conv1_1 ~ conv5_1 activation maps."""
super(SqueezeNet1_0, self).__init__()
self.select_feats = ['maxpool1',
'maxpool4',
'fire8',
'maxpool8',
'fire9', ]
self.select_classifier = ['conv11', 'avgpool13']
self.feat_list = self.select_feats + self.select_classifier
self.sqnet_feats = models.squeezenet1_0(pretrained=True).features
self.sqnet_classifier = models.squeezenet1_0(
pretrained=True).classifier
def initialize_squeezenet(model_name,
num_classes,
feature_extract,
use_pretrained=True):
model_ft = None
model_params_optimized = None
input_size = 0
if model_name == 'squeezenet1_0':
model_ft = models.squeezenet1_0(pretrained=use_pretrained)
elif model_name == 'squeezenet1_1':
model_ft = models.squeezenet1_1(pretrained=use_pretrained)
else:
print(
'{} is not a valid squeeze net in torchvision'.format(model_name))
sys.exit(1)
set_parameter_requires_grad(model_ft, feature_extract)
model_ft.classifier[1] = nn.Conv2d(512,
num_classes,
kernel_size=(1, 1),
stride=(1, 1))
if feature_extract:
model_params_optimized = model_ft.classifier[1].paramters()
else:
model_params_optimized = model_ft.parameters()
input_size = 224
return model_ft, model_params_optimized, input_size
def setup_model(args):
"""Select CNN architecture, modified for transfer learning on
specific dataset."""
# Call model constructor
model = models.squeezenet1_0(pretrained=True)
# Freeze layer parameters if feature extracting
if args.feature_extract:
for param in model.parameters():
param.requires_grad = False
# Modify model to support fewer classes: (512, 1000) -> (512, 2)
# This will also unfreeze this layer's parameters, as the default value for
# (weight/bias).required_grad = True
model.classifier[1] = nn.Conv2d(512, args.num_classes, kernel_size=1)
model.num_classes = args.num_classes
# Send the model to GPU
model = model.to(device)
# Create list of parameters that should be updated (must be done after
# freezing, and after last layer is reshaped).
params_to_update = [p for p in model.parameters()
if p.requires_grad is True]
return model, params_to_update
def build_pretrained_model(model_name):
"""Returns a pre-trained model given its name"""
assert model_name in PRETRAINED_MODELS, "Choose model name in {}".format(
PRETRAINED_MODELS)
if model_name == "alexnet":
model = pretrain_models.alexnet(pretrained=True)
n_fts = model.classifier[6].in_features
model.classifier[6] = nn.Linear(n_fts, 1)
model.classifier.add_module("7", nn.Sigmoid())
elif model_name == "squeezenet":
model = pretrain_models.squeezenet1_0(pretrained=True)
model.classifier[1] = nn.Conv2d(512,
1,
kernel_size=(1, 1),
stride=(1, 1))
model.classifier.add_module("4", nn.Sigmoid())
model.num_classes = 1
elif model_name == "resnet":
model = pretrain_models.resnet18(pretrained=True)
n_fts = model.fc.in_features
model.fc = nn.Sequential(nn.Linear(n_fts, 1), nn.Sigmoid())
elif model_name == "vgg":
model = pretrain_models.vgg11_bn(pretrained=True)
n_fts = model.classifier[6].in_features
model.classifier[6] = nn.Linear(n_fts, 1)
model.classifier.add_module("7", nn.Sigmoid())
return model
def download_model(model_name):
if model_name == "mobilenet_v2":
model = models.mobilenet_v2(pretrained=True).eval()
elif model_name == "resnet18":
model = models.resnet18(pretrained=True).eval()
elif model_name == "resnet34":
model = models.resnet34(pretrained=True).eval()
elif model_name == "resnet50":
model = models.resnet50(pretrained=True).eval()
elif model_name == "resnet101":
model = models.resnet101(pretrained=True).eval()
elif model_name == "resnet152":
model = models.resnet152(pretrained=True).eval()
elif model_name == "shufflenet_v2_x0_5":
model = models.shufflenet_v2_x0_5(pretrained=True).eval()
elif model_name == "shufflenet_v2_x1_0":
model = models.shufflenet_v2_x1_0(pretrained=True).eval()
elif model_name == "squeezenet1_0":
model = models.squeezenet1_0(pretrained=True).eval()
elif model_name == "squeezenet1_1":
model = models.squeezenet1_1(pretrained=True).eval()
elif model_name == "vgg11":
model = models.vgg11(pretrained=True).eval()
elif model_name == "vgg13":
model = models.vgg13(pretrained=True).eval()
elif model_name == "vgg16":
model = models.vgg16(pretrained=True).eval()
elif model_name == "vgg19":
model = models.vgg19(pretrained=True).eval()
elif model_name == "wide_resnet50_2":
model = models.wide_resnet50_2(pretrained=True).eval()
return model
def choose_model(model_name):
# Choose a new pretrained model
model_name = input("Hello! Please choose what torchvision model you want to apply for pre-training the dataset. You can type vgg16 OR densenet OR squeezenet:").lower()
try:
if model_name == 'vgg16':
model = models.vgg16(pretrained=True)
for param in model.parameters():
param.requires_grad = False
elif model_name == 'squeezenet':
model = models.squeezenet1_0(pretrained=True)
for param in model.parameters():
param.requires_grad = False
elif model_name == 'densenet':
model = models.vgg16(pretrained=True)
for param in model.parameters():
param.requires_grad = False
else:
print("Sorry {} is not a valid model name. Let's move on.".format(model_name))
except ValueError as e:
print("Exception occurred: {}".format(e))
print('-'*40)
return model
def load_checkpoint(filepath):
pretrained_models = {
'resnet18': models.resnet18(pretrained=True),
'alexnet': models.alexnet(pretrained=True),
'squeezenet': models.squeezenet1_0(pretrained=True),
'vgg16': models.vgg16(pretrained=True),
'densenet': models.densenet121(pretrained=True),
'inception': models.inception_v3(pretrained=True)
}
checkpoint = torch.load(filepath)
#model = pretrained_models[arch]
input_size = checkpoint['input_size']
output_size = checkpoint['output_size']
hidden_units = 512 #checkpoint['hidden_units']
drpout = 0.2
pretrained_model = pretrained_models[checkpoint['pretrained_model']]
#print('Pretrained model', pretrained_model)
model = build_network(input_size, hidden_units, pretrained_model,
output_size, drpout)
#model.classifier = checkpoint['classifier']
model.load_state_dict(checkpoint['state_dict'])
#print("Trained model")
#print('Trained model',model)
return model
def setModel(self, arch):
print("\nsetting up model...\n")
# alexnet = models.alexnet(pretrained=True)
# squeezenet = models.squeezenet1_0(pretrained=True)
# vgg11 = models.vgg11(pretrained=True)
# vgg13 = models.vgg13(pretrained=True)
# vgg16 = models.vgg16(pretrained=True)
# vgg19 = models.vgg19(pretrained=True)
# densenet = models.densenet161(pretrained=True)
# inception = models.inception_v3(pretrained=True)
# googlenet = models.googlenet(pretrained=False)
# shufflenet = models.shufflenet_v2_x1_0(pretrained=True)
# mobilenet = models.mobilenet_v2(pretrained=True)
# resnext50_32x4d = models.resnext50_32x4d(pretrained=True)
switcher = {
'alexnet': models.alexnet(pretrained=True),
'squeezenet': models.squeezenet1_0(pretrained=True),
'vgg16': models.vgg16(pretrained=True),
'inception': models.inception_v3(pretrained=True),
#'googlenet': models.googlenet(pretrained=False),
#'mobilenet': models.mobilenet_v2(pretrained=True),
#'resnext50_32x4d' : models.resnext50_32x4d(pretrained=True),
'vgg11': models.vgg11(pretrained=True),
'vgg13': models.vgg13(pretrained=True),
'vgg16': models.vgg16(pretrained=True),
}
error = "\nThat model is not supported yet. The supported models are : 'alexnet', squeezenet',\
'vgg11', 'vgg13', 'vgg16', 'vgg19', 'inception', 'googlenet', 'mobilenet', 'resnext50_32x4d' "
self.model = switcher.get(arch, error)
if self.model == error:
print(error)
else:
print('\nmodel successfully set to {}'.format(arch))
def initialize_model(model_name,
num_classes,
feature_extract,
use_pretrained=True):
# Initialize these variables which will be set in this if statement. Each of these, variables is model specific.
model_ft = None
input_size = 0
if model_name == 'resnet':
model_ft = models.resnet18(pretrained=use_pretrained)
set_parameter_grad(model_ft, feature_extract)
num_ftrs = model_ft.fc.in_features
model_ft.fc = nn.Linear(num_ftrs, num_classes)
input_size = 224
elif model_name == 'alexnet':
model_ft = models.alexnet(pretrained=use_pretrained)
set_parameter_grad(model_ft, feature_extract)
num_ftrs = model_ft.classifier[6].in_features
model_ft.classifier[6] = nn.Linear(num_ftrs, num_classes)
input_size = 224
elif model_name == 'vgg':
model_ft = models.vgg11_bn(pretrained=use_pretrained)
set_parameter_grad(model_ft, feature_extract)
num_ftrs = model_ft.classifier[6].in_features
model_ft.classifier[6] = nn.Linear(num_ftrs, num_classes)
input_size = 224
elif model_name == 'squeezenet':
model_ft = models.squeezenet1_0(pretrained=use_pretrained)
set_parameter_grad(model_ft, feature_extract)
model_ft.classifier[1] = nn.Conv2d(512,
num_classes,
kernel_size=(1, 1),
stride=(1, 1))
model_ft.num_classes = num_classes
input_size = 224
elif model_name == 'densenet':
model_ft = models.densenet121(pretrained=use_pretrained)
set_parameter_grad(model_ft, feature_extract)
num_ftrs = model_ft.classifier.in_features
model_ft.classifier = nn.Linear(num_ftrs, num_classes)
input_size = 224
elif model_name == 'inception':
model_ft = models.inception_v3(pretrained=use_pretrained)
set_parameter_grad(model_ft, feature_extract)
num_ftrs = model_ft.AuxLogits.fc.in_features
model_ft.AuxLogits.fc = nn.Linear(num_ftrs, num_classes)
num_ftrs = model_ft.fc.in_features
model_ft.fc = nn.Linear(num_ftrs, num_classes)
input_size = 299
else:
print('Invalid model name, exiting......')
exit()
return model_ft, input_size
def predict_class(imgPath, model):
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
img = Image.open(imgPath)
trans = transforms.Compose([
transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])
image = trans(img)
image.unsqueeze_(0)
image = image.to(device)
nnmodel = models.resnet18(pretrained=False)
if model == 'resnet':
nnmodel = models.resnet18(pretrained=False)
nnmodel.load_state_dict(
torch.load('dataset/resnet18-epoch100-acc70.316.pth'))
elif model == 'alexnet':
nnmodel = models.alexnet(pretrained=False)
nnmodel.load_state_dict(torch.load('dataset/alexnet-owt-4df8aa71.pth'))
elif model == 'vgg':
nnmodel = models.vgg11(pretrained=False)
nnmodel.load_state_dict(torch.load('dataset/vgg11-bbd30ac9.pth'))
elif model == 'squeezenet':
nnmodel = models.squeezenet1_0(pretrained=False)
nnmodel.load_state_dict(
torch.load('dataset/squeezenet1_0-a815701f.pth'))
nnmodel = torch.nn.DataParallel(nnmodel).to(device)
nnmodel.eval()
out = nnmodel(image)
_, indices = torch.sort(out, descending=True)
result = indices.detach().cpu().numpy()
return result
def squeezenet1_0(num_classes,
freeze_feature_extractor=False,
use_pretrained=True,
num_channels=3):
"""Wrapper for squeezenet architecture
input size = 224
Args:
num_classes:
number of output nodes for the final layer
freeze_feature_extractor:
if False (default), entire network will have gradients and can train
if True, feature block is frozen and only final layer is trained
use_pretrained:
if True, uses pre-trained ImageNet features from
Pytorch's model zoo.
num_channels:
specify channels in input sample, eg [channels h,w] sample shape
"""
model_ft = models.squeezenet1_0(pretrained=use_pretrained)
if freeze_feature_extractor:
freeze_params(model_ft)
model_ft.classifier[1] = nn.Conv2d(512,
num_classes,
kernel_size=(1, 1),
stride=(1, 1))
model_ft.num_classes = num_classes
# change input shape num_channels
if num_channels != 3:
model_ft.features[0] = nn.Conv2d(num_channels,
96,
kernel_size=7,
stride=2)
return model_ft
def get_classifier(classifier, pretrained, dataset='miniimagenet'):
if ((dataset == 'miniimagenet') or (dataset == 'miniimagenetgenerated')):
if classifier == 'vgg16_bn':
return models.vgg16(pretrained=pretrained)
elif classifier == 'alexnet':
return models.alexnet(pretrained=pretrained)
elif classifier == 'mnasnet1':
return models.mnasnet1_0(pretrained=pretrained)
elif classifier == 'resnet18':
return models.resnet18(pretrained=pretrained)
elif classifier == 'shufflenet':
return models.shufflenet_v2_x1_0(pretrained=pretrained)
elif classifier == 'resnet50':
return models.resnext50_32x4d(pretrained=pretrained)
elif classifier == 'squeezenet':
return models.squeezenet1_0(pretrained=pretrained)
elif classifier == 'densenet161':
return models.densenet161(pretrained=pretrained)
elif classifier == 'wide_resnet50_2':
model = models.wide_resnet50_2(pretrained=False)
if pretrained:
state_dict = torch.load(os.path.join(
get_project_root(),
'classifier/models/wide_resnet50_2.pt'),
map_location=torch.device('cuda'))
model.load_state_dict(state_dict)
return model
elif classifier == 'mobilenet_v2':
return models.mobilenet_v2(pretrained=pretrained)
elif classifier == 'googlenet':
return models.googlenet(pretrained=pretrained)
elif classifier == 'inception_v3':
return models.inception_v3(pretrained=pretrained)
else:
raise NameError('Please enter a valid classifier')
def getNetwork(args):
if (args.net_type == 'alexnet'):
net = models.alexnet(pretrained=args.finetune)
file_name = 'alexnet'
elif (args.net_type == 'vggnet'):
if (args.depth == 11):
net = models.vgg11(pretrained=args.finetune)
elif (args.depth == 13):
net = models.vgg13(pretrained=args.finetune)
elif (args.depth == 16):
net = models.vgg16(pretrained=args.finetune)
elif (args.depth == 19):
net = models.vgg19(pretrained=args.finetune)
else:
print(
'Error : VGGnet should have depth of either [11, 13, 16, 19]')
sys.exit(1)
file_name = 'vgg-%s' % (args.depth)
elif (args.net_type == 'squeezenet'):
net = models.squeezenet1_0(pretrained=args.finetune)
file_name = 'squeeze'
elif (args.net_type == 'resnet'):
net = resnet(args.finetune, args.depth)
file_name = 'resnet-%s' % (args.depth)
else:
print(
'Error : Network should be either [alexnet / squeezenet / vggnet / resnet]'
)
sys.exit(1)
return net, file_name
def load_checkpoint_general(filepath):
checkpoint = torch.load(filepath,
map_location=lambda storage, loc: storage)
arch = checkpoint['arch']
arch = arch.lower()
#print(arch)
if arch == "densenet121":
model = models.densenet121(pretrained=True)
elif arch == "resnet18":
model = models.resnet18(pretrained=True)
elif arch == "alexnet":
model = models.alexnet(pretrained=True)
elif arch == "squeezenet":
model = models.squeezenet1_0(pretrained=True)
elif arch == "vgg16":
model = models.vgg16(pretrained=True)
elif arch == "densenet161":
model = models.densenet161(pretrained=True)
elif arch == "inception":
model = models.inception_v3(pretrained=True)
else:
return False
model.classifier = utils.create_classifer(checkpoint['network'])
model.load_state_dict(checkpoint['state_dict'])
model.class_to_idx = checkpoint['class_to_idx']
return model
def execute(modelName, fileName):
# Creo un modello per caricare il checkpoint
if modelName == 'SqueezeNet':
model = squeezenet1_0()
model.classifier[1] = nn.Conv2d(512,
4,
kernel_size=(1, 1),
stride=(1, 1))
model.num_classes = 4
else:
model = vgg16()
model.classifier[6] = nn.Linear(4096, 4)
model.load_state_dict(
torch.load('checkpoint\\' + modelName +
'_checkpoint.pth')['state_dict'])
# Preprocessing dell'immagine di input
im = imresize(Image.open(fileName))
im = get_transform(im)
batch_t = torch.unsqueeze(im, 0)
model.eval()
out_predict = model(batch_t)
_, index = torch.max(out_predict, 1)
labels = ['Non premuto', 'Bottone A', 'Bottone B', 'Bottone C']
return labels[index[0]]
def set_up_model(self):
if self.opt.model_architecture == 'ResNet34':
self.net = models.resnet34(pretrained=True)
self.net.fc = nn.Linear(512, self.opt.num_classes)
elif self.opt.model_architecture == 'ResNet18':
self.net = models.resnet18(pretrained=True)
self.net.fc = nn.Linear(512, self.opt.num_classes)
elif self.opt.model_architecture == 'DesNet121':
self.net = models.densenet121(pretrained=True)
self.net.fc = nn.Linear(1024, self.opt.num_classes)
elif self.opt.model_architecture == 'inception_v3':
self.net = models.inception_v3(pretrained=True)
self.net.AuxLogits.fc = nn.Linear(768, self.opt.num_classes)
self.net.fc = nn.Linear(2048, self.opt.num_classes)
elif self.opt.model_architecture == 'squeezenet1_0':
self.net = models.squeezenet1_0(pretrained=True)
self.net.classifier[1] = nn.Conv2d(512,
self.opt.num_classes,
kernel_size=(1, 1),
stride=(1, 1))
else:
raise ValueError('Not implemented model.')
## set to cuda and train or test phase
self.net.to(self.opt.device)
if self.opt.phase == 'train':
self.net.train()
else:
self.net.eval()
def squeezenet_10(pretrained=False, imagenet_weight=False):
"""Constructs a squeeze-Net 10 model.
Args:
pretrained (bool): If True, returns a model pre-trained on ImageNet
"""
model = models.squeezenet1_0()
if pretrained:
if imagenet_weight:
print('=== use {} as backbone'.format(imagenet_weight))
state_dict = torch.load(imagenet_weight)['state_dict']
state_dict = exchange_weightkey_in_state_dict(state_dict)
model.load_state_dict(state_dict)
else:
print('=== use pytorch default backbone')
mode1l = models.squeezenet1_0(pretrained=True)
return model
def squeezenet1_0(num_classes,
freeze_feature_extractor=False,
use_pretrained=True):
"""Wrapper for squeezenet architecture
input size = 224
Args:
num_classes:
number of output nodes for the final layer
freeze_feature_extractor:
if False (default), entire network will have gradients and can train
if True, feature block is frozen and only final layer is trained
use_pretrained:
if True, uses pre-trained ImageNet features from
Pytorch's model zoo.
"""
model_ft = models.squeezenet1_0(pretrained=use_pretrained)
set_parameter_requires_grad(model_ft, freeze_feature_extractor)
model_ft.classifier[1] = nn.Conv2d(512,
num_classes,
kernel_size=(1, 1),
stride=(1, 1))
model_ft.num_classes = num_classes
return model_ft
def dl(arch):
if 'resnet18' in arch:
model = models.resnet18(pretrained=True)
elif 'alexnet' in arch:
model = models.alexnet(pretrained=True)
elif "squeezenet" in arch:
model = models.squeezenet1_0(pretrained=True)
elif 'vgg16' in arch:
model = models.vgg16(pretrained=True)
elif 'densenet161' in arch:
model = models.densenet161(pretrained=True)
elif 'inception' in arch:
model = models.inception_v3(pretrained=True)
elif 'googlenet' in arch:
model = models.googlenet(pretrained=True)
elif 'shufflenet' in arch:
model = models.shufflenet_v2_x1_0(pretrained=True)
elif "mobilenet" in arch:
model = models.mobilenet_v2(pretrained=True)
elif "resnext" in arch:
model = models.resnext50_32x4d(pretrained=True)
elif "wide_resnet" in arch:
model = models.wide_resnet50_2(pretrained=True)
elif "mnasnet" in arch:
model = models.mnasnet1_0(pretrained=True)
elif "vgg11" in arch:
model = models.vgg11(pretrained=True)
return model
def load_models():
alexnet = models.alexnet(pretrained=True)
vgg16 = models.vgg16(pretrained=True)
googlenet = models.googlenet(pretrained=True)
resnet18 = models.resnet18(pretrained=True)
vgg11 = models.vgg11(pretrained=True)
vgg13 = models.vgg13(pretrained=True)
vgg13_bn = models.vgg13_bn(pretrained=True)
vgg16_bn = models.vgg16_bn(pretrained=True)
vgg19 = models.vgg19(pretrained=True)
vgg19_bn = models.vgg19(pretrained=True)
resnet34 = models.resnet34(pretrained=True)
resnet50 = models.resnet50(pretrained=True)
resnet101 = models.resnet101(pretrained=True)
resnet152 = models.resnet152(pretrained=True)
squeezenet1_0 = models.squeezenet1_0(pretrained=True)
squeezenet1_1 = models.squeezenet1_1(pretrained=True)
densenet121 = models.densenet121(pretrained=True)
densenet161 = models.densenet161(pretrained=True)
densenet169 = models.densenet169(pretrained=True)
densenet201 = models.densenet201(pretrained=True)
inception_v3 = models.inception_v3(pretrained=True)
mobilenet_v2 = models.mobilenet_v2(pretrained=True)
resnext50_32x4d = models.resnext50_32x4d(pretrained=True)
resnext101_32x8d = models.resnext101_32x8d(pretrained=True)
return [
alexnet, vgg11, vgg13, vgg13_bn, vgg16, vgg16_bn, vgg19, vgg19_bn,
resnet18, resnet34, resnet50, resnet101, resnet152, squeezenet1_0,
squeezenet1_1, densenet121, densenet161, densenet169, densenet201,
inception_v3, googlenet, mobilenet_v2, resnext50_32x4d,
resnext101_32x8d
]
def load_all_NN():
"""
loads all networks and sets eval
:return: all networks
"""
import torchvision.models as models
resnet152 = models.resnet152(pretrained=True)
resnet152.eval()
alexnet = models.alexnet(pretrained=True)
alexnet.eval()
squeezenet = models.squeezenet1_0(pretrained=True)
squeezenet.eval()
vgg16 = models.vgg16(pretrained=True)
vgg16.eval()
densenet = models.densenet161(pretrained=True)
densenet.eval()
inception = models.inception_v3(pretrained=True)
inception.eval()
return resnet152, alexnet, squeezenet, vgg16, densenet, inception
def squeezenet(layer, use_gpu, pretrained):
model = models.squeezenet1_0(pretrained=pretrained)
# print(model)
# use_gpu = False
for parma in model.parameters():
parma.requires_grad = False
# model.avgpool = torch.nn.AdaptiveAvgPool2d((7,7))
model.features[0] = torch.nn.Sequential(
torch.nn.Conv2d(layer, 96, kernel_size=(7, 7), stride=(2, 2)))
model.classifier = torch.nn.Sequential(
torch.nn.Dropout(p=0.5, inplace=False),
torch.nn.Conv2d(512, 3, kernel_size=(1, 1), stride=(1, 1)),
torch.nn.ReLU(inplace=True),
torch.nn.AdaptiveAvgPool2d(output_size=(1, 1)))
# for indexf,parmaf in enumerate(model.features[0].parameters()):
# if indexf == 0:
# parmaf.requires_grad = True
# for index, parma in enumerate(model.classifier.parameters()):
# if index == 1:
# parma.requires_grad = True
if use_gpu:
# if torch.cuda.device_count()>1:
# model = nn.DataParallel(model)
model = model.cuda()
return model
def squeezenet1_0(num_classes=1000, pretrained='imagenet'):
r"""SqueezeNet model architecture from the `"SqueezeNet: AlexNet-level
accuracy with 50x fewer parameters and <0.5MB model size"
<https://arxiv.org/abs/1602.07360>`_ paper.
"""
model = models.squeezenet1_0(pretrained=False)
if pretrained is not None:
settings = pretrained_settings['squeezenet1_0'][pretrained]
model = load_pretrained(model, num_classes, settings)
return model
def initialize_model(model_name, num_classes, feature_extract, use_pretrained=True):
# Initialize these variables which will be set in this if statement. Each of these
# variables is model specific.
model_ft = None
input_size = 0
if model_name == "resnet":
""" Resnet18
"""
model_ft = models.resnet18(pretrained=use_pretrained)
set_parameter_requires_grad(model_ft, feature_extract)
num_ftrs = model_ft.fc.in_features
model_ft.fc = nn.Linear(num_ftrs, num_classes)
input_size = 224
elif model_name == "alexnet":
""" Alexnet
"""
model_ft = models.alexnet(pretrained=use_pretrained)
set_parameter_requires_grad(model_ft, feature_extract)
num_ftrs = model_ft.classifier[6].in_features
model_ft.classifier[6] = nn.Linear(num_ftrs,num_classes)
input_size = 224
elif model_name == "vgg":
""" VGG11_bn
"""
model_ft = models.vgg11_bn(pretrained=use_pretrained)
set_parameter_requires_grad(model_ft, feature_extract)
num_ftrs = model_ft.classifier[6].in_features
model_ft.classifier[6] = nn.Linear(num_ftrs,num_classes)
input_size = 224
elif model_name == "squeezenet":
""" Squeezenet
"""
model_ft = models.squeezenet1_0(pretrained=use_pretrained)
set_parameter_requires_grad(model_ft, feature_extract)
model_ft.classifier[1] = nn.Conv2d(512, num_classes, kernel_size=(1,1), stride=(1,1))
model_ft.num_classes = num_classes
input_size = 224
elif model_name == "densenet":
""" Densenet
"""
model_ft = models.densenet121(pretrained=use_pretrained)
set_parameter_requires_grad(model_ft, feature_extract)
num_ftrs = model_ft.classifier.in_features
model_ft.classifier = nn.Linear(num_ftrs, num_classes)
input_size = 224
elif model_name == "inception":
""" Inception v3
Be careful, expects (299,299) sized images and has auxiliary output
"""
model_ft = models.inception_v3(pretrained=use_pretrained)
set_parameter_requires_grad(model_ft, feature_extract)
# Handle the auxilary net
num_ftrs = model_ft.AuxLogits.fc.in_features
model_ft.AuxLogits.fc = nn.Linear(num_ftrs, num_classes)
# Handle the primary net
num_ftrs = model_ft.fc.in_features
model_ft.fc = nn.Linear(num_ftrs,num_classes)
input_size = 299
else:
print("Invalid model name, exiting...")
exit()
return model_ft, input_size
