AlphaFold2代码阅读（一）

2021SC@SDUSC

---

前言

  AlphaFold2的代码真的是太多了，而且几乎没有注释，难度真的很大。因为我对深度学习的了解也并不深刻，并且TensorFlow这个框架也没有用过，所以在看代码的过程，真的是不停地在学习。总之一句话就是痛并快乐着。
  今天要学习的是alphafold/alphafold/model/modules.py。modules.py 中主要是AlphaFold2 核心部分使用的模块和代码。

---

一、各种import

1、代码

import functools
from alphafold.common import residue_constants
from alphafold.model import all_atom
from alphafold.model import common_modules
from alphafold.model import folding
from alphafold.model import layer_stack
from alphafold.model import lddt
from alphafold.model import mapping
from alphafold.model import prng
from alphafold.model import quat_affine
from alphafold.model import utils
import haiku as hk
import jax
import jax.numpy as jnp

2、解析

1、导入了很多alphafold自己的模块

下面是各个模块的用途：

  residue_constants.py中写明了AlphaFold 中使用的常量

  all_atom.py中写明了所有原子表示的操作

  common_modules.py中写明了用于蛋白质折叠的常见俳句模块的集合

  folding.py中写明了结构模块的模块和实用程序

  layer_stack.py中写明了没有共享参数的堆叠层函数重复的函数

  lddt.py中写明了lDDT 蛋白质距离评分

  mapping.py中写明了专门的映射函数

  prng.py中写明了关于蛋白质折叠中 PRNG 使用的实用程序集合

  quat_affine.py中写明了四元数几何模块

  utils.py中写明了一组用于蛋白质折叠的 JAX 实用函数

2、导入了很多库

Haiku：
  Haiku是JAX的神经网络库，它允许用户使用熟悉的面向对象编程模型，同时允许完全访问JAX的纯函数转换。它提供了两个核心工具：模块抽象hk.Module，和一个简单的函数转换hk.transform。
  hk.Module是Python对象，包含对其自身参数、其他模块和对用户输入应用函数方法的引用。hk.transform允许完全访问JAX的纯函数转换。在后面的代码中就用到了hk.Module。

functools：

  functools是python的模块，用于高阶函数：指那些作用于函数或者返回其它函数的函数，通常只要是可以被当做函数调用的对象就是这个模块的目标。

  cmp_to_key，将一个比较函数转换关键字函数；

  partial，针对函数起作用，并且是部分的；

  reduce，与python内置的reduce函数功能一样；

  total_ordering，在类装饰器中按照缺失顺序，填充方法；

  update_wrapper，更新一个包裹（wrapper）函数，使其看起来更像被包裹（wrapped）的函数；

  wraps，可用作一个装饰器，简化调用update_wrapper的过程；

二、softmax_cross_entropy函数

1、代码

def softmax_cross_entropy(logits, labels):
  loss = -jnp.sum(labels * jax.nn.log_softmax(logits), axis=-1)
  return jnp.asarray(loss)

2、解析

上述代码是在给定 logits 和 one-hot class 标签的情况下计算 softmax 交叉熵

  softmax函数，又称归一化指数函数。它是二分类函数sigmoid在多分类上的推广，目的是将多分类的结果以概率的形式展现出来。

  交叉熵（Cross Entropy）是Shannon信息论中一个重要概念，主要用于度量两个概率分布间的差异性信息。

  有关softmax函数的详解，请见https://zhuanlan.zhihu.com/p/105722023

  有关交叉熵的详解，请见https://blog.youkuaiyun.com/qq_33542428/article/details/106482393

三、sigmoid_cross_entropy函