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