pytorch怎么加载数据

在 PyTorch 中，加载数据通常使用 torch.utils.data.Dataset 和 torch.utils.data.DataLoader。下面从基础到常见用法给你讲清楚。

一、核心概念

1. Dataset

• 定义“数据从哪里来”
• 需要继承 Dataset 类
• 实现三个方法：
  • __init__：初始化
  • __len__：数据总量
  • __getitem__：按索引取一条数据

2. DataLoader

• 负责“怎么批量、多线程、打乱加载”
• 输入 Dataset，输出可迭代的 batch

二、最基础示例（自定义 Dataset）

import torch
from torch.utils.data import Dataset, DataLoader

class MyDataset(Dataset):
    def __init__(self, data, labels):
        self.data = data
        self.labels = labels

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        return self.data[idx], self.labels[idx]

# 示例数据
x = torch.randn(100, 10)
y = torch.randint(0, 2, (100,))

dataset = MyDataset(x, y)
dataloader = DataLoader(dataset, batch_size=8, shuffle=True)

for batch_x, batch_y in dataloader:
    print(batch_x.shape, batch_y.shape)
    break

三、加载图片（常用）

1. 使用 ImageFolder（最简单）

目录结构：

data/
├── cat/
│   ├── 1.jpg
│   └── 2.jpg
├── dog/
│   ├── 1.jpg
│   └── 2.jpg

from torchvision import datasets, transforms

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

dataset = datasets.ImageFolder("data", transform=transform)
loader = DataLoader(dataset, batch_size=32, shuffle=True)

2. 自定义图片 Dataset（更灵活）

from PIL import Image
import os

class ImageDataset(Dataset):
    def __init__(self, img_dir, transform=None):
        self.img_dir = img_dir
        self.img_list = os.listdir(img_dir)
        self.transform = transform

    def __len__(self):
        return len(self.img_list)

    def __getitem__(self, idx):
        img_path = os.path.join(self.img_dir, self.img_list[idx])
        image = Image.open(img_path).convert("RGB")
        if self.transform:
            image = self.transform(image)
        return image

四、加载常见公开数据集

from torchvision import datasets

train_set = datasets.CIFAR10(
    root="./data",
    train=True,
    download=True,
    transform=transforms.ToTensor()
)

train_loader = DataLoader(train_set, batch_size=64, shuffle=True)

五、常见参数说明（DataLoader）

DataLoader(
    dataset,
    batch_size=32,        # 每批大小
    shuffle=True,        # 是否打乱
    num_workers=4,       # 多线程加载
    drop_last=False      # 是否丢弃最后不满的batch
)

六、典型训练循环

for epoch in range(epochs):
    for x, y in dataloader:
        x, y = x.to(device), y.to(device)
        pred = model(x)
        loss = criterion(pred, y)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

如果你愿意，可以告诉我：

• ✅ 是 图片 / 文本 / 表格数据？
• ✅ 是否使用 CUDA / 分布式训练？
• ✅ 是否用 HuggingFace / torchtext？

我可以直接给你对应的最优加载方案。