博客

Pytorch提取模型特征向量保存至csv的例子

  • Post category:Python

以下是详细的PyTorch提取模型特征向量并保存至CSV文件的完整攻略，包含两个示例。

安装PyTorch

在开始之前，我们需要先安装PyTorch。可以使用以下命令在Python中安装PyTorch：

pip install torch torchvision

加载模型

在进行特征提取之前，我们需要先加载模型。以下是一个使用PyTorch加载模型的示例：

import torch
import torchvision.models as models

# 加载模型
model = models.resnet18(pretrained=True)
model.eval()

在上面的代码中，我们使用PyTorch的models模块加载了一个预训练的ResNet-18模型，并将其设置为评估式。

加载数据

在进行特征提取之前，我们还需要加载数据。以下是一个使用PyTorch加载数据的示例：

import torchvision.datasets as datasets
import torchvision.transforms as transforms

# 定义数据转换
transform = 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])
])

# 加载数据集
dataset = datasets.ImageFolder('data', transform=transform)

在上面的代码中，我们使用PyTorch的datasets模块加载了一个图像数据集，并使用transforms模块定义了数据转换。

提取特征向量

在加载模型和数据之后，我们可以使用PyTorch提取模型特征向量。以下是一个使用PyTorch提取模型特征向量的示例：

import csv

# 打开CSV文件
with open('features.csv', 'w', newline='') as csvfile:
    writer = csv.writer(csvfile)

    # 遍历数据集
    for inputs, labels in dataset:
        # 提取特征向量
        outputs = model(inputs.unsqueeze(0))
        features = outputs.detach().numpy().flatten()

        # 写入CSV文件
        writer.writerow([labels] + list(features))

在上面的代码中，我们首先打开一个CSV文件，并创建一个CSV写入器。然后，我们遍历数据集，使用model函数提取特征向量，并将其保存至CSV文件中。

示例1：提取MNIST数据集的特征向量

以下是一个使用PyTorch提取MNIST数据集的特征向量并保存至CSV文件的示例：

import torch
import torchvision.datasets as datasets
import torchvision.transforms as transforms
import csv

# 加载数据集
dataset = datasets.MNIST(root='data', train=True, download=True, transform=transforms.ToTensor())

# 加载模型
model = torch.hub.load('pytorch/vision', 'resnet18', pretrained=True)
model.eval()

# 打开CSV文件
with open('mnist_features.csv', 'w', newline='') as csvfile:
    writer = csv.writer(csvfile)

    # 遍历数据集
    for inputs, labels in dataset:
        # 提取特征向量
        outputs = model(inputs.unsqueeze(0))
        features = outputs.detach().numpy().flatten()

        # 写入CSV文件
        writer.writerow([labels] + list(features))

在上面的代码中，我们首先使用PyTorch的datasets模块加载了MNIST数据集，并使用transforms模块定义了数据转换。接着，我们使用PyTorch的torch.hub模块加载了一个预训练的ResNet-18模型，并将其设置为评估式。最后，我们遍历数据集，使用model函数提取特征向量，并将其保存至CSV文件中。

示例2：提取CIFAR-10数据集的特征向量

以下是一个使用PyTorch提取CIFAR-10数据集的特征向量并保存至CSV文件的示例：

import torch
import torchvision.datasets as datasets
import torchvision.transforms as transforms
import csv

# 加载数据集
dataset = datasets.CIFAR10(root='data', train=True, download=True, transform=transforms.ToTensor())

# 加载模型
model = torch.hub.load('pytorch/vision', 'resnet18', pretrained=True)
model.eval()

# 打开CSV文件
with open('cifar10_features.csv', 'w', newline='') as csvfile:
    writer = csv.writer(csvfile)

    # 遍历数据集
    for inputs, labels in dataset:
        # 提取特征向量
        outputs = model(inputs.unsqueeze(0))
        features = outputs.detach().numpy().flatten()

        # 写入CSV文件
        writer.writerow([labels] + list(features))

在上面的代码中，我们首先使用PyTorch的datasets模块加载了CIFAR-10数据集，并使用transforms模块定义了数据转换。接着，我们使用PyTorch的torch.hub模块加载了一个预训练的ResNet-18模型，并将其设置为评估式。最后，我们遍历数据集，使用model函数提取特征向量，并将其保存至CSV文件中。

总结

本文详细讲解了如何使用PyTorch提取模型特征向量并保存至CSV文件的完整攻略。通过本文的学习，您可以了解如何使用PyTorch加载模型和数据，以及如何使用PyTorch提取模型特征向量并保存至CSV文件。同时，本文提供了两个示例，分别是使用PyTorch提取MNIST数据集的特征向量和使用PyTorch提取CIFAR-10数据集的特征向量。

Tags: