def __init__(self, device=None, jit=False):
super().__init__()
self.device = device
self.jit = jit
self.model = models.shufflenet_v2_x1_0().to(self.device)
self.eval_model = models.shufflenet_v2_x1_0().to(self.device)
self.example_inputs = (torch.randn((32, 3, 224, 224)).to(self.device),)
if self.jit:
self.model = torch.jit.script(self.model, example_inputs=[self.example_inputs, ])
self.eval_model = torch.jit.script(self.eval_model)
# model needs to in `eval`
# in order to be optimized for inference
self.eval_model.eval()
self.eval_model = torch.jit.optimize_for_inference(self.eval_model)
def build_model(self, n_classes):
if self.model_name == 'resnet18':
model_ft = models.resnet18(pretrained=self.img_init)
elif self.model_name == 'shufflenet':
model_ft = models.shufflenet_v2_x1_0(pretrained=self.img_init)
elif self.model_name == 'resnet34':
model_ft = models.resnet34(pretrained=self.img_init)
else:
raise ValueError('unrecognized model version: {}'.format(
self.model_name))
num_ftrs = model_ft.fc.in_features
model_ft.fc = nn.Linear(num_ftrs, n_classes)
layers = self.recursively_enumerate_model(model_ft)
layers = [
layer for layer in layers
if (type(layer) != nn.BatchNorm2d
and len(list(layer.parameters())) > 0)
]
for layer in layers[:round(self.fp * len(layers))]:
self.freeze_weights(layer)
# raise ValueError('incorrect mode setting: {}'.format(mode))
return model_ft
def get_example_params(example_index):
"""
Gets used variables for almost all visualizations, like the image, model etc.
Args:
example_index (int): Image id to use from examples
returns:
original_image (numpy arr): Original image read from the file
prep_img (numpy_arr): Processed image
target_class (int): Target class for the image
file_name_to_export (string): File name to export the visualizations
pretrained_model(Pytorch model): Model to use for the operations
"""
# Pick one of the examples
example_list = (('../input_images/snake.jpg',
56), ('../input_images/dog2.jpg', 243),
('../input_images/spider.png', 72))
img_path = example_list[example_index][0]
target_class = example_list[example_index][1]
file_name_to_export = img_path[img_path.rfind('/') + 1:img_path.rfind('.')]
# Read image
original_image = Image.open(img_path).convert('RGB')
# Process image
prep_img = preprocess_image(original_image)
# Define model
#pretrained_model = models.alexnet(pretrained=True)
pretrained_model = models.shufflenet_v2_x1_0(pretrained=True)
return (original_image, prep_img, target_class, file_name_to_export,
pretrained_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 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 pretrained_model(name):
if name == "resnet18":
return models.resnet18(pretrained=True)
elif name == "alexnet":
return models.alexnet(pretrained=True)
elif name == "squeezenet":
return models.squeezenet1_0(pretrained=True)
elif name == "vgg16":
return models.vgg16(pretrained=True)
elif name == "densenet":
return models.densenet161(pretrained=True)
elif name == "inception":
return models.inception_v3(pretrained=True)
elif name == "googlenet":
return models.googlenet(pretrained=True)
elif name == "shufflenet":
return models.shufflenet_v2_x1_0(pretrained=True)
elif name == "mobilenet":
return models.mobilenet_v2(pretrained=True)
elif name == "resnext50_32x4d":
return models.resnext50_32x4d(pretrained=True)
elif name == "wide_resnet50_2":
return models.wide_resnet50_2(pretrained=True)
elif name == "mnasnet":
return models.mnasnet1_0(pretrained=True)
else:
print("model {} don't know".format(name))
def 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 main():
ONNX_SAVE_ROOT.mkdir(parents=True, exist_ok=True)
# VGG11_bn (2.11ms)
# gen_onnx(models.vgg11_bn().features, "VGG11_bn")
# ResNet18 (0.86ms)
gen_onnx(models.resnet18(pretrained=True), "ResNet18")
# Densenet121 (5.2ms)
# gen_onnx(models.densenet121(), "Densenet121")
# shufflenet v2 x1.0 (0.87ms)
gen_onnx(models.shufflenet_v2_x1_0(pretrained=True), "ShuffleNetV2_x1.0")
""" # hard sigmoid 미지원
# MobileNet V3 Large
gen_onnx(models.mobilenet_v3_large(), "MobileNetV3_large")
# MobileNet V3 Small
gen_onnx(models.mobilenet_v3_small(), "MobileNetV3_small")
"""
# ResNeXt-50-32x4d (3.15ms)
# gen_onnx(models.resnext50_32x4d(), "ResNeXt50_32x4d")
# Wide ResNet-50-2 (3.37ms)
# gen_onnx(models.wide_resnet50_2(), "WideResNet_50_2")
# MNASNet 1.0 (0.99ms)
gen_onnx(models.mnasnet1_0(pretrained=True), "MNASNet_1_0")
# MNASNet 0.5 (0.6ms)
gen_onnx(models.mnasnet0_5(pretrained=True), "MNASNet_0_5")
def get_model(type):
"""
type should be one of (resnext|densenet)
"""
num_labels = 17
if type == "resnext":
model = models.resnext50_32x4d(pretrained=True)
# Replace the classifier, all layers will be trained
model.fc = nn.Sequential(nn.Linear(model.fc.in_features, num_labels),
nn.LogSoftmax(dim=1))
return model
elif type == "densenet":
model = models.densenet201(pretrained=True)
# Replace the classifier, all layers will be trained
model.classifier = nn.Sequential(
nn.Linear(model.classifier.in_features, num_labels),
nn.LogSoftmax(dim=1))
return model
elif type == "shufflenet":
model = models.shufflenet_v2_x1_0(pretrained=True)
# Replace the classifier, all layers will be trained
model.fc = nn.Sequential(
nn.Linear(in_features=model.fc.in_features,
out_features=num_labels), nn.LogSoftmax(dim=1))
return model
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__(self, device="cpu", jit=False):
self.device = device
self.jit = jit
self.model = models.shufflenet_v2_x1_0()
if self.jit:
self.model = torch.jit.script(self.model)
self.example_inputs = (torch.randn((32, 3, 224, 224)),)
def initialize_model(model_name, feature_extract, use_pretrained=True):
"""Pre-trained model head"""
model_ft = None
input_size = 0
if model_name == 'resnet50':
model_ft = models.resnext50_32x4d(pretrained=use_pretrained)
set_parameter_requires_grad(model_ft, feature_extract)
input_size = 224
elif model_name == 'mobilenet_v2':
model_ft = models.mobilenet_v2(pretrained=use_pretrained)
set_parameter_requires_grad(model_ft, feature_extract)
input_size = 224
elif model_name == 'shufflenet_v2':
model_ft = models.shufflenet_v2_x1_0(pretrained=use_pretrained)
set_parameter_requires_grad(model_ft, feature_extract)
input_size = 224
elif model_name == 'mixnet':
model_ft = timm.create_model('tf_mixnet_s', pretrained=True)
set_parameter_requires_grad(model_ft, feature_extract)
input_size = 224
elif model_name == 'facenet':
model_ft = InceptionResnetV1(pretrained='vggface2')
set_parameter_requires_grad(model_ft, feature_extract)
input_size = 160
else:
raise NameError("No such pre-trained model...")
return model_ft, input_size
def __init__(self, config):
self.model = models.shufflenet_v2_x1_0(pretrained=True)
# # freeze model parameters
# for param in self.model.parameters():
# param.requires_grad = False
self.model.fc = nn.Sequential(nn.Linear(1024, config.class_num),
nn.Sigmoid())
# for param in self.model.fc.parameters():
# param.requires_grad = True
# # model check
# print(self.model)
# for name, param in self.model.named_parameters():
# if param.requires_grad:
# print("requires_grad: True ", name)
# else:
# print("requires_grad: False ", name)
self.device = torch.device("cpu")
if torch.cuda.is_available():
self.device = torch.device("cuda:%i" % config.device[0])
self.model = torch.nn.DataParallel(self.model,
device_ids=config.device)
self.model = self.model.to(self.device)
self.lr = config.lr
self.weight_decay = config.weight_decay
self.epoch = config.epoch
self.splits = config.n_splits
self.root = config.root
self.solver = Solver(self.model, self.device)
self.criterion = nn.CrossEntropyLoss()
self.TIME = "{0:%Y-%m-%dT%H-%M-%S}-classify".format(
datetime.datetime.now())
self.model_path = os.path.join(config.root, config.save_path,
config.model_name, self.TIME)
if not os.path.exists(self.model_path):
os.makedirs(self.model_path)
self.max_accuracy_valid = 0
self.seed = int(time.time())
# self.seed = 1570421136
seed_torch(self.seed)
self.train_transform = transforms.Compose([
transforms.Resize([256, 256]),
transforms.RandomCrop(224),
transforms.RandomRotation(degrees=(-40, 40)),
transforms.RandomHorizontalFlip(),
transforms.ToTensor()
])
self.test_transform = transforms.Compose(
[transforms.Resize(224),
transforms.ToTensor()])
def get_shufflenet_v2_x1_0(class_num):
model = models.shufflenet_v2_x1_0(pretrained=True)
set_parameter_requires_grad(model)
model.name = 'shufflenet_v2_x1_0'
n_inputs = model.fc.in_features
model.fc = nn.Linear(n_inputs, class_num)
return model, 244
def __init__(self, pretrained=True):
super().__init__()
model = models.shufflenet_v2_x1_0(pretrained=pretrained)
if pretrained:
for param in model.parameters():
param.requires_grad = False
# create new model by removing the last layer
self.model = torch.nn.Sequential(*(list(model.children())[:-1][:-2]))
def shufflenet_v2_x1_0_nc_1k():
m = shufflenet_v2_x1_0(pretrained=True)
m.fc = nn.Identity()
del m.conv5[2], m.conv5[1] # remove ReLU and BN
m.conv5[0] = nn.Conv2d(464,
1000,
kernel_size=(1, 1),
stride=(1, 1),
bias=True)
return m
def shufflenet_v2_x10(pretrained=False, imagenet_weight=False):
"""Constructs a Shufflenet-V2_x10 model.
Args:
pretrained (bool): If True, returns a model pre-trained on ImageNet
"""
print('=== Using Shufflenet-V2_x10 ===')
model = models.shufflenet_v2_x1_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')
model = models.shufflenet_v2_x1_0(pretrained=True)
else:
print('=== train from scratch ===')
return model
def get_model(model_name):
"""
Get specific modified pre-trained models by name
"""
model = None
if model_name == "AlexNet":
model = models.alexnet(pretrained=True)
in_features = model.fc.in_features
model.fc = nn.Linear(in_features, 26)
elif model_name == "VGG":
model = models.vgg16(pretrained=True)
in_features = model.fc.in_features
model.fc = nn.Linear(in_features, 26)
elif model_name == "ResNet":
model = models.resnet18(pretrained=True)
in_features = model.fc.in_features
model.fc = nn.Linear(in_features, 26)
elif model_name == "SqueezNet":
model = models.squeezenet1_1(pretrained=True)
in_features = model.fc.in_features
model.fc = nn.Linear(in_features, 26)
elif model_name == "DenseNet":
model = models.densenet121(pretrained=True)
in_features = model.classifier.in_features
model.classifier = nn.Linear(in_features, 26)
elif model_name == "Inception":
model = models.inception_v3(pretrained=True)
in_features = model.fc.in_features
model.fc = nn.Linear(in_features, 26)
elif model_name == "GoogleNet":
model = models.googlenet(pretrained=True)
in_features = model.fc.in_features
model.fc = nn.Linear(in_features, 26)
elif model_name == "ShuffleNet":
model = models.shufflenet_v2_x1_0(pretrained=True)
in_features = model.fc.in_features
model.fc = nn.Linear(in_features, 26)
elif model_name == "MobileNet":
model = models.mobilenet_v2(pretrained=True)
in_features = model.fc.in_features
model.fc = nn.Linear(in_features, 26)
elif model_name == "ResNext":
model = models.resnext101_32x8d(pretrained=True)
in_features = model.fc.in_features
model.fc = nn.Linear(in_features, 26)
elif model_name == "WResNet":
model = models.wide_resnet101_2(pretrained=True)
in_features = model.fc.in_features
model.fc = nn.Linear(in_features, 26)
elif model_name == "MNASNet":
model = models.mnasnet1_0(pretrained=True)
in_features = model.fc.in_features
model.fc = nn.Linear(in_features, 26)
return model
def __init__(self,
det_threshold,
checkpoint_path,
device,
use_apex=False,
full_precision=False,
top_k=750,
nms_threshold=0.4):
"""
parameters:
config: dict (contains training/inference backbone`s parameters)
det_threshold: float (threshold for bbox`s confidence)
checkpoint_path: str (path to pretrained model)
device: str (cpu or gpu inference)
use_apex: bool (need only for not full precision and inference optimization)
full_precision: bool ()
top_k: int (max for detections)
nms_threshold: int (threshold for non-maximum supression)
"""
self.device_str = device
self.device = torch.device(device)
self.top_k = top_k
self.nms_threshold = nms_threshold
self.det_threshold = det_threshold
torch.set_grad_enabled(False)
model = shufflenet_v2_x1_0(pretrained=False)
self.model = RetinaFace(model, 'inference')
self.__load_ckpt(checkpoint_path)
cudnn.enabled = True
cudnn.benchmark = True
if use_apex:
from apex import amp
self.model = amp.initialize(
model.to(device),
None,
opt_level='O0' if full_precision else 'O2',
keep_batchnorm_fp32=True,
verbosity=0)
self.model = self.model.to(self.device)
self.model.eval()
self.priorbox = None
self.priors = None
self.prior_data = None
self.width, self.height = None, None
self.scale, self.second_scale = None, None
def get_net(net_name, weight_path=None):
"""
根据网络名称获取模型
:param net_name: 网络名称
:param weight_path: 与训练权重路径
:return:
"""
pretrain = weight_path is None # 没有指定权重路径，则加载默认的预训练权重
if net_name in ['vgg16']:
net = models.vgg16(pretrained=pretrain)
elif net_name == 'vgg19':
net = models.vgg19(pretrained=pretrain)
elif net_name in ['resnet18']:
net = models.resnet18(pretrained=pretrain)
elif net_name in ['frcnn']:
net = models.frcnn(pretrained=pretrain)
elif net_name in ['resnet50']:
net = models.resnet50(pretrained=pretrain)
elif net_name == 'resnet101':
net = models.resnet101(pretrained=pretrain)
elif net_name in ['densenet121']:
net = models.densenet121(pretrained=pretrain)
elif net_name in ['densenet169']:
net = models.densenet169(pretrained=pretrain)
elif net_name in ['inception']:
net = models.inception_v3(pretrained=pretrain)
elif net_name in ['mobilenet_v2']:
net = models.mobilenet_v2(pretrained=pretrain)
elif net_name in ['shufflenet_v2']:
net = models.shufflenet_v2_x1_0(pretrained=pretrain)
elif net_name == 'efficientnet':
net = EfficientNet.from_name('efficientnet-b0')
feature = net._fc.in_features
net._fc = nn.Linear(in_features=feature, out_features=2,
bias=True) # 修改分类层结构
net.load_state_dict(torch.load("efficientNet-b0.pt")) # 加载CT集上训练参数
else:
raise ValueError('invalid network name:{}'.format(net_name))
# 加载指定路径的权重参数
if weight_path is not None and net_name.startswith('densenet'):
pattern = re.compile(
r'^(.*denselayer\d+\.(?:norm|relu|conv))\.((?:[12])\.(?:weight|bias|running_mean|running_var))$'
)
state_dict = torch.load(weight_path)
for key in list(state_dict.keys()):
res = pattern.match(key)
if res:
new_key = res.group(1) + res.group(2)
state_dict[new_key] = state_dict[key]
del state_dict[key]
net.load_state_dict(state_dict)
elif weight_path is not None:
net.load_state_dict(torch.load(weight_path))
return net
def __init__(self, n_classes=3):
super(LSTMDeepClassification, self).__init__()
# backbone = models.resnet34(pretrained=True)
backbone = models.shufflenet_v2_x1_0(pretrained=True)
layers = list(backbone.children())
self.backbone = nn.Sequential(*layers[:-1])
# self.backbone = backbone.features
# self.pool = nn.AdaptiveAvgPool2d((1, 1))
self.embedding_dim = layers[-1].in_features
self.lstm = nn.LSTM(self.embedding_dim, 512, num_layers=2)
self.fc = nn.Linear(512, n_classes)
def _get_network_model(self, model_name, class_num, device):
return {
'CNN': ConvNetwork(num_class=class_num).to(device),
'CustomResNet': ResNet('imagenet', 18, class_num).to(device),
'RES18': torch_models.resnet18().to(device),
'RES34': torch_models.resnet34().to(device),
'RES50': torch_models.resnet50().to(device),
'RES101': torch_models.resnet101().to(device),
'GoogleNet': torch_models.googlenet().to(device),
'AlexNet': torch_models.alexnet().to(device),
'ShuffleNet': torch_models.shufflenet_v2_x1_0().to(device)
}.get(model_name, None)
def __init__(self):
super(AAA, self).__init__()
import torchvision.models as models
model = models.shufflenet_v2_x1_0()
self.conv1 = model.conv1 #112x112 24 2s
self.maxpool = model.maxpool #56x56 24 4s
self.stage2 = model.stage2 #28x28 58 残差连接,,两个58 116 8s
self.stage3 = model.stage3 #14x14 232 16s
self.stage4 = model.stage4 #7x7 464 32s
self.deconv32s = nn.ConvTranspose2d(464,
232,
kernel_size=3,
stride=2,
padding=1,
output_padding=1) #变16s
self.bn32s = nn.BatchNorm2d(232)
self.deconv16s = nn.ConvTranspose2d(232,
116,
kernel_size=3,
stride=2,
padding=1,
output_padding=1) #变8s
self.bn16s = nn.BatchNorm2d(116)
self.deconv8s = nn.ConvTranspose2d(116,
24,
kernel_size=3,
stride=2,
padding=1,
output_padding=1) #变成4s
self.bn8s = nn.BatchNorm2d(24)
self.deconv4s = nn.ConvTranspose2d(24,
24,
kernel_size=3,
stride=2,
padding=1,
output_padding=1) #变2s
self.bn4s = nn.BatchNorm2d(24)
self.deconv2s = nn.ConvTranspose2d(24,
64,
kernel_size=3,
stride=2,
padding=1,
output_padding=1) #变1s
self.bn2s = nn.BatchNorm2d(64)
self.relu = nn.ReLU(inplace=True)
self.cls = nn.Conv2d(64, 2, kernel_size=1, stride=1)
def __init__(self, config):
self.model = models.shufflenet_v2_x1_0(pretrained=False)
self.model.fc = nn.Sequential(nn.Linear(1024, config.class_num),
nn.Sigmoid())
self.device = torch.device("cpu")
if torch.cuda.is_available():
self.device = torch.device("cuda:%i" % config.device[0])
# self.model = torch.nn.DataParallel(self.model, device_ids=config.device)
self.model = self.model.to(self.device)
self.weight_path = config.weight_path
self.solver = Solver(self.model, self.device)
def test1():
resnet18 = models.resnet18()
alexnet = models.alexnet()
vgg16 = models.vgg16()
squeezenet = models.squeezenet1_0()
densenet = models.densenet161()
inception = models.inception_v3()
googlenet = models.googlenet()
shufflenet = models.shufflenet_v2_x1_0()
mobilenet = models.mobilenet_v2()
resnext50_32x4d = models.resnext50_32x4d()
wide_resnet50_2 = models.wide_resnet50_2()
mnasnet = models.mnasnet1_0()
def create_model(model_name, n_classes=10, freeze_feature=True):
input_size = 224
if model_name == 'resnet18':
model = models.resnet18(pretrained=True)
set_parameter_requires_grad(model, freeze_feature)
n_feat_in = model.fc.in_features
model.fc = nn.Linear(n_feat_in, n_classes)
elif model_name == 'vgg19':
model = models.vgg19(pretrained=True)
set_parameter_requires_grad(model, freeze_feature)
model.classifier = nn.Sequential(nn.Linear(512 * 7 * 7, 1024),
nn.ReLU(True), nn.Dropout(),
nn.Linear(1024, n_classes))
elif model_name == 'vgg19_bn':
model = models.vgg19_bn(pretrained=True)
set_parameter_requires_grad(model, freeze_feature)
model.classifier = nn.Sequential(nn.Linear(512 * 7 * 7, 1024),
nn.ReLU(True), nn.Dropout(),
nn.Linear(1024, n_classes))
elif model_name == 'squeezenet1_1':
model = models.squeezenet1_1(pretrained=True)
set_parameter_requires_grad(model, freeze_feature)
model.num_classes = n_classes
model.classifier = nn.Sequential(
nn.Dropout(p=0.5), nn.Conv2d(512, n_classes, kernel_size=1),
nn.ReLU(inplace=True), nn.AvgPool2d(13, stride=1))
elif model_name == 'shufflenet_v2_x1_0':
model = models.shufflenet_v2_x1_0(pretrained=True)
set_parameter_requires_grad(model, freeze_feature)
n_feat_in = model.fc.in_features
model.fc = nn.Linear(n_feat_in, n_classes)
elif model_name == 'mobilenet_v2':
model = models.mobilenet_v2(pretrained=True)
set_parameter_requires_grad(model, freeze_feature)
n_feat_in = model.last_channel
model.classifier = nn.Sequential(nn.Dropout(0.2),
nn.Linear(n_feat_in, n_classes))
elif model_name == 'inception_v3':
# it expects (299,299) sized images and has auxiliary output
input_size = 299
model = models.inception_v3(pretrained=True)
set_parameter_requires_grad(model, freeze_feature)
# Handle the auxilary net
n_feat_in = model.AuxLogits.fc.in_features
model.AuxLogits.fc = nn.Linear(n_feat_in, n_classes)
# Handle the primary net
n_feat_in = model.fc.in_features
model.fc = nn.Linear(n_feat_in, n_classes)
else:
raise (ValueError, 'invalid model')
return model, input_size
def get_net(net_name, weight_path=None):
"""
根据网络名称获取模型
:param net_name: 网络名称
:param weight_path: 与训练权重路径
:return:
"""
pretrain = weight_path is None # 没有指定权重路径，则加载默认的预训练权重
if net_name in ['vgg', 'vgg16']:
net = models.vgg16(pretrained=pretrain)
elif net_name == 'vgg19':
net = models.vgg19(pretrained=pretrain)
elif net_name in ['resnet', 'resnet50']:
net = models.resnet50(pretrained=pretrain)
elif net_name == 'resnet101':
net = models.resnet101(pretrained=pretrain)
elif net_name in ['densenet', 'densenet121']:
net = models.densenet121(pretrained=pretrain)
elif net_name in ['inception']:
net = models.inception_v3(pretrained=pretrain)
elif net_name in ['mobilenet_v2']:
net = models.mobilenet_v2(pretrained=pretrain)
elif net_name in ['shufflenet_v2']:
net = models.shufflenet_v2_x1_0(pretrained=pretrain)
elif net_name in ['effnet']:
net = Effnet(
args.enet_type,
out_dim=args.out_dim,
pretrained=True,
freeze_cnn=args.freeze_cnn
)
else:
raise ValueError('invalid network name:{}'.format(net_name))
# 加载指定路径的权重参数
if weight_path is not None and net_name.startswith('densenet'):
pattern = re.compile(
r'^(.*denselayer\d+\.(?:norm|relu|conv))\.((?:[12])\.(?:weight|bias|running_mean|running_var))$')
state_dict = torch.load(weight_path)
for key in list(state_dict.keys()):
res = pattern.match(key)
if res:
new_key = res.group(1) + res.group(2)
state_dict[new_key] = state_dict[key]
del state_dict[key]
net.load_state_dict(state_dict)
elif weight_path is not None:
net.load_state_dict(torch.load(weight_path))
return net
def neural_network(self, model="resnet18", pre_trained=True):
# reference : https://pytorch.org/docs/stable/torchvision/models.html
model_list = [
"resnet18", "alexnet", "squeezenet", "vgg16", "densenet",
"inception", "googlenet", "shufflenet", "mobilenet", "resnet34",
"wide_resnet50_2", "mnasnet"
]
if model in model_list:
rospy.loginfo(
"[{}] You use model [{}]. Need some time to load model...".
format(self.node_name, model))
start_time = rospy.get_time()
if model == "resnet18":
self.model = models.resnet18(pretrained=pre_trained)
self.model.fc = torch.nn.Linear(512, 2)
elif model == "alexnet":
self.model = models.alexnet(pretrained=pre_trained)
elif model == "squeezenet":
self.model = models.squeezenet1_1(pretrained=pre_trained)
elif model == "vgg16":
self.model = models.vgg16(pretrained=pre_trained)
elif model == "densenet":
self.model = models.densenet161(pretrained=pre_trained)
elif model == "inception":
self.model = models.inception_v3(pretrained=pre_trained)
elif model == "googlenet":
self.model = models.googlenet(pretrained=pre_trained)
elif model == "shufflenet":
self.model = models.shufflenet_v2_x1_0(pretrained=pre_trained)
elif model == "mobilenet":
self.model = models.mobilenet_v2(pretrained=pre_trained)
elif model == "resnext50_32x4d":
self.model = models.resnext50_32x4d(pretrained=pre_trained)
elif model == "resnet34":
self.model = models.resnet34(pretrained=pre_trained)
elif model == "wide_resnet50_2":
self.model = models.wide_resnet50_2(pretrained=pre_trained)
elif model == "mnasnet":
self.model = models.mnasnet1_0(pretrained=pre_trained)
interval = rospy.get_time() - start_time
rospy.loginfo(
"[{}] loading modle done! Use {:.2f} seconds.".format(
self.node_name, interval))
else:
rospy.loginfo(
"[{}] We don't have the model in this rpoject. Please choose one of below: "
.format(self.node_name))
rospy.loginfo("[{}] {}".format(self.node_name, model_list))
self.on_shutdown()
def shufflenet(pretrained: bool = True, input_channels: int = 1):
model = models.shufflenet_v2_x1_0(pretrained=pretrained)
model.conv1 = nn.Sequential(
nn.Conv2d(input_channels,
24,
kernel_size=(3, 3),
stride=2,
padding=1,
bias=False), *model.conv1[1:])
model = nn.Sequential(*list(model.children())[:-1])
model.input_size = 224
model.out_features = 1024
model.name = 'shufflenet'
return model
def paraWC(model_class, batch_size, EPOCH, lr):
if model_class == 'resnet':
net = models.resnet18()
if model_class == 'alexnet':
net = models.alexnet()
if model_class == 'vgg':
net = models.vgg16()
if model_class == 'googlenet':
net = models.googlenet()
if model_class == 'densenet':
net = models.densenet161()
if model_class == 'shufflenet':
net = models.shufflenet_v2_x1_0()
if model_class == 'ours':
net = SSNet()
return net, batch_size, EPOCH, lr
