如何在PyTorch中按苹果样本组划分图像数据集为训练/验证集

要实现每个苹果的6张图像作为整体划分，核心思路是先按苹果ID分组，再对苹果分组进行训练/验证划分，而非直接拆分单个样本。具体步骤如下：

1. 读取数据集并按苹果分组 #

先用Pandas读取CSV文件，通过Apple字段将同属一个苹果的所有样本归为一组：

import pandas as pd

# 读取CSV数据集
df = pd.read_csv("your_dataset.csv")
# 按Apple字段分组，得到每个苹果对应的所有样本索引
apple_groups = df.groupby("Apple").groups
# 提取所有苹果的唯一ID列表
apple_unique_ids = list(apple_groups.keys())

2. 划分训练/验证用的苹果ID #

对苹果的唯一ID进行划分，确保每个苹果整体进入训练或验证集：

from sklearn.model_selection import train_test_split

# 按8:2比例划分苹果ID，random_state保证划分可复现
train_apple_ids, val_apple_ids = train_test_split(
    apple_unique_ids, test_size=0.2, random_state=42
)

如果不想依赖sklearn，也可以手动随机划分：

import random

random.seed(42)
random.shuffle(apple_unique_ids)
split_idx = int(len(apple_unique_ids) * 0.8)
train_apple_ids = apple_unique_ids[:split_idx]
val_apple_ids = apple_unique_ids[split_idx:]

3. 生成训练/验证集的样本索引 #

将划分好的苹果ID对应的所有样本索引收集起来：

# 收集训练集所有样本的索引
train_indices = []
for apple_id in train_apple_ids:
    train_indices.extend(apple_groups[apple_id])

# 收集验证集所有样本的索引
val_indices = []
for apple_id in val_apple_ids:
    val_indices.extend(apple_groups[apple_id])

4. 用PyTorch Subset创建数据集 #

基于自定义数据集类，通过Subset快速生成训练和验证集：

from torch.utils.data import Dataset, Subset

# 自定义苹果数据集类（根据你的实际需求调整）
class AppleDataset(Dataset):
    def __init__(self, csv_path):
        self.df = pd.read_csv(csv_path)
        # 可在此提前加载图像路径、预处理配置等

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

    def __getitem__(self, idx):
        # 根据索引获取单样本数据：图像、糖度标签等
        row = self.df.iloc[idx]
        img_path = row["image_path"]
        sugar_level = row["sugar_content"]
        # 此处添加图像加载、预处理逻辑
        # img = load_and_preprocess(img_path)
        return img, sugar_level

# 创建完整数据集
full_dataset = AppleDataset("your_dataset.csv")

# 生成训练集和验证集
train_dataset = Subset(full_dataset, train_indices)
val_dataset = Subset(full_dataset, val_indices)

5. 用DataLoader加载数据集 #

最后用PyTorch的DataLoader封装，用于模型训练：

from torch.utils.data import DataLoader

train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=32, shuffle=False)

关键说明 #

这种方式的核心是以苹果为单位划分，避免了单个苹果的图像被拆分到不同集合，完全符合你的需求。如果需要调整划分比例，只需修改test_size参数或手动拆分的比例即可。

内容的提问来源于stack exchange，提问作者M.K.