AF3 TransformDataset类解读

AlphaFold3 protein_datamodule 模块 TransformDataset 类是通用的包装类，用于对已有的 torch.utils.data.Dataset 数据集对象进行统一的转换操作（transform），特别适合在数据预处理和增强过程中使用。

源代码：

class TransformDataset(torch.utils.data.Dataset):
    """A convenience class that applies arbitrary transforms to torch.utils.data.Dataset objects."""

    def __init__(
            self,
            dataset: Dataset,
            transform: Callable):
        super().__init__()
        self.dataset = dataset
        self.transform = transform

    def __getitem__(self, idx):
        sample = self.dataset[idx]
        return self.transform(sample)

    def __len__(self):
        return self.dataset.__len__()

源码解读：

类定义与用途

class TransformDataset(torch.utils.data.Dataset):
    """A convenience class that applies arbitrary