pytorch中如何使用自己的数据集进行训练

需求

cifar10中有十个类别的图像，我需要其中的第一类和第二类作为数据集，重新构建训练集和测试集，用这份小数据集来训练一个diffusion model

get data

import os
import pickle
import numpy as np
from PIL import Image

# 替换为你的 CIFAR-10 pickle 文件路径
pickle_file_path = "train_batch_class12.pickle"

# 读取 CIFAR-10 pickle 文件内容
with open(pickle_file_path, "rb") as f:
    cifar10_data = pickle.load(f)

# 提取图像数据和标签
image_data = cifar10_data['data']
labels = cifar10_data['labels']

# 类别名称的映射字典
class_names = {
   
    0: "airplane",
    1: "automobile",
    # 添加其他类别的映射
}

# 循环处理每个图像
for i in range(len(image_data)):
    # 获取单个图像和对应标签
    img_array = image_data[i]
    label = labels[i]

    # 调整图像数据形状和数据类型
    img_array = img_array.reshape((3, 32, 32)).transpose((1, 2, 0))  # 调整形状和通道顺序
    img_array = img_array.astype(np.uint8)

    # 将 NumPy 数组转换为 PIL Image
    image = Image.fromarray(img_array)

    # 获取类别名称
    class_name = class_names.get(label, f"class{
     label}")

    # 创建存储文件夹的路径
    class_folder_path = f"train/{
     class_name}"
    os.makedirs(class_folder_path, exist_ok=True)

    # 保存图像到对应的类别文件夹
    image.save(os.path.join(class_folder_path, f"cifar10_image_{
     i}_label_{
     label}.png"))
    # 如果需要显示图像，取消注释下面这行
    # image.show()

import os
import pickle
import numpy as np
from PIL import Image

# 替换为你的 CIFAR-10 pickle 文件路径
pickle_file_path = "test_batch_class12.pickle"

# 读取 CIFAR-10 pickle 文件内容
with open(pickle_file_path, "rb") as f:
    cifar10_data = pickle.load(f)

# 提取图像数据和标签
image_data = cifar10_data['data']
labels