1

2

3

4

5

6

7

8

9

10

11

12

13

import torch

from torch.utils.data import Subset

from torchvision.datasets import MNIST

# 载入MNIST数据集

dataset = MNIST(root='data/', download=True, train=True)

# 定义子集中的索引

indices = [0, 1, 2]

# 创建子集

subset = Subset(dataset, indices)

# 打印子集中的元素

for i, (image, label) in enumerate(subset):

    print(f"Index: {i}, Label: {label}")

    # 这里可以加入图像展示代码，如：image.show()

1

2

3

4

5

6

7

8

9

10

11

12

13

from torch.utils.data import DataLoader, random_split

# 假设我们有一个较大的数据集

large_dataset = MNIST(root='data/', download=True, train=True)

# 随机划分数据集为训练集和验证集

train_size = int(0.8 * len(large_dataset))

val_size = len(large_dataset) - train_size

train_dataset, val_dataset = random_split(large_dataset, [train_size, val_size])

# 使用Subset类来进一步细化训练集或验证集

train_indices = range(100)  # 假设我们只用前100个样本来训练

train_subset = Subset(train_dataset, train_indices)

# 创建DataLoader

train_loader = DataLoader(train_subset, batch_size=10, shuffle=True)

# 现在可以使用train_loader来训练模型了