使用torchvision.datasets.ImageFolder生成X_train, X_test, y_train, y_test,X_cal,Y_cal并保存为data.pkl文件

在 PyTorch 中，我们可以使用 Pickle 库将数据保存为 pkl 文件。具体实现如下：

import torchvision.datasets as datasets
import torch.utils.data as data
import pickle

# 定义数据集路径
data_path = 'path/to/data'

# 定义数据增强和标准化操作
data_transforms = transforms.Compose([
    transforms.RandomResizedCrop(224),
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

# 加载数据集
image_dataset = datasets.ImageFolder(root=data_path, transform=data_transforms)

# 将数据集分割为训练集、测试集和验证集
train_size = int(0.7 * len(image_dataset))
test_size = int(0.2 * len(image_dataset))
cal_size = len(image_dataset) - train_size - test_size
train_dataset, test_dataset, cal_dataset = data.random_split(image_dataset, [train_size, test_size, cal_size])

# 创建数据加载器
train_loader = data.DataLoader(train_dataset, batch_size=32, shuffle=True, num_workers=4)
test_loader = data.DataLoader(test_dataset, batch_size=32, shuffle=True, num_workers=4)
cal_loader = data.DataLoader(cal_dataset, batch_size=32, shuffle=True, num_workers=4)

# 获取训练集、测试集和验证集的数据和标签
X_train, y_train = next(iter(train_loader))
X_test, y_test = next(iter(test_loader))
X_cal, y_cal = next(iter(cal_loader))

# 将数据保存为 pkl 文件
with open('data.pkl', 'wb') as f:
    pickle.dump((X_train, y_train, X_test, y_test, X_cal, y_cal), f)

在代码中，我们首先定义数据集路径和数据增强、标准化操作。然后使用 torchvision.datasets.ImageFolder 函数加载图像数据集，再使用 torch.utils.data.random_split 函数将数据集分割为训练集、测试集和验证集。接着，我们创建数据加载器，并使用 next(iter(loader)) 函数获取训练集、测试集和验证集的数据和标签。最后，我们使用 Pickle 库将数据保存为 pkl 文件。