在PyTorch中定义批量大小为1的手动排序的MNIST数据集

import torch
from torch.utils.data import Dataset, DataLoader
from torchvision import datasets, transforms

# 定义自定义数据集类
class SortedMNIST(Dataset):
    def __init__(self, root, train=True, transform=None, download=False):
        self.mnist = datasets.MNIST(root=root, train=train, transform=transform, download=download)
        self.sorted_indices = self.sort_indices()

    def sort_indices(self):
        # 获取所有标签
        labels = self.mnist.targets
        # 获取排序后的索引
        sorted_indices = torch.argsort(labels)
        return sorted_indices

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

    def __getitem__(self, idx):
        sorted_idx = self.sorted_indices[idx]
        image, label = self.mnist[sorted_idx]
        return image, label

# 定义数据转换
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.1307,), (0.3081,))
])

# 创建自定义数据集
sorted_mnist_train = SortedMNIST(root='./data', train=True, transform=transform, download=True)
sorted_mnist_test = SortedMNIST(root='./data', train=False, transform=transform, download=True)

# 创建数据加载器，批量大小为1
train_loader = DataLoader(sorted_mnist_train, batch_size=1, shuffle=False)
test_loader = DataLoader(sorted_mnist_test, batch_size=1, shuffle=False)

# 测试数据加载器
for batch_idx, (data, target) in enumerate(train_loader):
    print(f'Batch {batch_idx}: Label {target.item()}')
    if batch_idx >= 10:  # 只打印前10个批次
        break