由于 PyTorch 与 MindSpore 在数据加载、数据增强、数据集格式、分布式读取等机制上存在差异，若直接复用原生代码会出现数据不匹配、精度丢失、训练卡死等问题，因此必须按照 MindSpore 数据 pipeline 规范完成标准化改造。聚焦小模型迁移中的数据准备全流程，从差异分析、数据集构建、增强适配、加载优化四个维度展开，并提供可直接运行的迁移对照代码，覆盖图像分类、文本分类等主流小模型场景。

一、PyTorch 与 MindSpore 数据机制核心差异

小模型迁移的数据准备，本质是将PyTorch 的 DataLoader、Dataset、Transforms平滑迁移到MindSpore 的 Dataset、GeneratorDataset、vision / 文本算子体系。二者核心差异集中在三点：

1. 数据管道构建方式：PyTorch 使用类继承 Dataset + 重写__getitem__，MindSpore 采用算子链式调用设计；
2. 数据增强接口：视觉 / 文本增强函数名称、参数顺序、输入格式不同，无法直接兼容；
3. 批处理与格式：PyTorch 默认通道在前（CHW），MindSpore 自动适配但需统一归一化、格式转换；
4. 读取效率：MindSpore 支持多进程、异步加载、MindRecord 优化格式，适合昇腾 NPU 加速。

掌握差异后，数据准备迁移可遵循“数据集定义→数据清洗→格式转换→增强适配→管道构建”五步流程，确保迁移后数据分布与原 PyTorch 项目完全一致，保证模型精度无损。

二、小模型迁移数据准备核心流程

数据准备是小模型迁移的基础工程，目标是让迁移后模型输入与原模型保持相同尺寸、相同分布、相同增强逻辑。完整流程包括：

1. 数据集路径与结构统一：保持目录结构与 PyTorch 一致，降低迁移成本；
2. 自定义数据集适配：将 PyTorch 的 Dataset 类改写为 MindSpore 可识别格式；
3. 数据增强对齐：替换 Resize、Normalize、Flip、Padding 等算子；
4. 批处理、 shuffle、多线程加载配置：对齐训练超参；
5. 输出校验：验证 shape、dtype、数值范围与原项目一致。

该流程无需修改模型结构，仅通过数据层适配即可完成迁移前的核心准备工作。

三、环境准备与依赖安装

在昇腾环境中配置 PyTorch+MindSpore 双框架，用于数据对齐验证：

# 安装依赖
pip install torch torchvision
pip install mindspore mindvision  # 昇思主框架+视觉工具
pip install pillow numpy pandas

环境支持图像分类、文本分类等小模型迁移，适配 openEuler + 鲲鹏 / 昇腾硬件。

四、完整迁移代码：从 PyTorch 到 MindSpore（数据准备）

下面以图像分类小模型为例，提供PyTorch 原版数据代码与昇思 MindSpore 迁移代码，实现 1:1 对齐，可直接用于迁移任务。

1. PyTorch 原版数据准备代码（基准）

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

# 1. 数据增强（与训练对齐）
transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.RandomHorizontalFlip(p=0.5),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

# 2. 加载数据集
train_dataset = datasets.ImageFolder(root="./data/train", transform=transform)

# 3. 构建DataLoader
train_loader = DataLoader(
    train_dataset,
    batch_size=32,
    shuffle=True,
    num_workers=4,
    drop_last=True
)

# 测试输出
for images, labels in train_loader:
    print("PyTorch 数据shape:", images.shape)  # (32,3,224,224)
    print("PyTorch 标签shape:", labels.shape)
    break

2. MindSpore 迁移版数据准备代码（核心）

import mindspore as ms
import mindspore.dataset as ds
import mindspore.dataset.vision as vision
from mindspore.dataset import GeneratorDataset

# 1. 自定义数据集（对齐PyTorch格式）
class CustomDataset:
    def __init__(self, data_dir):
        # 模拟加载图像路径+标签（可替换为真实数据集）
        self.data_list = self.load_data(data_dir)

    def load_data(self, data_dir):
        # 实际项目中读取ImageFolder结构
        data = [("./test.jpg", 0) for _ in range(1000)]
        return data

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

    def __getitem__(self, index):
        img_path, label = self.data_list[index]
        # 读取图像（必须返回numpy格式）
        img = ms.dataset.vision.read_file(img_path)
        img = ms.dataset.vision.decode_image(img)
        return img, label

# 2. 数据增强（与PyTorch完全对齐）
def create_dataset(data_dir, batch_size=32):
    # 加载自定义数据
    dataset = GeneratorDataset(
        source=CustomDataset(data_dir),
        column_names=["image", "label"]
    )

    # 对齐增强算子
    trans = [
        vision.Resize((224, 224)),
        vision.RandomHorizontalFlip(prob=0.5),
        vision.ToTensor(),
        vision.Normalize(mean=[0.485, 0.456, 0.406],
                         std=[0.229, 0.224, 0.225])
    ]

    # 数据管道
    dataset = dataset.map(operations=trans, input_columns="image")
    dataset = dataset.batch(batch_size, drop_remainder=True)
    dataset = dataset.shuffle(buffer_size=100)
    return dataset

# 3. 构建数据集
if __name__ == "__main__":
    ms.set_context(device_target="Ascend")  # 运行在昇腾NPU
    train_dataset = create_dataset(data_dir="./data/train")

    # 验证输出shape（与PyTorch完全一致）
    for images, labels in train_dataset.create_tuple_iterator():
        print("MindSpore 数据shape:", images.shape)  # (32,3,224,224)
        print("MindSpore 标签shape:", labels.shape)
        break

五、文本分类小模型数据准备（扩展）

若迁移文本类小模型，数据准备流程一致，仅需替换为文本处理算子：

import mindspore.dataset.text as text

# 文本处理流程
vocab = text.Vocab.from_list(["我", "你", "爱", "中国"])
trans = [
    text.BertTokenizer(vocab=vocab),
    text.PadEnd(max_sequence_length=128, pad_value=0)
]
dataset = dataset.map(operations=trans, input_columns="text")

MindSpore 文本数据 pipeline 可完美对齐 PyTorch 的 Tokenizer、Pad、Truncate 逻辑。

六、迁移关键要点与精度保障

1. 数据分布严格对齐：Normalize 均值、方差、图像尺寸必须与 PyTorch 完全相同；
2. 增强概率一致：RandomFlip、RandomCrop 等概率参数保持不变；
3. Batch 与 Shuffle 对齐：保证训练过程数据读取顺序、批大小无差异；
4. 昇腾硬件优化：开启多线程、异步加载，提升 NPU 数据吞吐量；
5. 输出校验：通过 shape、dtype、数值范围三重校验，确保无损迁移。

数据准备完成后，即可进入模型结构迁移、权重加载、训练流程适配阶段，大幅降低整体迁移难度。

七、总结

在昇思 MindSpore 与 PyTorch 小模型迁移工程中，数据准备是决定迁移成败的首要环节。通过对齐数据集结构、数据增强、批处理、格式规范，可实现数据 pipeline 1:1 复刻，确保模型训练精度无损。本文提供的图像 / 文本双场景代码，覆盖自定义数据集、增强适配、昇腾硬件优化等核心能力，无需改造业务逻辑，即可快速完成小模型数据迁移。