Python torchvision.models.mobilenet_v2() Examples

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)
    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 test():
    """ Compare the classification result of WarmMobileNetV2 versus torchvision mobilenet_v2. """
    new = WarmMobileNetV2()
    from torchvision.models import mobilenet_v2
    old = mobilenet_v2()
    state = old.state_dict()
    for k in list(state.keys()): # Map parameters of old, e.g. layer2.0.conv1.weight
        s = k.split('.') # to parameters of new, e.g. layer2-0-conv1.weight
        s = '-'.join(s[:-1])+'.'+s[-1]
        state[s] = state.pop(k)
    new.load_state_dict(state)
    warm.util.summary(old)
    warm.util.summary(new)
    x = torch.randn(1, 3, 224, 224)
    with torch.no_grad():
        old.eval()
        y_old = old(x)
        new.eval()
        y_new = new(x)
        if torch.equal(y_old, y_new):
            print('Success! Same results from old and new.')
        else:
            print('Warning! New and old produce different results.') 

def get_trained_model(args, device, train_loader, val_loader, criterion):
    if args.model == 'LeNet':
        model = LeNet().to(device)
        optimizer = torch.optim.Adadelta(model.parameters(), lr=1)
        scheduler = StepLR(optimizer, step_size=1, gamma=0.7)
        for epoch in range(args.pretrain_epochs):
            train(args, model, device, train_loader,
                  criterion, optimizer, epoch)
            scheduler.step()
    elif args.model == 'vgg16':
        model = VGG(depth=16).to(device)
        optimizer = torch.optim.SGD(model.parameters(), lr=0.01,
                                    momentum=0.9,
                                    weight_decay=5e-4)
        scheduler = MultiStepLR(
            optimizer, milestones=[int(args.pretrain_epochs*0.5), int(args.pretrain_epochs*0.75)], gamma=0.1)
        for epoch in range(args.pretrain_epochs):
            train(args, model, device, train_loader,
                  criterion, optimizer, epoch)
            scheduler.step()
    elif args.model == 'resnet18':
        model = models.resnet18(pretrained=False, num_classes=10).to(device)
        optimizer = torch.optim.SGD(model.parameters(), lr=0.01,
                                    momentum=0.9,
                                    weight_decay=5e-4)
        scheduler = MultiStepLR(
            optimizer, milestones=[int(args.pretrain_epochs*0.5), int(args.pretrain_epochs*0.75)], gamma=0.1)
        for epoch in range(args.pretrain_epochs):
            train(args, model, device, train_loader,
                  criterion, optimizer, epoch)
            scheduler.step()
    elif args.model == 'mobilenet_v2':
        model = models.mobilenet_v2(pretrained=True).to(device)

    if args.save_model:
        torch.save(model.state_dict(), os.path.join(
            args.experiment_data_dir, 'model_trained.pth'))
        print('Model trained saved to %s', args.experiment_data_dir)

    return model, optimizer 

def test_mobilenet_v2_model(input_var):
    original_model = torchvision_models.mobilenet_v2(pretrained=True)
    finetune_model = make_model(
        'mobilenet_v2',
        num_classes=1000,
        pool=default,
        pretrained=True,
    )
    copy_module_weights(original_model.classifier[-1], finetune_model._classifier)
    assert_equal_model_outputs(input_var, original_model, finetune_model) 

def test_mobilenet_v2_model_with_another_input_size(input_var):
    model = make_model('mobilenet_v2', num_classes=1000, pretrained=True)
    model(input_var) 

def __init__(self, classes, model='resnet18'):
        super(CNN, self).__init__()
        if (model == 'resnet18'):
            self.cnn = models.resnet18(pretrained=True)
            self.cnn.fc = nn.Linear(512, classes)
        elif (model == 'resnext50_32x4d'):

            self.cnn = models.resnext50_32x4d(pretrained=True)
            self.cnn.classifier = nn.Linear(1280, classes)
        elif (model == 'mobilenet_v2'):

            self.cnn = models.mobilenet_v2(pretrained=True)
            self.cnn.classifier = nn.Linear(1280, classes)