pytorch加载自己的数据集源码分享

下面是关于pytorch加载自己的数据集的完整攻略。

1. 准备数据集

在使用pytorch训练模型需要一个自己的数据集，这里以图像分类任务为例，准备一个包含训练集和测试集的数据集，其中每个图像都分好了类别并放在对应的文件夹中，例如：

dataset
├── train
│   ├── cat
│   │   ├── cat1.jpg
│   │   ├── cat2.jpg
│   │   └── ...
│   ├── dog
│   │   ├── dog1.jpg
│   │   ├── dog2.jpg
│   │   └── ...
│   └── ...
└── test
    ├── cat
    │   ├── cat1.jpg
    │   ├── cat2.jpg
    │   └── ...
    ├── dog
    │   ├── dog1.jpg
    │   ├── dog2.jpg
    │   └── ...
    └── ...

2. 定义Dataset类

接下来需要定义一个torch.utils.data.Dataset的子类，在其中实现数据集的加载、预处理等操作。以下是一个基本的示例：

import torch
from PIL import Image
from torch.utils.data import Dataset

class MyDataset(Dataset):
    def __init__(self, root_dir, transform=None):
        self.root_dir = root_dir
        self.transform = transform
        self.labels = sorted(os.listdir(root_dir))

    def __len__(self):
        return sum([len(files) for _, _, files in os.walk(self.root_dir)])

    def __getitem__(self, index):
        label = self.labels[index // len(self)]
        img_path = glob.glob(f"{self.root_dir}/{label}/*")[index % len(self)]
        img = Image.open(img_path).convert("RGB")
        if self.transform:
            img = self.transform(img)
        return img, label

在上面的代码中，我们定义了一个名为MyDataset的子类，实现了4个方法：

• __init__：初始化方法，接收两个参数，一个是数据集的根目录，另一个是数据集上的转换操作（transform），这里使用了PIL.Image库来加载图像。
• __len__：返回数据集的长度，这里是遍历数据集中所有图片的数量。
• __getitem__：根据索引返回对应的图像和标签，并且进行预处理，这里直接返回了图像的张量和标签字符串。
• labels：这个属性用于保存数据集中所有类别的名称，使用了Python内置的os.listdir方法。

3. 定义DataLoader

接下来可以使用torch.utils.data.DataLoader类来加载数据集，并使用pytorch进行训练。以下是一个基本的示例：

from torchvision import transforms
from torch.utils.data import DataLoader

transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
])

train_set = MyDataset("dataset/train", transform=transform)
train_loader = DataLoader(train_set, batch_size=64, shuffle=True)

test_set = MyDataset("dataset/test", transform=transform)
test_loader = DataLoader(test_set, batch_size=64, shuffle=False)

在上面的代码中，我们使用torchvision.transforms库来定义了一组转换操作，包括将图像大小转换为224x224、转换为张量和标准化操作，并将其传递给MyDataset类的实例中。

然后我们分别创建了训练集和测试集的DataLoader，其中MyDataset是传入数据集的实例，batch_size表示每个batch的大小，shuffle=True表示在每个epoch开始时打乱数据。

至此，我们已经完成了pytorch加载自己的数据集的完整攻略。

4. 示例说明

示例一

假设我们要训练一个resnet18网络来对上述示例中的图片分类，可以按照以下步骤定义并训练这个模型：

import torch.nn as nn
import torch.optim as optim
from torchvision.models import resnet18

device = "cuda" if torch.cuda.is_available() else "cpu"
model = resnet18(pretrained=True)
model.fc = nn.Linear(512, 2)
model = model.to(device)

criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.01)

for epoch in range(10):
    running_loss = 0
    for inputs, labels in train_loader:
        inputs, labels = inputs.to(device), labels.to(device)
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        running_loss += loss.item()
    print(f"Epoch: {epoch+1} Loss: {running_loss/len(train_loader):.4f}")

在上述代码中，我们使用了torchvision内置的resnet18网络，将其输出层改为2个节点的全连接层，用于分类2个类别的图像。

我们使用了交叉熵损失函数和随机梯度下降优化器，每个epoch训练完成后输出当前loss值。

示例二

接下来我们可以使用训练好的模型对测试集中的图像进行预测，示例代码如下：

correct = 0
total = 0
with torch.no_grad():
    for inputs, labels in test_loader:
        inputs, labels = inputs.to(device), labels.to(device)
        outputs = model(inputs)
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()
print(f"Accuracy: {100*correct/total:.2f}%")

在上述代码中，我们使用torch.no_grad()上下文管理器来关闭梯度计算，防止内存溢出。然后遍历测试集中的所有图像，进行前向预测，并计算准确率。