def __init__(self, subtype='shufflenet_v2_x1_0', out_stages=[2,3,4], backbone_path=None):
super(ShuffleNetV2, self).__init__()
self.out_stages = out_stages
self.backbone_path = backbone_path
if subtype == 'shufflenet_v2_x0_5':
self.backbone = shufflenet_v2_x0_5(pretrained=not self.backbone_path)
self.out_channels = [24, 48, 96, 192, 1024]
elif subtype == 'shufflenet_v2_x1_0':
self.backbone = shufflenet_v2_x1_0(pretrained=not self.backbone_path)
self.out_channels = [24, 116, 232, 464, 1024]
elif subtype == 'shufflenet_v2_x1_5':
self.backbone = shufflenet_v2_x1_5(pretrained=not self.backbone_path)
self.out_channels = [24, 176, 352, 704, 1024]
elif subtype == 'shufflenet_v2_x2_0':
self.backbone = shufflenet_v2_x2_0(pretrained=not self.backbone_path)
self.out_channels = [24, 244, 488, 976, 2048]
else:
raise NotImplementedError
self.out_channels = self.out_channels[self.out_stages[0]:self.out_stages[-1] + 1]
if self.backbone_path:
self.backbone.load_state_dict(torch.load(self.backbone_path))
else:
self.init_weights()
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 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 shufflenet_v2_x05(pretrained=False, imagenet_weight=False):
"""Constructs a Shufflenet-V2_x05 model.
Args:
pretrained (bool): If True, returns a model pre-trained on ImageNet
"""
print('=== Using Shufflenet-V2_x05 ===')
model = models.shufflenet_v2_x0_5()
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.shufflenet_v2_x0_5(pretrained=True)
return model
def __init__(self, num_classes=2):
super(ShuffleNet, self).__init__()
self.model_name = "ShuffleNet"
model = models.shufflenet_v2_x0_5(pretrained=True)
for param in model.parameters():
param.requires_grad = False
model.fc = nn.Linear(model.fc.in_features, num_classes)
self.model = model
def get_shufflenet_v2_x0_5(class_num):
model = models.shufflenet_v2_x0_5(pretrained=True)
set_parameter_requires_grad(model)
model.name = 'shufflenet_v2_x0_5'
n_inputs = model.fc.in_features
model.fc = nn.Linear(n_inputs, class_num)
return model, 244
def test_shufflenet_v2_dynamic_axes(self):
model = models.shufflenet_v2_x0_5(pretrained=True)
dummy_input = torch.randn(1, 3, 224, 224, requires_grad=True)
test_inputs = torch.cat([dummy_input, dummy_input, dummy_input], 0)
self.run_model(model, [(dummy_input,), (test_inputs,)],
input_names=["input_images"],
output_names=["output"],
dynamic_axes={"input_images": {0: 'batch_size'}, "output": {0: 'batch_size'}},
tolerate_small_mismatch=True)
def shufflenet_v2_x0_5_nc_1k():
m = shufflenet_v2_x0_5(pretrained=True)
m.fc = nn.Identity()
del m.conv5[2], m.conv5[1] # remove ReLU and BN
m.conv5[0] = nn.Conv2d(192,
1000,
kernel_size=(1, 1),
stride=(1, 1),
bias=True)
return m
def __init__(self, pretrained=True):
super().__init__()
self.shuffle = models.shufflenet_v2_x0_5(pretrained=pretrained)
# self.features = self.mobile.features
self.conv1 = self.shuffle.conv1
self.maxpool = self.shuffle.maxpool
self.stage2 = self.shuffle.stage2
self.stage3 = self.shuffle.stage3
self.stage4 = self.shuffle.stage4
self.conv5 = self.shuffle.conv5
self.fc = nn.Linear(self.shuffle._stage_out_channels[-1], 6)
def get_shufflenet_model(
model_name, output_dim,
):
""" Helper function to get a model
:param model_name: Name of the model
:type model_name: str
:param output_dim: Output dimension of the model
:type output_dim: int
:return model: The model
:rtype: ShuffleNetV2
"""
# Getting the model
if model_name == "shufflenet_v2_x0_5":
model = shufflenet_v2_x0_5(num_classes=output_dim)
pretrained_model = models.shufflenet_v2_x0_5(pretrained=True)
elif model_name == "shufflenet_v2_x1_0":
model = shufflenet_v2_x1_0(num_classes=output_dim)
pretrained_model = models.shufflenet_v2_x1_0(pretrained=True)
elif model_name == "shufflenet_v2_x1_5":
model = shufflenet_v2_x1_5(num_classes=output_dim)
pretrained_model = models.shufflenet_v2_x1_5(pretrained=True)
elif model_name == "shufflenet_v2_x2_0":
model = shufflenet_v2_x2_0(num_classes=output_dim)
pretrained_model = models.shufflenet_v2_x2_0(pretrained=True)
else:
model = shufflenet_v2_x0_5(num_classes=output_dim)
pretrained_model = models.shufflenet_v2_x0_5(pretrained=True)
IN_FEATURES = pretrained_model.fc.in_features
fc = nn.Linear(IN_FEATURES, output_dim)
pretrained_model.fc = fc
model.load_state_dict(pretrained_model.state_dict())
return model
def get_shufflenet_v2_0_5(output_dim=2, weight=None):
model = models.shufflenet_v2_x0_5(pretrained=True)
for parma in model.parameters():
parma.requires_grad = False
model.fc = torch.nn.Linear(1024, output_dim)
for parma in model.fc.parameters():
parma.requires_grad = True
if weight:
model.load_state_dict(torch.load(weight))
return model
def pytorch_shufflenetv2(request: Any) -> ModelAndData:
if request.config.option.skipslow:
pytest.skip()
import torchvision.models as models
model = models.shufflenet_v2_x0_5(pretrained=True).eval()
preprocessing = dict(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225], axis=-3)
fmodel = fbn.PyTorchModel(model, bounds=(0, 1), preprocessing=preprocessing)
x, y = fbn.samples(fmodel, dataset="imagenet", batchsize=8)
x = ep.astensor(x)
y = ep.astensor(y)
return fmodel, x, y
def convert_shufflenetv2():
# model= models.shufflenet_v2_x1_0(pretrained=False, num_classes=9)
# path = r"checkpoint/data_11_16/shufflenetv2/pre/111/shufflenetv2_1_my_acc=97.8541.pth"
# to_path = r"checkpoint/data_11_16/shufflenetv2/pre/111/shufflenetv2_1.onnx"
model = models.shufflenet_v2_x0_5(pretrained=False, num_classes=9)
path = r"checkpoint/data_11_16/shufflenetv2/pre/333/shufflenetv2_05_my_acc=99.5708.pth"
to_path = r"checkpoint/data_11_16/shufflenetv2/pre/333/shufflenetv2_05_my.onnx"
checkpoint = torch.load(path)
model.load_state_dict(checkpoint["net"])
print("loaded model with acc:{}".format(checkpoint["acc"]))
model.cuda()
# dummy_input = torch.randn(10, 3, 224, 224)
dummy_input = torch.randn(1, 3, 224, 224, device='cuda')
torch.onnx.export(model, dummy_input, to_path, verbose=True)
def main():
parser = argparse.ArgumentParser(description='Process args for Triplet Net Training')
parser.add_argument("--val_dir", type=str, help="Val directory",
default="/home/delta_one/Metric_Learning/tiny-imagenet-200/val")
parser.add_argument("--val_annotation_file", type=str, help="Val Annotation File",
default="/home/delta_one/Metric_Learning/tiny-imagenet-200/val/val_annotations.txt")
parser.add_argument('--classes', type=str, nargs='+', default=\
['n01443537', 'n01629819', 'n01641577', 'n01644900', 'n01698640', 'n01742172', 'n01768244', 'n01770393', 'n01774384', 'n01774750'],
help='Classes used for Train and Eval')
parser.add_argument('--samples_per_class', type=int, default=50,
help='# Samples used for Test')
parser.add_argument('--batch_size_test', type=int, default=30,
help='validation set input batch size')
args = parser.parse_args()
transform_test = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])
val_dataset = TinyImagnet_Val(root_dir=args.val_dir, transform=transform_test,
annotation_file=args.val_annotation_file, given_classes=args.classes,
samples_per_class=args.samples_per_class)
nw = 4 if torch.cuda.is_available() else 0
val_dataloader = DataLoader(val_dataset, batch_size=args.batch_size_test,
num_workers=nw, shuffle=False)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
writer = SummaryWriter(log_dir="logs/")
model_dict = {
"AlexNet": models.alexnet(pretrained=True),
"DenseNet": models.densenet121(pretrained=True),
"MNASNet": models.mnasnet0_5(pretrained=True),
"MobileNetV2": models.mobilenet_v2(pretrained=True),
"ResNet18": models.resnet18(pretrained=True),
"ResNet101": models.resnet101(pretrained=True),
"ShuffleNet": models.shufflenet_v2_x0_5(pretrained=True),
"SqueezeNet": models.squeezenet1_0(pretrained=True),
"VGG19": models.vgg19(pretrained=True)
}
test(val_dataloader, model_dict, writer, device)
def test_shufflenet_v2_dynamic_axes(self):
model = models.shufflenet_v2_x0_5(
weights=models.ShuffleNet_V2_X0_5_Weights.DEFAULT)
dummy_input = torch.randn(1, 3, 224, 224, requires_grad=True)
test_inputs = torch.cat([dummy_input, dummy_input, dummy_input], 0)
self.run_model(
model,
[(dummy_input, ), (test_inputs, )],
input_names=["input_images"],
output_names=["output"],
dynamic_axes={
"input_images": {
0: "batch_size"
},
"output": {
0: "batch_size"
}
},
tolerate_small_mismatch=True,
)
def shufflenet_v2(name, pretrained):
from torchvision.models import shufflenet_v2_x0_5, shufflenet_v2_x1_0, shufflenet_v2_x1_5, shufflenet_v2_x2_0
imagenet_pretrained = pretrained == 'imagenet'
if name == 'shufflenet_v2_x0_5':
net = shufflenet_v2_x0_5(imagenet_pretrained)
elif name == 'shufflenet_v2_x1_0':
net = shufflenet_v2_x1_0(imagenet_pretrained)
elif name == 'shufflenet_v2_x1_5':
net = shufflenet_v2_x1_5(imagenet_pretrained)
elif name == 'shufflenet_v2_x2_0':
net = shufflenet_v2_x2_0(imagenet_pretrained)
else:
raise NotImplementedError(name)
stage1 = nn.Sequential(net.conv1, net.maxpool)
layers = [stage1, net.stage2, net.stage3, net.stage4, net.conv5]
for stage, out in zip(layers, net._stage_out_channels):
stage.out_channels = out
n_pretrained = len(layers) if imagenet_pretrained else 0
return layers, True, n_pretrained
def shufflenet(layer, use_gpu, pretrained):
model = models.shufflenet_v2_x0_5(pretrained=pretrained)
# print(model)
# use_gpu = False
for parma in model.parameters():
parma.requires_grad = False
# model.avgpool = torch.nn.AdaptiveAvgPool2d((7,7))
model.conv1 = torch.nn.Sequential(
torch.nn.Conv2d(layer,
24,
kernel_size=(3, 3),
stride=(2, 2),
padding=(1, 1),
bias=False),
torch.nn.BatchNorm2d(24,
eps=1e-05,
momentum=0.1,
affine=True,
track_running_stats=True),
torch.nn.ReLU(inplace=True))
model.fc = torch.nn.Sequential(
torch.nn.Linear(in_features=1024, out_features=3, bias=True))
# for indexf,parmaf in enumerate(model.conv1.parameters()):
# if indexf == 0:
# parmaf.requires_grad = True
# for index, parma in enumerate(model.fc.parameters()):
# if index == 0:
# parma.requires_grad = True
if use_gpu:
# if torch.cuda.device_count()>1:
# model = nn.DataParallel(model)
model = model.cuda()
return model
pre_model = models.densenet121(pretrained=True, progress=True) elif TopModelName=="densenet169": pre_model = models.densenet169(pretrained=True, progress=True) elif TopModelName=="densenet161": pre_model = models.densenet161(pretrained=True, progress=True) elif TopModelName=="densenet201": pre_model = models.densenet201(pretrained=True, progress=True) elif TopModelName=="inception": pre_model = models.inception_v3(pretrained=True, progress=True) elif TopModelName=="googlenet": pre_model = models.googlenet(pretrained=True, progress=True) elif TopModelName=="shufflenet_v2_x0_5": pre_model = models.shufflenet_v2_x0_5(pretrained=True, progress=True) elif TopModelName=="shufflenet_v2_x1_0": pre_model = models.shufflenet_v2_x0_5(pretrained=True, progress=True) elif TopModelName=="shufflenet_v2_x1_5": pre_model = models.shufflenet_v2_x0_5(pretrained=True, progress=True) elif TopModelName=="shufflenet_v2_x2_0": pre_model = models.shufflenet_v2_x0_5(pretrained=True, progress=True) elif TopModelName=="mobilenet_v2": pre_model = models.mobilenet_v2(pretrained=True, progress=True) ######### ResNet's ########### elif TopModelName=='resnet18': pre_model = models.resnet18(pretrained=True, progress=True) elif TopModelName=='resnet34': pre_model = models.resnet34(pretrained=True, progress=True)
def __init__(self, model, n_classes, pretrained=True):
super(classifie, self).__init__()
if (model == 'efficientnet-b3'):
if (pretrained == True):
self.cnn_arch = EfficientNet.from_pretrained('efficientnet-b3')
else:
self.cnn_arch = EfficientNet.from_name('efficientnet-b3')
if (model == 'efficientnet-b2'):
if (pretrained == True):
self.cnn_arch = EfficientNet.from_pretrained('efficientnet-b2')
else:
self.cnn_arch = EfficientNet.from_name('efficientnet-b2')
if (model == 'efficientnet-b1'):
if (pretrained == True):
self.cnn_arch = EfficientNet.from_pretrained('efficientnet-b1')
else:
self.cnn_arch = EfficientNet.from_name('efficientnet-b1')
if (model == 'efficientnet-b0'):
if (pretrained == True):
self.cnn_arch = EfficientNet.from_pretrained('efficientnet-b0')
else:
self.cnn_arch = EfficientNet.from_name('efficientnet-b0')
if (model == 'resnet18'):
self.cnn_arch = models.resnet18(pretrained=pretrained)
if (model == 'resnet34'):
self.cnn_arch = models.resnet34(pretrained=pretrained)
if (model == 'resnet50'):
self.cnn_arch = models.resnet50(pretrained=pretrained)
if (model == 'resnet101'):
self.cnn_arch = models.resnet101(pretrained=pretrained)
if (model == 'resnet152'):
self.cnn_arch = models.resnet152(pretrained=pretrained)
if (model == 'densenet121'):
self.cnn_arch = models.densenet121(pretrained=pretrained)
if (model == 'densenet161'):
self.cnn_arch = models.densenet161(pretrained=pretrained)
if (model == 'densenet169'):
self.cnn_arch = models.densenet169(pretrained=pretrained)
if (model == 'densenet201'):
self.cnn_arch = models.densenet201(pretrained=pretrained)
if (model == 'squeezenet1_0'):
self.cnn_arch = models.squeezenet1_0(pretrained=pretrained)
if (model == 'squeezenet1_1'):
self.cnn_arch = models.squeezenet1_1(pretrained=pretrained)
if (model == 'shufflenet_v2_x0_5'):
self.cnn_arch = models.shufflenet_v2_x0_5(pretrained=pretrained)
if (model == 'shufflenet_v2_x1_0'):
self.cnn_arch = models.shufflenet_v2_x1_0(pretrained=pretrained)
if (model == 'shufflenet_v2_x1_5'):
self.cnn_arch = models.shufflenet_v2_x1_5(pretrained=pretrained)
if (model == 'shufflenet_v2_x2_0'):
self.cnn_arch = models.shufflenet_v2_x2_0(pretrained=pretrained)
if (model == 'resnext50_32x4d'):
self.cnn_arch = models.resnext50_32x4d(pretrained=pretrained)
if (model == 'resnext101_32x8d'):
self.cnn_arch = models.resnext101_32x8d(pretrained=pretrained)
self.linear1 = nn.Linear(1000, 256)
self.relu = nn.LeakyReLU()
self.linear2 = nn.Linear(256, n_classes)
#self.softmax = nn.Softmax(dim=1)
self.dropout = nn.Dropout(0.7)
def import_all(download):
"""
Commented out models do not have pre-trained variants available.
This function will load and cache all of these models if they are not already cached.
On CRC we can just run this initially as its own python script to download and cache the models.
Afterwards, this function will be implemented into some other testing code that will call it in
a similar fashion to what is shown below.
"""
print("Loading or checking if models are cached...\n")
print(
"This may take a while. Progress will be shown if models are being downloaded.\n "
)
############################################################################################
# Image Classification
alexnet = models.alexnet(pretrained=download, progress=True)
squeezenet1_0 = models.squeezenet1_0(pretrained=download, progress=True)
squeezenet1_1 = models.squeezenet1_1(pretrained=download, progress=True)
vgg16 = models.vgg16_bn(pretrained=download, progress=True)
vgg19 = models.vgg19_bn(pretrained=download, progress=True)
resnet18 = models.resnet18(pretrained=download, progress=True)
resnet34 = models.resnet34(pretrained=download, progress=True)
resnet50 = models.resnet50(pretrained=download, progress=True)
resnet101 = models.resnet101(pretrained=download, progress=True)
resnet152 = models.resnet152(pretrained=download, progress=True)
densenet121 = models.densenet121(pretrained=download,
progress=True,
memory_efficient=False)
densenet161 = models.densenet161(pretrained=download,
progress=True,
memory_efficient=False)
densenet201 = models.densenet201(pretrained=download,
progress=True,
memory_efficient=False)
densenet121_efficient = models.densenet121(pretrained=download,
progress=True,
memory_efficient=True)
densenet161_efficient = models.densenet161(pretrained=download,
progress=True,
memory_efficient=True)
densenet201_efficient = models.densenet201(pretrained=download,
progress=True,
memory_efficient=True)
googlenet = models.googlenet(pretrained=download, progress=True)
shufflenet_v2_1 = models.shufflenet_v2_x1_0(pretrained=download,
progress=True)
shufflenet_v2_0_5 = models.shufflenet_v2_x0_5(pretrained=download,
progress=True)
#shufflenet_v2_1_5 = models.shufflenet_v2_x1_5(pretrained=download, progress=True)
#shufflenet_v2_2 = models.shufflenet_v2_x2_0(pretrained=download, progress=True)
mobilenet_v2 = models.mobilenet_v2(pretrained=download, progress=True)
resnext50_32x4d = models.resnext50_32x4d(pretrained=download,
progress=True)
resnext101_32x8d = models.resnext101_32x8d(pretrained=download,
progress=True)
wide_resnet50_2 = models.wide_resnet50_2(pretrained=download,
progress=True)
wide_resnet101_2 = models.wide_resnet101_2(pretrained=download,
progress=True)
mnasnet1_0 = models.mnasnet1_0(pretrained=download, progress=True)
mnasnet0_5 = models.mnasnet0_5(pretrained=download, progress=True)
#mnasnet0_75 = models.mnasnet0_75(pretrained=download, progress=True)
#mnasnet1_3 = models.mnasnet1_3(pretrained=download, progress=True)
efficientnet_b0 = EfficientNet.from_pretrained('efficientnet-b0')
efficientnet_b1 = EfficientNet.from_pretrained('efficientnet-b1')
efficientnet_b2 = EfficientNet.from_pretrained('efficientnet-b2')
efficientnet_b3 = EfficientNet.from_pretrained('efficientnet-b3')
efficientnet_b4 = EfficientNet.from_pretrained('efficientnet-b4')
efficientnet_b5 = EfficientNet.from_pretrained('efficientnet-b5')
efficientnet_b6 = EfficientNet.from_pretrained('efficientnet-b6')
efficientnet_b7 = EfficientNet.from_pretrained('efficientnet-b7')
#
#
# ###########################################################################################
# # Video Classification
# # resnet_3d = models.video.r3d_18(pretrained=download, progress=True)
# # resnet_mixed_conv = models.video.mc3_18(pretrained=download, progress=True)
# # resnet_2_1D = models.video.r2plus1d_18(pretrained=download, progress=True)
#
# ###########################################################################################
# # Object Detection
#
fasterrcnn_resnet50 = models.detection.fasterrcnn_resnet50_fpn(
pretrained=download,
progress=True,
num_classes=91,
pretrained_backbone=True)
maskcnn_resnet50 = models.detection.maskrcnn_resnet50_fpn(
pretrained=download,
progress=True,
num_classes=91,
pretrained_backbone=True)
keypointrcnn_resnet50 = models.detection.keypointrcnn_resnet50_fpn(
pretrained=download,
progress=True,
num_classes=2,
num_keypoints=17,
pretrained_backbone=True)
#
# ###########################################################################################
# # Semantic Segmentation
#
fcn_resnet50 = models.segmentation.fcn_resnet50(pretrained=False,
progress=True,
num_classes=21,
aux_loss=None)
fcn_resnet101 = models.segmentation.fcn_resnet101(pretrained=download,
progress=True,
num_classes=21,
aux_loss=None)
deeplabv3_resnet50 = models.segmentation.deeplabv3_resnet50(
pretrained=False, progress=True, num_classes=21, aux_loss=None)
deeplabv3_resnet101 = models.segmentation.deeplabv3_resnet101(
pretrained=False, progress=True, num_classes=21, aux_loss=None)
###########################################################################################
# Generative Adversarial Networks
###########################################################################################
models_dict = {
"alexnet": alexnet,
"fcn_resnet50": fcn_resnet50,
"fcn_resnet101": fcn_resnet101,
"deeplabv3_resnet50": deeplabv3_resnet50,
"deeplabv3_resnet101": deeplabv3_resnet101,
"squeezenet1_0": squeezenet1_0,
"squeezenet1_1": squeezenet1_1,
"vgg16": vgg16,
"vgg19": vgg19,
"resnet18": resnet18,
"resnet34": resnet34,
"resnet50": resnet50,
"resnet101": resnet101,
"resnet152": resnet152,
"densenet121": densenet121,
"densenet161": densenet161,
"densenet201": densenet201,
"densenet121_efficient": densenet121_efficient,
"densenet161_efficient": densenet161_efficient,
"densenet201_efficient": densenet201_efficient,
"googlenet": googlenet,
"shufflenet_v2_1": shufflenet_v2_1,
"shufflenet_v2_0_5": shufflenet_v2_0_5,
#"shufflenet_v2_1_5": shufflenet_v2_1_5,
#"shufflenet_v2_2": shufflenet_v2_2,
"mobilenet_v2": mobilenet_v2,
"resnext50_32x4d": resnext50_32x4d,
"resnext101_32x4d": resnext101_32x8d,
"wide_resnet50_2": wide_resnet50_2,
"wide_resnet101_2": wide_resnet101_2,
"mnasnet1_0": mnasnet1_0,
#"mnasnet1_3": mnasnet1_3,
"mnasnet0_5": mnasnet0_5,
#"mnasnet0_75": mnasnet0_75,
"efficientnet_b0": efficientnet_b0,
"efficientnet_b1": efficientnet_b1,
"efficientnet_b2": efficientnet_b2,
"efficientnet_b3": efficientnet_b3,
"efficientnet_b4": efficientnet_b4,
"efficientnet_b5": efficientnet_b5,
"efficientnet_b6": efficientnet_b6,
"efficientnet_b7": efficientnet_b7,
}
# "maskcnn_resnet50" : maskcnn_resnet50,
# "keypointrcnn_resnet50" : keypointrcnn_resnet50,
# "fasterrcnn_resnet50" : fasterrcnn_resnet50,
# "resnet_3d" : resnet_3d,
# "resnet_mixed_conv" : resnet_mixed_conv,
# "resnet_2_1D" : resnet_2_1D,
return models_dict
def main(train_loader, train_datasets, valid_loader):
'''args: model_arch, train_loader, train_datasets, valid_loader, epochs, learning_rate, hidden_units, device, save_dir
saves checkpoint under save_dir
returns: checkpoint with {'epoch_tot',
'model',
'criterion',
'optimizer',
'optimizer.state_dict',
'model.state_dict',
'model.class_to_idx'}'''
arg = input_args()
model_arch = arg.arch
hidden_units = arg.hidden_units
learning_rate = arg.learning_rate
device = arg.device
epochs = arg.epochs
save_dir = arg.save_dir
if model_arch == 'alexnet':
gs_vgg = models.alexnet(pretrained=True)
elif model_arch == 'vgg11':
gs_vgg = models.vgg11(pretrained=True)
elif model_arch == 'vgg11_bn':
gs_vgg = models.vgg11_bn(pretrained=True)
elif model_arch == 'vgg13':
gs_vgg = models.vgg13(pretrained=True)
elif model_arch == 'vgg13_bn':
gs_vgg = models.vgg13_bn(pretrained=True)
elif model_arch == 'vgg16':
gs_vgg = models.vgg16(pretrained=True)
elif model_arch == 'vgg16_bn':
gs_vgg = models.vgg16_bn(pretrained=True)
elif model_arch == 'vgg19':
gs_vgg = models.vgg19(pretrained=True)
elif model_arch == 'vgg19_bn':
gs_vgg = models.vgg19_bn(pretrained=True)
elif model_arch == 'resnet18':
gs_vgg = models.resnet18(pretrained=True)
elif model_arch == 'resnet34':
gs_vgg = models.resnet34(pretrained=True)
elif model_arch == 'resnet50':
gs_vgg = models.resnet50(pretrained=True)
elif model_arch == 'resnet101':
gs_vgg = models.resnet101(pretrained=True)
elif model_arch == 'resnet152':
gs_vgg = models.resnet152(pretrained=True)
elif model_arch == 'squeezenet1_0':
gs_vgg = models.squeezenet1_0(pretrained=True)
elif model_arch == 'squeezenet1_1':
gs_vgg = models.squeezenet1_1(pretrained=True)
elif model_arch == 'densenet121':
gs_vgg = models.densenet121(pretrained=True)
elif model_arch == 'densenet169':
gs_vgg = models.densenet169(pretrained=True)
elif model_arch == 'densenet161':
gs_vgg = models.densenet161(pretrained=True)
elif model_arch == 'densenet201':
gs_vgg = models.densenet201(pretrained=True)
elif model_arch == 'inception_v3':
gs_vgg = models.inception_v3(pretrained=True)
elif model_arch == 'googlenet':
gs_vgg = models.googlenet(pretrained=True)
elif model_arch == 'shufflenet_v2_x0_5':
gs_vgg = models.shufflenet_v2_x0_5(pretrained=True)
elif model_arch == 'shufflenet_v2_x1_0':
gs_vgg = models.shufflenet_v2_x1_0(pretrained=True)
elif model_arch == 'shufflenet_v2_x1_5':
gs_vgg = models.shufflenet_v2_x1_5(pretrained=True)
elif model_arch == 'shufflenet_v2_x2_0':
gs_vgg = models.shufflenet_v2_x2_0(pretrained=True)
elif model_arch == 'mobilenet_v2':
gs_vgg = models.mobilenet_v2(pretrained=True)
elif model_arch == 'resnext50_32x4d':
gs_vgg = models.resnext50_32x4d(pretrained=True)
elif model_arch == 'resnext101_32x8d':
gs_vgg = models.resnext101_32x8d(pretrained=True)
elif model_arch == 'wide_resnet50_2':
gs_vgg = models.wide_resnet50_2(pretrained=True)
elif model_arch == 'wide_resnet101_2':
gs_vgg = models.wide_resnet101_2(pretrained=True)
elif model_arch == 'mnasnet0_5':
gs_vgg = models.mnasnet0_5(pretrained=True)
elif model_arch == 'mnasnet0_75':
gs_vgg = models.mnasnet0_75(pretrained=True)
elif model_arch == 'mnasnet1_0':
gs_vgg = models.mnasnet1_0(pretrained=True)
elif model_arch == 'mnasnet1_3':
gs_vgg = models.mnasnet1_3(pretrained=True)
epoch_tot = 0
for parameters in gs_vgg.parameters():
parameters.requires_grad = False
try:
input_layer = gs_vgg.classifier[0].in_features
except:
input_layer = gs_vgg.classifier.in_features
hidden_layers = [(int(hidden_units * 0.68)), (int(hidden_units * 0.32))]
output_layer = len(train_loader)
gs_vgg.classifier = nn.Sequential(
nn.Linear(input_layer, hidden_layers[0]),
nn.ReLU(),
nn.Dropout(p=0.3),
nn.Linear(hidden_layers[0], hidden_layers[1]),
nn.ReLU(),
#nn.Linear(hidden_layers[1], output_layer),
nn.Linear(hidden_layers[1], 102),
nn.LogSoftmax(dim=1))
criterion = nn.NLLLoss()
optimizer = optim.Adam(gs_vgg.classifier.parameters(), lr=learning_rate)
gs_vgg.to(device)
step_num = 0
epoch = epochs
running_loss = 0
print_every = 10
for e in range(epoch):
epoch_tot += 1
for images, labels in train_loader:
gs_vgg.train()
images, labels = images.to(device), labels.to(device)
optimizer.zero_grad()
output = gs_vgg.forward(images)
loss = criterion(output, labels)
loss.backward()
optimizer.step()
running_loss += loss.item()
step_num += 1
if step_num % print_every == 0:
test_loss = 0
accuracy = 0
gs_vgg.eval()
for images, labels in valid_loader:
images, labels = images.to(device), labels.to(device)
output = gs_vgg.forward(images)
loss = criterion(output, labels)
test_loss += loss.item()
prob = torch.exp(output)
top_p, top_class = prob.topk(1, dim=1)
equals = top_class == labels.view(*top_class.shape)
accuracy += torch.mean(equals.type(
torch.FloatTensor)).item()
print(
f"Total Epochs: {epoch_tot}.. "
f"Train loss: {running_loss/print_every:.3f}.. "
f"Test loss: {test_loss/len(valid_loader):.3f}.. "
f"Test accuracy: {(accuracy/len(valid_loader))*100:.1f}%")
running_loss = 0
gs_vgg.train()
gs_vgg.class_to_idx = train_datasets.class_to_idx
gs_checkpoint = {
'epoch_tot': epoch_tot,
'model': gs_vgg,
'criterion': criterion,
'optimizer': optimizer,
'optimizer.state_dict': optimizer.state_dict(),
'model.state_dict': gs_vgg.state_dict(),
'model.class_to_idx': gs_vgg.class_to_idx
}
torch.save(gs_checkpoint, save_dir)
return gs_checkpoint
def get_model(name, device):
if name == "normal_cnn":
model = Net().to(device)
return model
elif name == "alexnet":
model = models.alexnet().to(device)
num_features = model.classifier[6].in_features
model.classifier[6] = nn.Linear(num_features, 1).to(device)
return model
elif name == "vgg11":
model = models.vgg11().to(device)
num_features = model.classifier[6].in_features
model.classifier[6] = nn.Linear(num_features, 1).to(device)
return model
elif name == "vgg13":
model = models.vgg13().to(device)
num_features = model.classifier[6].in_features
model.classifier[6] = nn.Linear(num_features, 1).to(device)
return model
elif name == "vgg16":
model = models.vgg16().to(device)
num_features = model.classifier[6].in_features
model.classifier[6] = nn.Linear(num_features, 1).to(device)
return model
elif name == "vgg19":
model = models.vgg19().to(device)
num_features = model.classifier[6].in_features
model.classifier[6] = nn.Linear(num_features, 1).to(device)
return model
elif name == "vgg11_bn":
model = models.vgg11_bn().to(device)
num_features = model.classifier[6].in_features
model.classifier[6] = nn.Linear(num_features, 1).to(device)
return model
elif name == "vgg13_bn":
model = models.vgg13_bn().to(device)
num_features = model.classifier[6].in_features
model.classifier[6] = nn.Linear(num_features, 1).to(device)
return model
elif name == "vgg16_bn":
model = models.vgg16_bn().to(device)
num_features = model.classifier[6].in_features
model.classifier[6] = nn.Linear(num_features, 1).to(device)
return model
elif name == "vgg19_bn":
model = models.vgg19_bn().to(device)
num_features = model.classifier[6].in_features
model.classifier[6] = nn.Linear(num_features, 1).to(device)
return model
elif name == "densenet121":
model = models.densenet121().to(device)
num_features = model.classifier.in_features
model.classifier = nn.Linear(num_features, 1).to(device)
return model
elif name == "densenet161":
model = models.densenet161().to(device)
num_features = model.classifier.in_features
model.classifier = nn.Linear(num_features, 1).to(device)
return model
elif name == "densenet169":
model = models.densenet169().to(device)
num_features = model.classifier.in_features
model.classifier = nn.Linear(num_features, 1).to(device)
return model
elif name == "densenet201":
model = models.densenet201().to(device)
num_features = model.classifier.in_features
model.classifier = nn.Linear(num_features, 1).to(device)
return model
elif name == "resnet18":
model = models.resnet18().to(device)
num_features = model.fc.in_features
model.fc = nn.Linear(num_features, 1).to(device)
return model
elif name == "resnet34":
model = models.resnet34().to(device)
num_features = model.fc.in_features
model.fc = nn.Linear(num_features, 1).to(device)
return model
elif name == "resnet50":
model = models.resnet50().to(device)
num_features = model.fc.in_features
model.fc = nn.Linear(num_features, 1).to(device)
return model
elif name == "resnet101":
model = models.resnet101().to(device)
num_features = model.fc.in_features
model.fc = nn.Linear(num_features, 1).to(device)
return model
elif name == "resnet152":
model = models.resnet152().to(device)
num_features = model.fc.in_features
model.fc = nn.Linear(num_features, 1).to(device)
return model
elif name == "googlenet":
model = models.googlenet(aux_logits=False).to(device)
num_features = model.fc.in_features
model.fc = nn.Linear(num_features, 1).to(device)
return model
elif name == "inception_v3":
model = models.inception_v3(aux_logits=False).to(device)
num_features = model.fc.in_features
model.fc = nn.Linear(num_features, 1).to(device)
return model
elif name == "shufflenet_v2_x0_5":
model = models.shufflenet_v2_x0_5().to(device)
num_features = model.fc.in_features
model.fc = nn.Linear(num_features, 1).to(device)
return model
elif name == "shufflenet_v2_x1_0":
model = models.shufflenet_v2_x1_0().to(device)
num_features = model.fc.in_features
model.fc = nn.Linear(num_features, 1).to(device)
return model
elif name == "shufflenet_v2_x1_5":
model = models.shufflenet_v2_x1_5().to(device)
num_features = model.fc.in_features
model.fc = nn.Linear(num_features, 1).to(device)
return model
elif name == "shufflenet_v2_x2_0":
model = models.shufflenet_v2_x2_0().to(device)
num_features = model.fc.in_features
model.fc = nn.Linear(num_features, 1).to(device)
return model
elif name == "mobilenet_v2":
model = models.mobilenet_v2().to(device)
num_features = model.classifier[1].in_features
model.classifier[1] = nn.Linear(num_features, 1).to(device)
return model
elif name == "resnext50_32x4d":
model = models.resnext50_32x4d().to(device)
num_features = model.fc.in_features
model.fc = nn.Linear(num_features, 1).to(device)
return model
elif name == "resnext101_32x8d":
model = models.resnext101_32x8d().to(device)
num_features = model.fc.in_features
model.fc = nn.Linear(num_features, 1).to(device)
return model
elif name == "wide_resnet50_2":
model = models.wide_resnet50_2().to(device)
num_features = model.fc.in_features
model.fc = nn.Linear(num_features, 1).to(device)
return model
elif name == "wide_resnet101_2":
model = models.wide_resnet101_2().to(device)
num_features = model.fc.in_features
model.fc = nn.Linear(num_features, 1).to(device)
return model
elif name == "mnasnet0_5":
model = models.mnasnet0_5().to(device)
num_features = model.classifier[1].in_features
model.classifier[1] = nn.Linear(num_features, 1).to(device)
return model
elif name == "mnasnet0_75":
model = models.mnasnet0_75().to(device)
num_features = model.classifier[1].in_features
model.classifier[1] = nn.Linear(num_features, 1).to(device)
return model
elif name == "mnasnet1_0":
model = models.mnasnet1_0().to(device)
num_features = model.classifier[1].in_features
model.classifier[1] = nn.Linear(num_features, 1).to(device)
return model
elif name == "mnasnet1_3":
model = models.mnasnet1_3().to(device)
num_features = model.classifier[1].in_features
model.classifier[1] = nn.Linear(num_features, 1).to(device)
return model
def modeldict(mmd):
if mmd == 'alexnet_pt': out_model = models.alexnet(pretrained=True)
elif mmd == 'alexnet': out_model = models.alexnet(pretrained=False)
elif mmd == 'vgg11_pt': out_model = models.vgg11(pretrained=True)
elif mmd == 'vgg11': out_model = models.vgg11(pretrained=False)
elif mmd == 'vgg11bn_pt': out_model = models.vgg11_bn(pretrained=True)
elif mmd == 'vgg11bn': out_model = models.vgg11_bn(pretrained=False)
elif mmd == 'vgg13_pt': out_model = models.vgg13(pretrained=True)
elif mmd == 'vgg13': out_model = models.vgg13(pretrained=False)
elif mmd == 'vgg13bn_pt': out_model = models.vgg13_bn(pretrained=True)
elif mmd == 'vgg13bn': out_model = models.vgg13_bn(pretrained=False)
elif mmd == 'vgg16_pt': out_model = models.vgg16(pretrained=True)
elif mmd == 'vgg16': out_model = models.vgg16(pretrained=False)
elif mmd == 'vgg16bn_pt': out_model = models.vgg16_bn(pretrained=True)
elif mmd == 'vgg16bn': out_model = models.vgg16_bn(pretrained=False)
elif mmd == 'vgg19_pt': out_model = models.vgg19(pretrained=True)
elif mmd == 'vgg19': out_model = models.vgg19(pretrained=False)
elif mmd == 'vgg19bn_pt': out_model = models.vgg19_bn(pretrained=True)
elif mmd == 'vgg19bn': out_model = models.vgg19_bn(pretrained=False)
elif mmd == 'resnet18': out_model = models.resnet18(pretrained=False)
elif mmd == 'resnet18_pt': out_model = models.resnet18(pretrained=True)
elif mmd == 'resnet34': out_model = models.resnet34(pretrained=False)
elif mmd == 'resnet34_pt': out_model = models.resnet34(pretrained=True)
elif mmd == 'resnet50': out_model = models.resnet50(pretrained=False)
elif mmd == 'resnet50_pt': out_model = models.resnet50(pretrained=True)
elif mmd == 'resnet101': out_model = models.resnet101(pretrained=False)
elif mmd == 'resnet101_pt': out_model = models.resnet101(pretrained=True)
elif mmd == 'resnet152': out_model = models.resnet152(pretrained=False)
elif mmd == 'resnet152_pt': out_model = models.resnet152(pretrained=True)
elif mmd == 'squeezenet10_pt':
out_model = models.squeezenet1_0(pretrained=True)
elif mmd == 'squeezenet10':
out_model = models.squeezenet1_0(pretrained=False)
elif mmd == 'squeezenet11_pt':
out_model = models.squeezenet1_1(pretrained=True)
elif mmd == 'squeezenet11':
out_model = models.squeezenet1_1(pretrained=False)
elif mmd == 'densenet121_pt':
out_model = models.densenet121(pretrained=True)
elif mmd == 'densenet121':
out_model = models.densenet121(pretrained=False)
elif mmd == 'densenet161_pt':
out_model = models.densenet161(pretrained=True)
elif mmd == 'densenet161':
out_model = models.densenet161(pretrained=False)
elif mmd == 'densenet169_pt':
out_model = models.densenet169(pretrained=True)
elif mmd == 'densenet169':
out_model = models.densenet169(pretrained=False)
elif mmd == 'densenet201_pt':
out_model = models.densenet201(pretrained=True)
elif mmd == 'densenet201':
out_model = models.densenet201(pretrained=False)
elif mmd == 'inception_pt':
out_model = models.inception_v3(pretrained=True, aux_logits=False)
elif mmd == 'inception':
out_model = models.inception_v3(pretrained=False, aux_logits=False)
elif mmd == 'inceptionv4_pt':
out_model = pretrainedmodels.__dict__['inceptionv4'](
num_classes=1000, pretrained='imagenet')
elif mmd == 'inceptionv4':
out_model = pretrainedmodels.__dict__['inceptionv4'](num_classes=1000,
pretrained=None)
elif mmd == 'inceptionresnetv2_pt':
out_model = pretrainedmodels.__dict__['inceptionresnetv2'](
num_classes=1000, pretrained='imagenet')
elif mmd == 'inceptionresnetv2':
out_model = pretrainedmodels.__dict__['inceptionresnetv2'](
num_classes=1000, pretrained=None)
elif mmd == 'googlenet_pt':
out_model = models.googlenet(pretrained=True, aux_logits=False)
elif mmd == 'googlenet':
out_model = models.googlenet(pretrained=False, aux_logits=False)
elif mmd == 'shufflenet05_pt':
out_model = models.shufflenet_v2_x0_5(pretrained=True)
elif mmd == 'shufflenet05':
out_model = models.shufflenet_v2_x0_5(pretrained=False)
elif mmd == 'shufflenet10_pt':
out_model = models.shufflenet_v2_x1_0(pretrained=True)
elif mmd == 'shufflenet10':
out_model = models.shufflenet_v2_x1_0(pretrained=False)
elif mmd == 'shufflenet20':
out_model = models.shufflenet_v2_x2_0(pretrained=False)
elif mmd == 'mobilenet_pt':
out_model = models.mobilenet_v2(pretrained=True)
elif mmd == 'mobilenet':
out_model = models.mobilenet_v2(pretrained=False)
elif mmd == 'resnext50_32x4d_pt':
out_model = models.resnext50_32x4d(pretrained=True)
elif mmd == 'resnext50_32x4d':
out_model = models.resnext50_32x4d(pretrained=False)
elif mmd == 'resnext101_32x8d_pt':
out_model = models.resnext101_32x8d(pretrained=True)
elif mmd == 'resnext101_32x8d':
out_model = models.resnext101_32x8d(pretrained=False)
elif mmd == 'wide_resnet50_2_pt':
out_model = models.wide_resnet50_2(pretrained=True)
elif mmd == 'wide_resnet50_2':
out_model = models.wide_resnet50_2(pretrained=False)
elif mmd == 'wide_resnet101_2_pt':
out_model = models.wide_resnet101_2(pretrained=True)
elif mmd == 'wide_resnet101_2':
out_model = models.wide_resnet101_2(pretrained=False)
elif mmd == 'mnasnet05_pt':
out_model = models.mnasnet0_5(pretrained=True)
elif mmd == 'mnasnet05':
out_model = models.mnasnet0_5(pretrained=False)
elif mmd == 'mnasnet075':
out_model = models.mnasnet0_75(pretrained=False)
elif mmd == 'mnasnet10_pt':
out_model = models.mnasnet1_0(pretrained=True)
elif mmd == 'mnasnet10':
out_model = models.mnasnet1_0(pretrained=False)
elif mmd == 'mnasnet13':
out_model = models.mnasnet1_3(pretrained=False)
elif mmd == 'xception_pt':
out_model = pretrainedmodels.__dict__['xception'](
num_classes=1000, pretrained='imagenet')
elif mmd == 'xception':
out_model = pretrainedmodels.__dict__['xception'](num_classes=1000,
pretrained=None)
elif mmd == 'dpn68_pt':
out_model = pretrainedmodels.__dict__['dpn68'](num_classes=1000,
pretrained='imagenet')
elif mmd == 'dpn68':
out_model = pretrainedmodels.__dict__['dpn68'](num_classes=1000,
pretrained=None)
elif mmd == 'dpn98_pt':
out_model = pretrainedmodels.__dict__['dpn98'](num_classes=1000,
pretrained='imagenet')
elif mmd == 'dpn98':
out_model = pretrainedmodels.__dict__['dpn98'](num_classes=1000,
pretrained=None)
elif mmd == 'dpn131_pt':
out_model = pretrainedmodels.__dict__['dpn131'](num_classes=1000,
pretrained='imagenet')
elif mmd == 'dpn131':
out_model = pretrainedmodels.__dict__['dpn131'](num_classes=1000,
pretrained=None)
elif mmd == 'dpn68b_pt':
out_model = pretrainedmodels.__dict__['dpn68b'](
num_classes=1000, pretrained='imagenet+5k')
elif mmd == 'dpn68b':
out_model = pretrainedmodels.__dict__['dpn68b'](num_classes=1000,
pretrained=None)
elif mmd == 'dpn92_pt':
out_model = pretrainedmodels.__dict__['dpn92'](
num_classes=1000, pretrained='imagenet+5k')
elif mmd == 'dpn92':
out_model = pretrainedmodels.__dict__['dpn92'](num_classes=1000,
pretrained=None)
elif mmd == 'dpn107_pt':
out_model = pretrainedmodels.__dict__['dpn107'](
num_classes=1000, pretrained='imagenet+5k')
elif mmd == 'dpn107':
out_model = pretrainedmodels.__dict__['dpn107'](num_classes=1000,
pretrained=None)
elif mmd == 'fbresnet152_pt':
out_model = pretrainedmodels.__dict__['fbresnet152'](
num_classes=1000, pretrained='imagenet')
elif mmd == 'fbresnet152':
out_model = pretrainedmodels.__dict__['fbresnet152'](num_classes=1000,
pretrained=None)
else:
out_model = None
print('Invalid model name. Terminated.')
exit(0)
return out_model
import matplotlib.pyplot as plt
import numpy as np
import torch
from torchvision.models import alexnet, vgg11, mobilenet, mnasnet0_5, resnet18, resnet34, mobilenet_v2, densenet121, shufflenet_v2_x0_5, squeezenet1_0, vgg19
from ai.ai import ArithmeticIntensity
from thop import profile
import timeit
models = [
alexnet(),
resnet18(),
resnet34(),
mobilenet_v2(),
densenet121(),
shufflenet_v2_x0_5(),
squeezenet1_0(),
vgg19()
]
model_names = [
alexnet.__name__, resnet18.__name__, resnet34.__name__,
mobilenet_v2.__name__, densenet121.__name__, shufflenet_v2_x0_5.__name__,
squeezenet1_0.__name__, vgg19.__name__
]
metrics = []
for model in models:
ai_profiler = ArithmeticIntensity(model=model, input_dims=(1, 3, 224, 224))
ai, macs = ai_profiler.get_metrics()
tensor = torch.randn(5, 3, 224, 224)
_, params = profile(model, inputs=(tensor, ))
def __init__(self):
super(ShuffleneTiny, self).__init__()
self.model = models.shufflenet_v2_x0_5()
self.model.fc = nn.Sequential(nn.BatchNorm1d(1024), nn.Linear(1024, 3))
def shufflenetv2(num_classes, pretrained=True, small=False):
net = models.shufflenet_v2_x0_5(pretrained) if small else models.shufflenet_v2_x1_0(pretrained)
net.fc = nn.Linear(1024, num_classes)
return net
def feature_extractor(cls):
shuffle = models.shufflenet_v2_x0_5(pretrained=True)
return nn.Sequential(*list(shuffle.children())[:-1])
def shufflenet():
sf1=models.shufflenet_v2_x0_5()
sf1.conv1[0]=nn.Conv2d(1, 24, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)
sf1.fc=nn.Linear(1024,1,True)
return sf1.to('cuda')
def __init__(self, model_type=8, cnn_type='x0_5', cnn_layer_test=0):
super(LightWeight_feat, self).__init__()
self.model_type = model_type
self.cnn_type = cnn_type
self.use_layer = cnn_layer_test
self.bn = True # BenchNorm
# select cnn_model from SqueezeNet, MobileNet, ShuffleNet, MnasNet
if self.model_type == LightCNN.SqueezeNet.value: # use SqueezeNet
if cnn_type == '1_0': # V1_1_0
self.model = models.squeezenet1_0(pretrained=True)
self.sel_feat_layer = {0, 4, 8,
12} # conv, conv2_2, conv3_2, cam
else: # V1_x1_1
self.cnn_type = '1_1'
self.model = models.squeezenet1_1(pretrained=True)
self.sel_feat_layer = {0, 4, 7,
12} # conv, conv2_2, conv3_2, cam
print('use SqueezeNet', self.cnn_type)
elif self.model_type == LightCNN.MobileNet.value: # use MobileNet
if cnn_type != 'Small':
self.cnn_type = 'Large'
self.model = MobileNet_feat(self.cnn_type, True,
1000) # Large, Small
elif self.model_type == LightCNN.ShuffleNet.value: # use ShuffleNet
if cnn_type == 'x0_5': # V2_x0.5
self.model = models.shufflenet_v2_x0_5(pretrained=True)
else: # V2_x1.0
self.cnn_type = 'x1_0'
self.model = models.shufflenet_v2_x1_0(pretrained=True)
print('use ShuffleV2', self.cnn_type)
elif self.model_type == LightCNN.MnasNet.value: # use MnasNet
if cnn_type == '0_5': # MnasNet_0_5
self.model = models.mnasnet0_5(pretrained=True)
else: # MnasNet_1_0
self.cnn_type = '1_0'
self.model = models.mnasnet1_0(pretrained=True)
self.sel_feat_layer = {6, 8, 10, 16} # conv1_3, conv2, conv4, cam
print('use MnasNet', self.cnn_type)
self.model.eval()
# using Grad-CAM with softmax weight (FC)
params = list(self.model.parameters())
self.weight_softmax = np.squeeze(params[-2].data.numpy())
if use_gpu:
print('use cuda model')
self.model = self.model.cuda()
if self.model_type == LightCNN.SqueezeNet.value:
self.features = nn.ModuleList(list(self.model.features)).eval()
elif self.model_type == LightCNN.MnasNet.value:
self.features = nn.ModuleList(list(self.model.layers)).eval()
# test feature
if self.use_layer == 0: # run combine cnn_features
self.use_layer = -4
print('test linear combined cnn_features')
elif self.use_layer > 4: # run combine features with grad-CAM
print('test combined cnn_features with grad-CAM')
else: # test each conv layers
print('test single conv layer', self.use_layer)
logging.info("VERSION %s" % pytorch_metric_learning.__version__)
##############################
########## Training ##########
##############################
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# In this example, we'll take multiple trunks and embedders,
# and encapsulate them into a single trunk and embedder model.
# Note that CascadedEmbeddings does not necessarily require a complicated setup like this.
# CascadedEmbeddings just assumes that the output of your embedder
# should be partitioned into different sections, as specified by the init argument
# "embedding_sizes".
trunk1 = models.shufflenet_v2_x0_5(pretrained=True)
trunk2 = models.shufflenet_v2_x1_0(pretrained=True)
trunk3 = models.resnet18(pretrained=True)
all_trunks = [trunk1, trunk2, trunk3]
trunk_output_sizes = []
for T in all_trunks:
trunk_output_sizes.append(T.fc.in_features)
T.fc = Identity()
trunk = ListOfModels(all_trunks)
trunk = torch.nn.DataParallel(trunk.to(device))
# Set the embedders. Each embedder takes a corresponding trunk model output, and outputs 64-dim embeddings.
all_embedders = []
for s in trunk_output_sizes:
