AlphaFold2代码阅读（六）

最新推荐文章于 2024-04-23 21:01:17 发布

xiaofengzihhh

最新推荐文章于 2024-04-23 21:01:17 发布

阅读量697

点赞数

分类专栏：蛋白质结构预测文章标签： python 深度学习机器学习

本文链接：https://blog.csdn.net/xiaofengzihhh/article/details/121442687

版权

2021SC@SDUSC

1、class Transition


class Transition(hk.Module):

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

  def __call__(self, act, mask, is_training=True):

    _, _, nc = act.shape

    num_intermediate = int(nc * self.config.num_intermediate_factor)
    mask = jnp.expand_dims(mask, axis=-1)

    act = hk.LayerNorm(
        axis=[-1],
        create_scale=True,
        create_offset