AF3 OpenFoldDataLoader类_prep_batch_properties_probs方法解读

AlphaFold3 data_modules 模块的 OpenFoldDataLoader 类的 _prep_batch_properties_probs 方法是为每个批次数据准备 recycling 维度 的概率分布。它将根据配置文件中的设定为每个批次数据生成 recycling 轮次的概率分布，并存储到 prop_probs_tensor 中，用于后续抽样选择特定recycling维度的数据（ _add_batch_properties方法中实现）。

源代码：

    def _prep_batch_properties_probs(self):
        keyed_probs = []
        stage_cfg = self.config[self.stage]

        max_iters = self.config.common.max_recycling_iters

        if stage_cfg.uniform_recycling:
            recycling_probs = [
                1. / (max_iters + 1) for _ in range(max_iters + 1)
            ]
        else:
            recycling_probs = [
                0. for _ in range(max_iters + 1)
            ]
            recycling_probs[-1] = 1.

        keyed_probs.append(
            ("no_recycling_iters", recycling_probs)
        )

        keys, probs = zip(*keyed_probs)
        max_len = max([len(p) for p in probs])
        padding = [[0.] * (max_len - len(p)) for p in probs]

        self.prop_keys = keys
        self.prop_probs_tensor = torch.tensor(
            [p + pad for p, pad in zip(probs, padding)],
            dtype=torch.float32,
        )

源码解读：

def _prep_batch_properties_probs(self):
    keyed_probs = []
    stage_cfg = self.config[self.stage]

    max_iters = self.config.common.max_recycling_iters

• keyed_probs = []:

  • 初始化一个空的列表 keyed_probs，它将用来存储特征名（例如 no_recycling_iters）和与该特征相关的概率分布（例如每次