大小未知的自定义采样PyTorch Dataset如何重写__len__与__getitem__

你描述的这种带状态（像素被选中次数）、动态拒绝采样、自动跳转下一张样本的逻辑，本来就不适合PyTorch默认的map-style Dataset（基于索引访问的数据集，需要固定__len__和__getitem__），优先换用IterableDataset适配，如果你必须使用map-style格式，也可以通过预生成采样索引的方式兼容。

方案1：使用IterableDataset（推荐） #

IterableDataset不需要实现固定索引的__getitem__和精准__len__，只需要实现__iter__方法定义迭代逻辑，完美适配你的动态采样需求，示例代码如下：

import torch
from torch.utils.data import IterableDataset, DataLoader
import cv2
import numpy as np

class CustomPatchDataset(IterableDataset):
    def __init__(self, img_paths, mask_paths, crop_size=352, score_thresh=0.3, max_retry=4):
        super().__init__()
        self.img_paths = img_paths
        self.mask_paths = mask_paths
        self.crop_size = crop_size
        self.score_thresh = score_thresh
        self.max_retry = max_retry
        # 初始化像素选中计数存储，key为图像路径，value为和原图同尺寸的计数矩阵
        self.pixel_counts = {p: None for p in img_paths}
    
    def _cal_score(self, mask_patch, count_patch):
        # 可根据需求调整评分公式，此处为示例：class_A占比越高、平均选中次数越低，评分越高
        class_a_ratio = np.mean(mask_patch == 1) # 假设class_A的掩码标注值为1
        avg_selected_count = np.mean(count_patch)
        return class_a_ratio * (1 - avg_selected_count / 10) # 10为预设的单像素最大允许选中次数
    
    def __iter__(self):
        # 逐张遍历所有训练图像
        for img_path, mask_path in zip(self.img_paths, self.mask_paths):
            # 加载当前图像、class_A掩码
            img = cv2.cvtColor(cv2.imread(img_path), cv2.COLOR_BGR2RGB)
            mask = cv2.imread(mask_path, 0)
            h, w = img.shape[:2]
            # 首次加载当前图像时初始化计数矩阵
            if self.pixel_counts[img_path] is None:
                self.pixel_counts[img_path] = np.zeros((h, w), dtype=np.int32)
            count_mat = self.pixel_counts[img_path]
            
            consecutive_fail = 0
            while consecutive_fail < self.max_retry:
                # 随机生成裁剪坐标
                x1 = np.random.randint(0, w - self.crop_size + 1)
                y1 = np.random.randint(0, h - self.crop_size + 1)
                x2, y2 = x1 + self.crop_size, y1 + self.crop_size
                
                # 裁剪得到图像块、掩码块、计数块
                img_patch = img[y1:y2, x1:x2]
                mask_patch = mask[y1:y2, x1:x2]
                count_patch = count_mat[y1:y2, x1:x2]
                
                # 评分校验
                score = self._cal_score(mask_patch, count_patch)
                if score >= self.score_thresh:
                    # 更新选中区域的计数
                    count_mat[y1:y2, x1:x2] += 1
                    # 转换为训练用格式返回
                    img_tensor = torch.from_numpy(img_patch).permute(2,0,1).float() / 255.
                    mask_tensor = torch.from_numpy(mask_patch).long()
                    yield img_tensor, mask_tensor
                    consecutive_fail = 0
                else:
                    consecutive_fail += 1
    
    def __len__(self):
        # 可返回预估的每轮样本数，不需要完全精准，仅用于训练进度展示
        return len(self.img_paths) * 2

注意事项 #

如果使用多进程加载，需要额外处理进程间的计数同步和样本分流，避免重复采样，简单场景下直接用单进程加载即可满足性能需求。

方案2：适配map-style Dataset（兼容现有训练代码用） #

如果你的训练代码必须使用基于索引的map-style Dataset，可以把动态采样逻辑提前到每轮epoch开始前执行：

1. 每个epoch启动前，按照你的采样规则遍历所有图像，预先生成所有合法采样项的列表，每个采样项存储(图像路径, 裁剪左上角坐标x, 裁剪左上角坐标y)
2. __len__直接返回预生成列表的长度
3. __getitem__根据索引取出对应采样项，加载图像裁剪后返回即可
4. 每轮epoch结束后清空预生成列表，下一轮重新生成，同时更新全局像素计数矩阵

这种方式完全兼容原生DataLoader逻辑，不需要修改训练代码，仅需要额外增加每轮开始前的预采样开销，对大部分数据集来说可忽略。

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