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

使用Python实现的torchvision.modelsmobilenet_v2模型的图像分类精度评估

发布时间:2023-12-12 08:26:47

MobileNetV2是一种轻量级的深度卷积神经网络模型,经过特殊设计可以在计算资源有限的设备上进行高效的图像分类。在PyTorch的torchvision.models模块中,我们可以使用预先训练好的MobileNetV2模型进行图像分类任务。

以下是一个使用Python实现的torchvision.models.mobilenet_v2模型的图像分类精度评估的示例:

1. 导入所需的库和模块:

import torch
import torch.nn as nn
import torchvision.transforms as transforms
import torchvision.models as models
import torchvision.datasets as datasets

2. 定义数据预处理的转换:

preprocess = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406],
                         std=[0.229, 0.224, 0.225])
])

此处预处理的转换将图像调整为224x224大小,然后进行张量转换和标准化处理。

3. 加载预训练的MobileNetV2模型:

model = models.mobilenet_v2(pretrained=True)

这将下载预训练的MobileNetV2模型并加载到内存中。

4. 加载用于评估的数据集:

data_dir = 'path_to_dataset_directory'
eval_dataset = datasets.ImageFolder(
    root=data_dir,
    transform=preprocess
)
eval_loader = torch.utils.data.DataLoader(
    eval_dataset,
    batch_size=4,
    shuffle=False,
    num_workers=1
)

此处假设我们的数据集放置在名为path_to_dataset_directory的文件夹中。我们使用ImageFolder数据集类加载数据集,并使用DataLoader进行数据集的批处理和并行加载。

5. 定义评估函数:

def evaluate(model, eval_loader):
    model.eval()
    correct = 0
    total = 0

    with torch.no_grad():
        for images, labels in eval_loader:
            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 = 100 * correct / total
    return accuracy

评估函数将模型设为评估模式(model.eval()),并在测试数据集上进行迭代,计算模型预测的准确性。

6. 使用evaluate函数评估模型精度:

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = model.to(device)

accuracy = evaluate(model, eval_loader)
print("Accuracy: {:.2f}%".format(accuracy))

在此处,我们将模型移动到可用的计算设备(GPU或CPU),然后调用evaluate函数获取模型在测试数据集上的准确性。

完整的代码如下所示:

import torch
import torch.nn as nn
import torchvision.transforms as transforms
import torchvision.models as models
import torchvision.datasets as datasets

# Define data preprocessing transforms
preprocess = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406],
                         std=[0.229, 0.224, 0.225])
])

# Load pre-trained MobileNetV2 model
model = models.mobilenet_v2(pretrained=True)

# Load evaluation dataset
data_dir = 'path_to_dataset_directory'
eval_dataset = datasets.ImageFolder(
    root=data_dir,
    transform=preprocess
)
eval_loader = torch.utils.data.DataLoader(
    eval_dataset,
    batch_size=4,
    shuffle=False,
    num_workers=1
)

# Define evaluation function
def evaluate(model, eval_loader):
    model.eval()
    correct = 0
    total = 0

    with torch.no_grad():
        for images, labels in eval_loader:
            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 = 100 * correct / total
    return accuracy

# Evaluate model accuracy
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = model.to(device)

accuracy = evaluate(model, eval_loader)
print("Accuracy: {:.2f}%".format(accuracy))

请替换path_to_dataset_directory为实际存储数据集的目录,并确保已经安装了PyTorch和torchvision库。运行此代码将计算预训练的MobileNetV2模型在给定测试数据集上的分类精度。精度结果将以百分比形式打印出来。