欢迎访问宙启技术站
智能推送

使用Python创建的torchvision.modelsmobilenet_v2模型的迁移学习方法

发布时间:2023-12-12 08:28:45

迁移学习是指通过在一个任务上训练好的模型,将其应用于另一个相关任务上。这种方法的优势在于可以利用预训练模型在大规模数据集上学到的特征,加快新任务的训练速度,并提升模型的性能。

torchvision.models包含了一些经过预训练的深度学习模型,其中包括MobileNetV2。MobileNetV2是一种轻量级的卷积神经网络模型,适用于移动设备等资源受限的场景。

下面是一种使用torchvision.models.mobilenet_v2进行迁移学习的方法,包括模型的加载、微调和预测过程。假设我们要训练模型进行猫和狗的分类任务。

首先,我们需要导入所需的库。

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, models, transforms

接下来,我们定义一些超参数和数据变换。

# 定义超参数
batch_size = 32
num_classes = 2
num_epochs = 10
feature_extract = True

# 定义图像变换
data_transforms = {
    'train': transforms.Compose([
        transforms.RandomResizedCrop(224),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ]),
    'val': transforms.Compose([
        transforms.Resize(256),
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ]),
}

然后,我们加载数据集。

# 加载数据集
data_dir = 'path_to_dataset'
image_datasets = {x: datasets.ImageFolder(os.path.join(data_dir, x), data_transforms[x])
                  for x in ['train', 'val']}
dataloaders = {x: torch.utils.data.DataLoader(image_datasets[x], batch_size=batch_size, shuffle=True)
               for x in ['train', 'val']}

接下来,我们定义一个函数用于微调模型。

def set_parameter_requires_grad(model, feature_extracting):
    if feature_extracting:
        for param in model.parameters():
            param.requires_grad = False

然后我们加载预训练的MobileNetV2模型,并替换最后一层全连接层。

def initialize_model(num_classes, feature_extract, use_pretrained=True):
    model = models.mobilenet_v2(pretrained=use_pretrained)
    set_parameter_requires_grad(model, feature_extract)
    num_features = model.classifier[1].in_features
    model.classifier = nn.Linear(num_features, num_classes)
    return model

定义损失函数和优化器。

model = initialize_model(num_classes, feature_extract)
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)

最后,我们训练模型。

# 训练模型
def train_model(model, criterion, optimizer, num_epochs=25):
    for epoch in range(num_epochs):
        for phase in ['train', 'val']:
            if phase == 'train':
                model.train()
            else:
                model.eval()
            
            running_loss = 0.0
            running_corrects = 0
            
            for inputs, labels in dataloaders[phase]:
                optimizer.zero_grad()
                
                with torch.set_grad_enabled(phase == 'train'):
                    outputs = model(inputs)
                    _, preds = torch.max(outputs, 1)
                    loss = criterion(outputs, labels)
                    
                    if phase == 'train':
                        loss.backward()
                        optimizer.step()
                
                running_loss += loss.item() * inputs.size(0)
                running_corrects += torch.sum(preds == labels.data)
                
        epoch_loss = running_loss / len(image_datasets[phase])
        epoch_acc = running_corrects.double() / len(image_datasets[phase])
        
        print('{} Loss: {:.4f} Acc: {:.4f}'.format(phase, epoch_loss, epoch_acc))

调用训练函数开始训练。

train_model(model, criterion, optimizer, num_epochs=num_epochs)

以上就是使用torchvision.models.mobilenet_v2进行迁移学习的方法。在实际应用中,你可以根据自己的需求修改超参数、数据变换和训练方式,以获得更好的模型性能。