Python中MNIST数据集不同加载方式的互转实现咨询

一、不同MNIST加载格式互转完全可行 #

不管是torchvision的Dataset对象、原始numpy数组格式，还是CSV存储的格式，核心是理清每种格式的数据结构，就能实现互相转换。

1. torchvision MNIST Dataset → 原始MNIST格式（numpy数组） #

torchvision的MNIST类实例（比如trainset）是自定义Dataset对象，内部直接存储了图像和标签的张量。要转成原始MNIST的numpy格式（图像为(n_samples,28,28)的uint8数组，标签为(n_samples,)的int数组），直接提取内部属性即可：

import torch
import numpy as np
from torchvision.datasets import MNIST

# 加载torchvision的MNIST
transform = None  # 根据需求添加transform，比如ToTensor()
trainset = MNIST('./', download=True, train=True, transform=transform)

# 转成numpy格式
# 未加ToTensor时，trainset.data本身就是uint8的28x28张量
mnist_images = trainset.data.numpy()  # shape (60000,28,28)
mnist_labels = trainset.targets.numpy()  # shape (60000,)

# 如果加了ToTensor（图像变为float32，值在0-1之间），需要转回原始格式
# mnist_images = (trainset.data * 255).numpy().astype(np.uint8)

2. 原始MNIST格式 → torchvision风格的Dataset对象 #

如果手里有原始的图像numpy数组和标签数组，想要转成和trainset功能一致的Dataset对象（支持索引、transform等），可以自定义一个Dataset类：

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

class CustomMNISTDataset(Dataset):
    def __init__(self, images, labels, transform=None):
        self.images = images  # shape (n_samples,28,28)
        self.labels = labels  # shape (n_samples,)
        self.transform = transform

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

    def __getitem__(self, idx):
        img = self.images[idx]
        label = self.labels[idx]

        if self.transform:
            # 多数transform需要PIL图像格式，先转换
            img = Image.fromarray(img.astype('uint8'))
            img = self.transform(img)
        else:
            # 直接转成张量并添加通道维度（符合PyTorch的图像格式要求）
            img = torch.tensor(img, dtype=torch.float32).unsqueeze(0)

        return img, torch.tensor(label, dtype=torch.long)

# 使用示例：假设已有mnist_images和mnist_labels数组
custom_trainset = CustomMNISTDataset(mnist_images, mnist_labels, transform=transform)
# 现在custom_trainset和torchvision的trainset用法完全一致

二、CSV导出与还原的正确姿势 #

你之前失败的核心原因是：CSV保存的是扁平化的图像数据，而torchvision的Dataset需要带结构的图像（28x28）+标签，且必须封装成Dataset对象，不是读CSV就能自动变成trainset格式。

1. torchvision Dataset → CSV #

把图像展平成784维，和标签一起存为CSV：

import pandas as pd

# 提取数据并展平
images_flat = trainset.data.numpy().reshape(-1, 28*28)  # (60000,784)
labels = trainset.targets.numpy().reshape(-1,1)  # (60000,1)

# 合并成DataFrame并导出
df = pd.DataFrame(np.hstack([labels, images_flat]), columns=['label'] + [f'pixel_{i}' for i in range(784)])
df.to_csv('mnist_train.csv', index=False)

2. CSV → torchvision风格Dataset #

读取CSV后，需要把图像列重新reshape成28x28，再用自定义Dataset封装：

# 读取CSV
df = pd.read_csv('mnist_train.csv')
labels = df['label'].values
images_flat = df.drop('label', axis=1).values

# 把扁平化图像恢复成28x28结构
images = images_flat.reshape(-1,28,28).astype(np.uint8)

# 用自定义Dataset封装
restored_trainset = CustomMNISTDataset(images, labels, transform=transform)
# 现在restored_trainset和原始trainset功能完全一致

三、关于Python类、属性的直白解释 #

• 类：相当于一张"图纸"，定义了这类事物该有什么数据（比如数据集要有图像、标签）和能做什么操作（比如获取样本总数、按索引取样本）。torchvision.datasets.MNIST就是一张数据集图纸，规定了怎么下载、读取MNIST数据。
• 对象：是用类这个"图纸"造出来的具体实物。比如trainset = MNIST(...)，trainset就是一个实实在在的数据集对象，它有自己的data（存储所有图像）、targets（存储所有标签）这些专属数据。
• 属性赋值：就是给对象的专属数据赋值。可以在类的初始化方法（__init__）里自动完成（比如MNIST类在初始化时会把下载好的图像存到self.data里，self指的就是当前这个对象），也可以在外面手动给对象加新属性，比如trainset.my_tag = "自定义训练集"。

内容的提问来源于stack exchange，提问作者user27357269