haiku实现TemplatePairStack类

TemplatePairStack是实现蛋白质结构模版pair_act特征表示的类:
通过layer_stack.layer_stack(c.num_block)(block) 堆叠c.num_block(配置文件中为2)block 函数,每个block对输入pair_act 和 pair_mask执行计算流程:TriangleAttention ---> dropout ->TriangleAttention ---> dropout -> TriangleMultiplication ---> dropout -> TriangleMultiplication ---> dropout -> Transition

复制代码
import haiku as hk


class TemplatePairStack(hk.Module):
  """Pair stack for the templates.

  Jumper et al. (2021) Suppl. Alg. 16 "TemplatePairStack"
  """

  def __init__(self, config, global_config, name='template_pair_stack'):
    super().__init__(name=name)
    self.config = config
    self.global_config = global_config

  def __call__(self, pair_act, pair_mask, is_training, safe_key=None):
    """Builds TemplatePairStack module.

    Arguments:
      pair_act: Pair activations for single template, shape [N_res, N_res, c_t].
      pair_mask: Pair mask, shape [N_res, N_res].
      is_training: Whether the module is in training mode.
      safe_key: Safe key object encapsulating the random number generation key.

    Returns:
      Updated pair_act, shape [N_res, N_res, c_t].
    """

    if safe_key is None:
      safe_key = prng.SafeKey(hk.next_rng_key())

    gc = self.global_config
    c = self.config

    if not c.num_block:
      return pair_act

    def block(x):
      """One block of the template pair stack."""
      pair_act, safe_key = x

      dropout_wrapper_fn = functools.partial(
          dropout_wrapper, is_training=is_training, global_config=gc)

      safe_key, *sub_keys = safe_key.split(6)
      sub_keys = iter(sub_keys)

      pair_act = dropout_wrapper_fn(
          TriangleAttention(c.triangle_attention_starting_node, gc,
                            name='triangle_attention_starting_node'),
          pair_act,
          pair_mask,
          next(sub_keys))
      pair_act = dropout_wrapper_fn(
          TriangleAttention(c.triangle_attention_ending_node, gc,
                            name='triangle_attention_ending_node'),
          pair_act,
          pair_mask,
          next(sub_keys))
      pair_act = dropout_wrapper_fn(
          TriangleMultiplication(c.triangle_multiplication_outgoing, gc,
                                 name='triangle_multiplication_outgoing'),
          pair_act,
          pair_mask,
          next(sub_keys))
      pair_act = dropout_wrapper_fn(
          TriangleMultiplication(c.triangle_multiplication_incoming, gc,
                                 name='triangle_multiplication_incoming'),
          pair_act,
          pair_mask,
          next(sub_keys))
      pair_act = dropout_wrapper_fn(
          Transition(c.pair_transition, gc, name='pair_transition'),
          pair_act,
          pair_mask,
          next(sub_keys))

      return pair_act, safe_key

    if gc.use_remat:
      block = hk.remat(block)

    res_stack = layer_stack.layer_stack(c.num_block)(block)
    pair_act, safe_key = res_stack((pair_act, safe_key))
    return pair_act
相关推荐
飞哥数智坊2 小时前
从CodeBuddy翻车到MasterGo救场,我的小程序UI终于焕然一新
人工智能
AKAMAI5 小时前
跳过复杂环节:Akamai应用平台让Kubernetes生产就绪——现已正式发布
人工智能·云原生·云计算
databook6 小时前
Manim实现闪光轨迹特效
后端·python·动效
新智元6 小时前
阿里王牌 Agent 横扫 SOTA,全栈开源力压 OpenAI!博士级难题一键搞定
人工智能·openai
新智元6 小时前
刚刚,OpenAI/Gemini 共斩 ICPC 2025 金牌!OpenAI 满分碾压横扫全场
人工智能·openai
机器之心6 小时前
OneSearch,揭开快手电商搜索「一步到位」的秘技
人工智能·openai
阿里云大数据AI技术7 小时前
2025云栖大会·大数据AI参会攻略请查收!
大数据·人工智能
Juchecar7 小时前
解惑:NumPy 中 ndarray.ndim 到底是什么?
python
YourKing7 小时前
yolov11n.onnx格式模型转换与图像推理
人工智能
sans_7 小时前
NCCL的用户缓冲区注册
人工智能