def get_model(model_type, dataset):
if model_type == "resnet":
if dataset == 'document':
return resnet.resnet20(8)
elif dataset == 'corner':
return resnet.resnet20(2)
if model_type == "resnet8":
if dataset == 'document':
return resnet.resnet8(8)
elif dataset == 'corner':
return resnet.resnet8(2)
elif model_type == 'shallow':
if dataset == 'document':
return tm.MobileNet(8)
elif dataset == 'corner':
return tm.MobileNet(2)
elif model_type == "squeeze":
if dataset == 'document':
return models.squeezenet1_1(True)
elif dataset == 'corner':
return models.squeezenet1_1(True)
else:
print(
"Unsupported model; either implement the model in model/ModelFactory or choose a different model"
)
assert (False)
def _load_model(name='squeezenet'):
if name in _models_to_use:
return _models_to_use[name]
if name == 'densenet':
net = _models_to_use[name] = models.densenet169(pretrained=True)
elif name == 'squeezenet':
net = _models_to_use[name] = models.squeezenet1_1(pretrained=True)
else:
print('Unknown model name, using squeezenet')
net = _models_to_use[name] = models.squeezenet1_1(pretrained=True)
net._modules.get('features').register_forward_hook(_hook_feature)
return net
def score(image, file_name):
device = torch.device("cpu")
#base_model = models.alexnet(pretrained=False)
base_model = models.squeezenet1_1(pretrained=True)
#modules = list(base_model.children())[:-1] # delete the last fc layer.
#base_model = nn.Sequential(*modules)
model = NIMA(base_model)
model.load_state_dict(torch.load(file_name))
model.to(device)
model.eval()
test_transform = transforms.Compose([
transforms.Resize((224, 224)),
#transforms.RandomCrop(224),
transforms.ToTensor()
])
image = test_transform(image)
image = image.unsqueeze(0)
with torch.no_grad():
outputs = model(image)
outputs = outputs.view(-1, 5, 1)
predicted_mean, predicted_std = 0.0, 0.0
for i in range(5):
predicted_mean += i * outputs[:, i].cpu()
for i in range(5):
predicted_std += outputs[:, i].cpu() * (i - predicted_mean)**2
output_score = predicted_mean.numpy().flatten().tolist()[0]
output_score_std = predicted_std.numpy().flatten().tolist()[0]
return output_score, output_score_std
def dog_detector(img):
## : Complete the function.
'''
Use pre-trained VGG-16 model to obtain index corresponding to
predicted ImageNet class for image at specified path
Args:
img_path: path to an image
Returns:
Index corresponding to VGG-16 model's prediction
'''
squeezenet = models.squeezenet1_1(pretrained=True)
squeezenet.eval()
transform = transforms.Compose([
transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
])
img_t = transform(img)
batch_t = torch.unsqueeze(img_t, 0)
out = squeezenet(batch_t)
out = torch.nn.functional.softmax(out, dim=1)[0] * 100
return 150 < int(out.argmax()) < 269 # true/false
def __init__(self, num_outputs=2, max_velocity = 0.7, max_steering=np.pi/2):
"""
Parameters
----------
num_outputs : int
number of outputs of the action space (default 2)
max_velocity : float
the maximum velocity used by the teacher (default 0.7)
max_steering : float
maximum steering angle as we are predicting normalized [0-1] (default pi/2)
"""
super(Squeezenet, self).__init__()
self._device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
self.model = models.squeezenet1_1()
self.num_outputs = num_outputs
self.max_velocity_tensor = torch.tensor(max_velocity).to(self._device)
self.max_steering = max_steering
# using a subset of full squeezenet for input image features
self.model.features = nn.Sequential(*list(self.model.features.children())[:6])
self.final_conv = nn.Conv2d(32, self.num_outputs, kernel_size=1, stride=1)
self.model.classifier = nn.Sequential(
nn.Dropout(p=0.15),
nn.Conv2d(128, 64, kernel_size=3, stride=1),
nn.ReLU(inplace=True),
nn.Conv2d(64, 32, kernel_size=3, stride=1),
nn.Dropout(p=0.15),
self.final_conv,
nn.AdaptiveAvgPool2d((1, 1))
)
self.model.num_classes = self.num_outputs
self._init_weights()
def get_model(index):
switcher = {
1: models.googlenet(pretrained=True),
2: models.squeezenet1_1(pretrained=True),
3: models.mnasnet1_0(pretrained=True),
}
return switcher.get(index, alex_net())
def predict(*args, **kwargs):
# download image using sent url
response = requests.get(kwargs["image"])
# create a pillow image
image_pillow = Image.open(io.BytesIO(response.content))
# preprocess image
image_tensor = preprocess_image(image_pillow)
image_tensor.unsqueeze_(0)
# instantiate pre-trained squeezenet model
squeeze = models.squeezenet1_1(pretrained=True)
# select device for computation
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
squeeze.to(device)
image_variable = Variable(image_tensor)
image_variable = image_variable.to(device)
out = squeeze(image_variable)
# download labels from url
labels = {
int(key): value
for (key, value) in requests.get(kwargs["labels"]).json().items()
}
# get predicted class
predict_label = labels[out.data.cpu().numpy().argmax()]
return predict_label
def load_pretrained_weights(model_name):
'''
Load weights of a pretrained pytorch model for feature extraction
Example: For Alexnet, a torch.nn.Sequential model with everything
but the fully connected layers is returned
Input parameters:
model_name: name of the model to load. Options:
vgg16_bn
squeezenet1_1
resnet101
alexnet
Returns:
(torch.nn.Sequential) The first N layers of the pretrained model
that are useful for feature extraction. N depends on which model
'''
if model_name == 'vgg16_bn':
vgg16_bn = models.vgg16_bn(pretrained=True)
return torch.nn.Sequential(*list(vgg16_bn.features.children())[:-1])
elif model_name == 'squeezenet1_1':
fnet = models.squeezenet1_1(pretrained=True)
return torch.nn.Sequential(*list(fnet.features.children())[:-1])
elif model_name == 'resnet101':
fnet = models.resnet101(pretrained=True)
return torch.nn.Sequential(*list(fnet.children())[:-2])
elif model_name == 'alexnet':
fnet = models.alexnet(pretrained=True)
return torch.nn.Sequential(*list(fnet.features.children()))
else:
raise NotImplementedError
sys.exit()
def build_backbone(self, training):
"""
Builds the desired backbone
"""
if self.backbone == "resnet18":
resnet = models.resnet18(pretrained=training)
modules = list(resnet.children())[:-1]
local_net = nn.Sequential(*modules)
for param in local_net.parameters():
param.requires_grad = False
elif self.backbone == "resnet34":
resnet = models.resnet34(pretrained=training)
modules = list(resnet.children())[:-1]
local_net = nn.Sequential(*modules)
for param in local_net.parameters():
param.requires_grad = False
elif self.backbone == "alexnet":
local_net = models.alexnet(pretrained=training)
local_net = local_net.features
elif self.backbone == "vgg":
local_net = models.vgg11(pretrained=training)
local_net = local_net.features
elif self.backbone == "squeeze":
local_net = models.squeezenet1_1(pretrained=training)
local_net = local_net.features
return local_net
def train(epochs, num_class, trainloader, model_save_path):
resnet34 = models.squeezenet1_1(pretrained=True, num_classes=1000)
resnet34.fc = nn.Linear(512, num_class)
if t.cuda.is_available():
resnet34.cuda()
criterion = nn.CrossEntropyLoss() # 交叉熵损失函数
optimizer = optim.SGD(resnet34.parameters(), lr=0.001, momentum=0.9)
for epoch in range(epochs):
running_loss = 0.0
for i, data in enumerate(trainloader, 0):
inputs, labels = data
if t.cuda.is_available():
inputs, labels = inputs.cuda(), labels.cuda()
optimizer.zero_grad()
outputs = resnet34(inputs)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
running_loss += loss.item()
# 每32个batch打印一下loss
if i % 32 == 31:
print('epoch: %d, batch: %d loss: %f' %
(epoch + 1, i + 1, running_loss / 32))
running_loss = 0.0
t.save(resnet34, model_save_path)
def squeezenet_11(pretrained=False, imagenet_weight=False):
"""Constructs a squeeze-Net 11 model.
Args:
pretrained (bool): If True, returns a model pre-trained on ImageNet
"""
model = models.squeezenet1_1()
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')
model = models.squeezenet1_1(pretrained=True)
return model
def Predict(self, request, context):
# download image using sent image_url
response = requests.get(request.image)
# create a pillow image
img_pil = Image.open(io.BytesIO(response.content))
# preprocess image
img_tensor = preprocess_image(img_pil)
img_tensor.unsqueeze_(0)
# instantiate pre-trained squeezenet model
squeeze = models.squeezenet1_1(pretrained=True)
squeeze.to(device)
img_variable = Variable(img_tensor)
# transfer image to GPU
img_variable = img_variable.to(device)
with torch.autograd.profiler.profile() as prof:
# forward pass
for iterations in range(30):
fc_out = squeeze(img_variable)
# download labels to make output human-readable
labels = {
int(key): value
for (key, value) in requests.get(label_url).json().items()
}
# get predicted class
predict_label = labels[fc_out.data.cpu().numpy().argmax()]
print(prof.table())
return dlframework_pb2.DlReply(prediction=predict_label)
def sn01_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_1(pretrained=True)
squeeze._modules['features'][2] = torch.nn.modules.AdaptiveAvgPool2d(
(54, 54))
print('SqueezeNet01 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 __init__(self, n_layers=2, n_hidden=16, n_output=3):
super(SqueezenetLSTM, self).__init__()
self.n_layers = n_layers
self.n_hidden = n_hidden
inception = models.squeezenet1_1()
inception.avgpool = nn.AdaptiveAvgPool2d(1)
self.conv = nn.Sequential(*list(inception.children())[:-1])
# No more squeezenet
self.conv = nn.Sequential(
nn.Conv2d(1, 16, kernel_size=3, stride=1, padding=2),
nn.BatchNorm2d(16), nn.ReLU(), nn.MaxPool2d(kernel_size=2,
stride=2),
nn.Conv2d(16, 32, kernel_size=3, stride=1, padding=2),
nn.BatchNorm2d(32), nn.ReLU(), nn.MaxPool2d(kernel_size=2,
stride=2))
self.lstm = nn.LSTM(input_size=16928,
hidden_size=self.n_hidden,
num_layers=self.n_layers,
batch_first=True)
self.fc = nn.Sequential(nn.Linear(n_hidden, 8), nn.ReLU(),
nn.Dropout(0.2), nn.Linear(8, n_output))
def __init__(self, pretrained=True):
super(FeatExtractorSqueezeNetx16, self).__init__()
print("loading layers from squeezenet1_1...")
sq = models.squeezenet1_1(pretrained=pretrained)
self.conv1 = nn.Sequential(
sq.features[0],
sq.features[1],
)
self.conv2 = nn.Sequential(
nn.MaxPool2d(kernel_size=3, stride=2, padding=1),
sq.features[3],
sq.features[4],
)
self.conv3 = nn.Sequential(
nn.MaxPool2d(kernel_size=3, stride=2, padding=1),
sq.features[6],
sq.features[7],
)
self.conv4 = nn.Sequential(
nn.MaxPool2d(kernel_size=3, stride=2, padding=1),
sq.features[9],
sq.features[10],
sq.features[11],
sq.features[12],
)
self.conv1[0].padding = (1, 1)
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 __init__(self):
"""Select conv1_1 ~ conv5_1 activation maps."""
super(SqueezeNet1_1, self).__init__()
self.select_feats = ['maxpool1',
'maxpool3',
'maxpool5',
'fire7',
'fire9', ]
self.select_classifier = ['conv11', 'avgpool13']
self.feat_list = self.select_feats + self.select_classifier
self.sqnet_feats = models.squeezenet1_1(pretrained=True).features
self.sqnet_classifier = models.squeezenet1_1(
pretrained=True).classifier
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 __init__(self, feature_d=128):
super().__init__()
self.model_name = "squeezenet"
self.model = squeezenet1_1(pretrained=True)
self.extract = nn.Sequential( # fune-tuning
nn.Dropout(p=0.5), nn.Conv2d(512, feature_d, 1),
nn.ReLU(inplace=True), nn.AvgPool2d(13, stride=1))
def __init__(self):
super(Squeezenet, self).__init__()
self.model = tvm.squeezenet1_1(pretrained=True)
self.features = nn.Sequential(*list(self.model.children())[:-3]).to(
utils.torch_device())
def __init__(self, pretrained=True):
super(Squeezenet, self).__init__()
squeeze_pretrained_features = models.squeezenet1_1(
pretrained=pretrained).features
self.channels = [64, 128, 256, 384, 384, 512, 512]
self.featureMapSize = [111, 55, 27, 13, 13, 13, 13]
self.slice1 = torch.nn.Sequential()
self.slice2 = torch.nn.Sequential()
self.slice3 = torch.nn.Sequential()
self.slice4 = torch.nn.Sequential()
self.slice5 = torch.nn.Sequential()
self.slice6 = torch.nn.Sequential()
self.slice7 = torch.nn.Sequential()
self.N_slices = 7
for x in range(2):
self.slice1.add_module(str(x), squeeze_pretrained_features[x])
for x in range(2, 5):
self.slice2.add_module(str(x), squeeze_pretrained_features[x])
for x in range(5, 8):
self.slice3.add_module(str(x), squeeze_pretrained_features[x])
for x in range(8, 10):
self.slice4.add_module(str(x), squeeze_pretrained_features[x])
for x in range(10, 11):
self.slice5.add_module(str(x), squeeze_pretrained_features[x])
for x in range(11, 12):
self.slice6.add_module(str(x), squeeze_pretrained_features[x])
for x in range(12, 13):
self.slice7.add_module(str(x), squeeze_pretrained_features[x])
self.layerList = [
self.slice1, self.slice2, self.slice3, self.slice4, self.slice5,
self.slice6, self.slice7
]
def get_feature_net(net_name):
'''
Get the object representing the desired feature extraction network
Note: only the part of the network considered useful for feature
extraction is returned. i.e. everythnig but the fully
connected layers of AlexNet.
Input parameters:
net_name: (str) the name of the desired network
Availble net names:
vgg16_bn
squeezenet1_1
resenet101
alexnet
'''
if net_name == 'vgg16_bn':
fnet = models.vgg16_bn(pretrained=False)
return torch.nn.Sequential(
*list(fnet.features.children())[:-1]), 16, 512
elif net_name == 'squeezenet1_1':
fnet = models.squeezenet1_1(pretrained=False)
return torch.nn.Sequential(
*list(fnet.features.children())[:-1]), 16, 512
elif net_name == 'resnet101':
fnet = models.resnet101(pretrained=False)
return torch.nn.Sequential(*list(fnet.children())[:-2]), 32, 2048
elif net_name == 'alexnet':
fnet = models.alexnet(pretrained=False)
return torch.nn.Sequential(
*list(fnet.features.children())), 17, 256
else:
raise NotImplementedError
def __init__(self, requires_grad=False, pretrained=True):
super(squeezenet, self).__init__()
pretrained_features = tv.squeezenet1_1(pretrained=pretrained).features
self.slice1 = torch.nn.Sequential()
self.slice2 = torch.nn.Sequential()
self.slice3 = torch.nn.Sequential()
self.slice4 = torch.nn.Sequential()
self.slice5 = torch.nn.Sequential()
self.slice6 = torch.nn.Sequential()
self.slice7 = torch.nn.Sequential()
self.N_slices = 7
for x in range(2):
self.slice1.add_module(str(x), pretrained_features[x])
for x in range(2,5):
self.slice2.add_module(str(x), pretrained_features[x])
for x in range(5, 8):
self.slice3.add_module(str(x), pretrained_features[x])
for x in range(8, 10):
self.slice4.add_module(str(x), pretrained_features[x])
for x in range(10, 11):
self.slice5.add_module(str(x), pretrained_features[x])
for x in range(11, 12):
self.slice6.add_module(str(x), pretrained_features[x])
for x in range(12, 13):
self.slice7.add_module(str(x), pretrained_features[x])
if not requires_grad:
for param in self.parameters():
param.requires_grad = False
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 finetuneSqueeNet(train_dset, train_loader, val_loader, args):
model = models.squeezenet1_1(pretrained=True)
# print (model)
num_classes = len(train_dset.classes)
model.classifier[1] = nn.Conv2d(512,
num_classes,
kernel_size=(1, 1),
stride=(1, 1))
model.num_classes = num_classes
model.type(dtype)
loss_fn = nn.CrossEntropyLoss().type(dtype)
for param in model.parameters():
param.requires_grad = False
for param in model.classifier.parameters():
param.requires_grad = True
optimizer = optim.SGD(model.parameters(),
lr=args.lr,
momentum=args.momentum)
for epoch in range(args.epochs):
# Run an epoch over the training data.
print('Starting epoch %d / %d' % (epoch + 1, args.epochs))
run_epoch(model, loss_fn, train_loader, optimizer, dtype)
# Check accuracy on the train and val sets.
train_acc = check_accuracy(model, train_loader, dtype)
val_acc = check_accuracy(model, val_loader, dtype)
print('Train accuracy: ', train_acc)
print('Val accuracy: ', val_acc)
def __init__(self, device="cpu", jit=False):
self.device = device
self.jit = jit
self.model = models.squeezenet1_1()
if self.jit:
self.model = torch.jit.script(self.model)
self.example_inputs = (torch.randn((32, 3, 224, 224)), )
def __init__(self, requires_grad=False, pretrained=True):
super(squeezenet, self).__init__()
pretrained_features = models.squeezenet1_1(pretrained=pretrained).features
self.slice1 = torch.nn.Sequential()
self.slice2 = torch.nn.Sequential()
self.slice3 = torch.nn.Sequential()
self.slice4 = torch.nn.Sequential()
self.slice5 = torch.nn.Sequential()
self.slice6 = torch.nn.Sequential()
self.slice7 = torch.nn.Sequential()
self.N_slices = 7
for x in range(2):
self.slice1.add_module(str(x), pretrained_features[x])
for x in range(2,5):
self.slice2.add_module(str(x), pretrained_features[x])
for x in range(5, 8):
self.slice3.add_module(str(x), pretrained_features[x])
for x in range(8, 10):
self.slice4.add_module(str(x), pretrained_features[x])
for x in range(10, 11):
self.slice5.add_module(str(x), pretrained_features[x])
for x in range(11, 12):
self.slice6.add_module(str(x), pretrained_features[x])
for x in range(12, 13):
self.slice7.add_module(str(x), pretrained_features[x])
if not requires_grad:
for param in self.parameters():
param.requires_grad = False
def __init__(self,
dimension,
deconv=UpsampleDeConv,
activation=nn.LeakyReLU(),
drop_out: float = 0.5):
super(FreiburgSqueezeGenerator, self).__init__(dimension, deconv,
activation, drop_out)
pretrained_features = models.squeezenet1_1(pretrained=True).features
conv = ConvLayer(dimension, 64, kernel_size=3, stride=1, bias=True)
if dimension == 1 or dimension == 3:
weight = torch.FloatTensor(64, dimension, 3, 3)
parameters = list(pretrained_features.parameters())
for i in range(64):
if dimension == 1:
weight[i, :, :, :] = parameters[0].data[i].mean(0)
else:
weight[i, :, :, :] = parameters[0].data[i]
conv.weight.data.copy_(weight)
conv.bias.data.copy_(parameters[1].data)
self.enc1 = nn.Sequential(conv)
for x in range(1, 2):
self.enc1.add_module(str(x), pretrained_features[x])
self.enc2 = torch.nn.Sequential()
self.enc3 = torch.nn.Sequential()
self.enc4 = torch.nn.Sequential()
for x in range(3, 5):
self.enc2.add_module(str(x), pretrained_features[x])
for x in range(6, 8):
self.enc3.add_module(str(x), pretrained_features[x])
for x in range(9, 13):
self.enc4.add_module(str(x), pretrained_features[x])
def infer_image():
# img = torch.load("/tmp/image.pt")
img = torch.rand(1, 3, 224, 224)
model = models.squeezenet1_1(pretrained=True)
with torch.no_grad():
output = model(img)
print(f"Predicted class: {output.argmax()}")
def init_model(self):
"""initializes the model with the config file"""
if self.config["model"] == "resnet":
network = models.resnext50_32x4d(pretrained=True)
elif self.config["model"] == "vgg19":
network = models.vgg19(pretrained=True)
elif self.config["model"] == "densenet":
network = models.densenet121(pretrained=True)
elif self.config["model"] == "inception":
network = models.inception_v3(pretrained=True)
elif self.config["model"] == "mobile":
network = models.mobilenet_v2(pretrained=True)
elif self.config["model"] == "shuffle":
network = models.shufflenet_v2_x0_5(pretrained=True)
elif self.config["model"] == "squeeze":
network = models.squeezenet1_1(pretrained=True)
elif self.config["model"] == "resnetx":
network = models.resnext101_32x8d(pretrained=True)
elif self.config["model"] == "masnet":
network = models.mnasnet1_0(pretrained=True)
elif self.config["model"] == "googlenet":
network = models.googlenet(pretrained=True)
elif self.config["model"] == "alexnet":
network = models.alexnet(pretrained=True)
else:
print("Invalid Model")
network.eval()
self.model = network.to(self.device)
if self.config["fp16"]:
scaler = torch.cuda.amp.GradScaler()
self.model = self.model.half()
def __init__(self, hparams, *args, **kwargs):
super(SqueezeNet, self).__init__()
self.hparams = hparams
self.y_trues = []
self.predictions = []
self.model = squeezenet1_1(pretrained=True)
self.criterion = nn.CrossEntropyLoss()
def squeezenet1_1(num_classes=1000, pretrained='imagenet'):
r"""SqueezeNet 1.1 model from the `official SqueezeNet repo
<https://github.com/DeepScale/SqueezeNet/tree/master/SqueezeNet_v1.1>`_.
SqueezeNet 1.1 has 2.4x less computation and slightly fewer parameters
than SqueezeNet 1.0, without sacrificing accuracy.
"""
model = models.squeezenet1_1(pretrained=False)
if pretrained is not None:
settings = pretrained_settings['squeezenet1_1'][pretrained]
model = load_pretrained(model, num_classes, settings)
return model
def __init__(self, num_classes=2):
super(SqueezeNet, self).__init__()
self.model_name = 'squeezenet'
self.model = squeezenet1_1(pretrained=True)
# 修改 原始的num_class: 预训练模型是1000分类
self.model.num_classes = num_classes
self.model.classifier = nn.Sequential(
nn.Dropout(p=0.5),
nn.Conv2d(512, num_classes, 1),
nn.ReLU(inplace=True),
nn.AvgPool2d(13, stride=1)
)
# In[10]: # LABELS_URL is a JSON file that maps label indices to English descriptions of the ImageNet classes # IMG_URL can be any image you like. If it’s in one of the 1,000 ImageNet classes this code should correctly classify it. LABELS_URL = 'https://s3.amazonaws.com/outcome-blog/imagenet/labels.json' IMG_URL = 'https://s3.amazonaws.com/outcome-blog/wp-content/uploads/2017/02/25192225/cat.jpg' # In[11]: #Initialize the model squeeze = models.squeezenet1_1(pretrained=True) # In[12]: normalize = transforms.Normalize( mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225] ) preprocess = transforms.Compose([ transforms.Scale(256), transforms.CenterCrop(224), transforms.ToTensor(), normalize ])
import argparse
import os
from six import text_type as _text_type
from mmdnn.conversion.examples.imagenet_test import TestKit
import torch
import torchvision.models as models
NETWORKS_MAP = {
'inception_v3' : lambda : models.inception_v3(pretrained=True),
'vgg16' : lambda : models.vgg16(pretrained=True),
'vgg19' : lambda : models.vgg19(pretrained=True),
'resnet152' : lambda : models.resnet152(pretrained=True),
'densenet' : lambda : models.densenet201(pretrained=True),
'squeezenet' : lambda : models.squeezenet1_1(pretrained=True)
}
def _main():
parser = argparse.ArgumentParser()
parser.add_argument('-n', '--network',
type=_text_type, help='Model Type', required=True,
choices=NETWORKS_MAP.keys())
parser.add_argument('-i', '--image', type=_text_type, help='Test Image Path')
args = parser.parse_args()
file_name = "imagenet_{}.pth".format(args.network)
