极客笔记Pytorch 微调预训练模型MobileNet_V2
在本文中，我们将介绍如何使用Pytorch进行微调预训练模型MobileNet_V2。MobileNet_V2是一种轻量级神经网络模型，适用于计算资源有限的情况下进行深度学习任务。通过微调预训练的MobileNet_V2模型，我们可以在特定的任务上获得更好的性能。
阅读更多：Pytorch 教程
1. 导入所需库和数据集
在开始之前，我们首先需要导入所需的库和数据集。Pytorch库是一个基于Torch的机器学习库，提供了丰富的工具和函数用于实现深度学习模型。此外，我们还需要下载和加载MobileNet_V2的预训练权重，以及一个相关的数据集用于微调。
以下是导入库和数据集的代码示例：
import torch
import torchvision
from torchvision import datasets, models, transforms
# 检查CUDA是否可用
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# 数据预处理和增强
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=4, shuffle=True, num_workers=4)
for x in ['train', 'val']}
class_names = image_datasets['train'].classes
2. 加载预训练模型
接下来，我们需要加载预训练的MobileNet_V2模型，并且冻结它的所有层，以便在微调过程中不更新权重。这是因为MobileNet_V2的参数已经通过大规模的图像数据集进行了训练，具有良好的特征提取能力。
以下是加载预训练模型的代码示例：
# 加载预训练模型
model = models.mobilenet_v2(pretrained=True)
model.to(device)
# 冻结预训练模型的所有层
for param in model.parameters():
param.requires_grad = False
# 更改最后一层全连接层的输出大小
num_ftrs = model.classifier[1].in_features
model.classifier[1] = torch.nn.Linear(num_ftrs, len(class_names))
# 将模型参数移到CUDA设备上（如果可用）
model = model.to(device)
在上述代码中，我们首先加载预训练的MobileNet_V2模型，然后将其迁移到所选的设备上（CPU或GPU）。然后，我们冻结了所有层的参数，以防止在微调过程中对其进行更新。最后，我们更改了该模型的最后一层全连接层的输出大小，以适应我们数据集中的类别数量。同样，我们还将模型参数移动到CUDA设备上（如果可用）。
3. 定义损失函数和优化器
为了进行微调，我们需要定义损失函数和优化器。在分类任务中，常用的损失函数是交叉熵损失函数。优化器的选择可以根据具体情况进行，常见的选择是随机梯度下降（SGD）或Adam优化器。
以下是定义损失函数和优化器的代码示例：
# 定义损失函数和优化器
criterion = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.001, momentum=0.9)
在上述代码中，我们使用交叉熵损失函数作为分类任务的损失函数，并使用随机梯度下降（SGD）作为优化器。我们还设置了学习率为0.001和动量为0.9。
4. 微调模型
现在我们可以开始进行模型的微调了。我们需要迭代数据集中的所有样本，并根据模型的预测结果和真实标签来计算损失，并通过反向传播来更新模型的参数。
以下是微调模型的代码示例：
# 微调模型
num_epochs = 10
best_acc = 0.0
for epoch in range(num_epochs):
# 训练阶段
model.train()
running_loss = 0.0
running_corrects = 0
for inputs, labels in dataloaders['train']:
inputs = inputs.to(device)
labels = labels.to(device)
optimizer.zero_grad()
# 前向传播
with torch.set_grad_enabled(True):
outputs = model(inputs)
_, preds = torch.max(outputs, 1)
loss = criterion(outputs, labels)
# 反向传播和优化
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['train'])
epoch_acc = running_corrects.double() / len(image_datasets['train'])
# 验证阶段
model.eval()
running_loss = 0.0
running_corrects = 0
for inputs, labels in dataloaders['val']:
inputs = inputs.to(device)
labels = labels.to(device)
# 前向传播
with torch.no_grad():
outputs = model(inputs)
_, preds = torch.max(outputs, 1)
loss = criterion(outputs, labels)
running_loss += loss.item() * inputs.size(0)
running_corrects += torch.sum(preds == labels.data)
epoch_loss_val = running_loss / len(image_datasets['val'])
epoch_acc_val = running_corrects.double() / len(image_datasets['val'])
# 输出训练过程中的损失和准确率
print('Epoch [{}/{}], Loss: {:.4f}, Acc: {:.4f}, Val Loss: {:.4f}, Val Acc: {:.4f}'
.format(epoch+1, num_epochs, epoch_loss, epoch_acc, epoch_loss_val, epoch_acc_val))
# 保存在验证集上获得最佳准确率的模型
if epoch_acc_val > best_acc:
best_acc = epoch_acc_val
best_model_wts = model.state_dict()
# 加载在验证集上获得最佳准确率的模型
model.load_state_dict(best_model_wts)
在上述代码中，我们首先设置了训练的总轮数（num_epochs）和最佳准确率（best_acc）。然后，我们使用两个嵌套的循环来进行训练和验证。在训练阶段，我们首先将模型设置为训练模式，然后对每个输入进行前向传播和反向传播以更新模型的参数。在验证阶段，我们将模型设置为评估模式，并计算验证集上的损失和准确率。我们还输出每个epoch的训练和验证过程中的损失和准确率，并保存在验证集上获得最佳准确率的模型参数。最后，我们加载在验证集上获得最佳准确率的模型参数，以便后续的预测和测试。
5. 模型预测和测试
在微调完模型后，我们可以使用该模型进行预测和测试了。首先，我们需要定义一个函数来使用训练好的模型对单个图像进行预测。
以下是定义预测函数的代码示例：
def predict_image(image_path, model):
image = Image.open(image_path)
image_tensor = data_transforms['val'](image).unsqueeze(0)
image_tensor = image_tensor.to(device)
with torch.no_grad():
model.eval()
outputs = model(image_tensor)
_, preds = torch.max(outputs, 1)
predicted_class = class_names[preds.item()]
return predicted_class
在上述代码中，我们首先打开图像文件，然后将其转换为对应的张量并移动到所选的设备上。接下来，我们将模型设置为评估模式，并使用模型对图像进行预测。最后，我们通过输出的类别索引找到对应的类别名称，并返回预测结果。
除了单个图像的预测，我们还可以对整个测试集进行预测，并计算预测准确率。
以下是对测试集进行预测并计算准确率的代码示例：
correct = 0
total = 0
model.eval()
with torch.no_grad():
for images, labels in dataloaders['val']:
images = images.to(device)
labels = labels.to(device)
outputs = model(images)
_, predicted = torch.max(outputs.data, 1)
total += labels.size(0)
correct += (predicted == labels).sum().item()
accuracy = correct / total
print('Accuracy on test set: {:.4f}'.format(accuracy))
在上述代码中，我们首先将模型设置为评估模式，并使用torch.no_grad()来关闭梯度计算。然后，我们对测试集的每个图像进行预测，并统计正确预测的数量。最后，我们计算预测准确率并输出在测试集上的准确率。
总结
本文介绍了如何使用Pytorch进行微调预训练模型MobileNet_V2。我们首先导入所需的库和数据集，然后加载预训练的MobileNet_V2模型，并冻结其所有层。接下来，我们定义了损失函数和优化器，并进行了模型的微调。最后，我们展示了如何使用训练好的模型进行预测和测试，并计算了在测试集上的准确率。通过本文的介绍，希望读者能够了解和掌握使用Pytorch进行预训练模型微调的方法。