2021SC@SDUSC基于人工智能的多肽药物分析问题（九）

基于人工智能的多肽药物分析问题（九）

2021SC@SDUSC

1. 前言

在调用模型时根据传入的参数选择构建适当的模型

if use_templ:
    self.templ_emb = Templ_emb(d_templ=d_templ, n_att_head=n_head_templ, r_ff=r_ff, 
                               performer_opts=performer_L_opts, p_drop=0.0)
    self.pair_emb = Pair_emb_w_templ(d_model=d_pair, d_templ=d_templ, p_drop=p_drop)

前几篇文章已经分析逐层剖析了Templ_emb嵌入层的结构以及作用，接下来是其余层的分析解释

2. 代码分析

2.1 Pair_emb_w_templ层

该层的大多数模型已经在之前的文章中介绍过了，包括了嵌入层，正则层，全连接层以及二维位置信息编码层，在此处就不再详细介绍，链接如下 :

    def __init__(self, d_model=128, d_seq=21, d_templ=64, p_drop=0.1):   # 288 21 64 0.1
        super(Pair_emb_w_templ, self).__init__()
        self.d_model = d_model
        self.d_emb = d_model // 2
        self.emb = nn.Embedding(d_seq, self.d_emb)
        self.norm_templ = LayerNorm(d_templ)
        self.projection = nn.Linear(d_model + d_templ + 1, d_model)
        self.pos = PositionalEncoding2D(d_model, p_drop=p_drop)

2.2 特征提取模型

self.initial = Pair2Pair(n_layer=n_layer, n_att_head=n_head_pair,
                         n_feat=d_pair, r_ff=r_ff, p_drop=p_drop,
                         performer_L_opts=performer_L_opts)

在该模型中，首先是定义了一个自定义的Pair2Pair模型，该模型是一个两层模型的神经网络模型，第一层是坐标轴编码层AxialEncoderLayer，第二层是编码层Encoder，根据forward函数我们可以了解到该模型的作用是将传入的MODEL_PARAM模型参数编码后返回。（Pair2Pair的两层神经网络具体解释已经在上篇文章中了）

    def __init__(self, n_layer=1, n_att_head=8, n_feat=128, r_ff=4, p_drop=0.1,
                 performer_L_opts=None):
        super(Pair2Pair, self).__init__()
        enc_layer = AxialEncoderLayer(d_model=n_feat, d_ff=n_feat*r_ff,
                                      heads=n_att_head, p_drop=p_drop,
                                      performer_opts=performer_L_opts)
        self.encoder = Encoder(enc_layer, n_layer)
    
    def forward(self, x):
        return self.encoder(x)

之后，根据n_module的正负性调用之前说过的_get_clones方法，深拷贝了一份IterBlock模型的参数并保存在最外层模型中，IterBlock模型介绍如下：

2.3 IterBlock模型

该模型主要定义了msa和pair之间的转化层，共有四层神经网络，模型输入最初始的MSA提取的特征数据以及初始残基对数据，通过MSA2MSA层处理MSA特征，在经过MSA2Pair层利用处理后的MSA更新残基对数据，之后是一个往返过程，在上一层利用MSA数据更新了残基对数据之后，再分别经过Pair2Pair和Pair2MSA两层最终将MSA的参数进行更新。最终是将msa数据和pair数据均返回，代码见下方：

class IterBlock(nn.Module):
    def __init__(self, n_layer=1, d_msa=64, d_pair=128, n_head_msa=4, n_head_pair=8, r_ff=4,
                 n_resblock=1, p_drop=0.1, performer_L_opts=None, performer_N_opts=None):
        super(IterBlock, self).__init__()
        
        self.msa2msa = MSA2MSA(n_layer=n_layer, n_att_head=n_head_msa, n_feat=d_msa,
                               r_ff=r_ff, p_drop=p_drop,
                               performer_N_opts=performer_N_opts,
                               performer_L_opts=performer_L_opts)
        self.msa2pair = MSA2Pair(n_feat=d_msa, n_feat_out=d_pair, n_feat_proj=32,
                                 n_resblock=n_resblock, p_drop=p_drop, n_att_head=n_head_msa)
        self.pair2pair = Pair2Pair(n_layer=n_layer, n_att_head=n_head_pair,
                                   n_feat=d_pair, r_ff=r_ff, p_drop=p_drop,
                                   performer_L_opts=performer_L_opts)
        self.pair2msa = Pair2MSA(n_layer=n_layer, n_att_head=4, 
                                 n_feat_in=d_pair, n_feat_out=d_msa, r_ff=r_ff, p_drop=p_drop)
    def forward(self, msa, pair):
        msa, att = self.msa2msa(msa)
        pair = self.msa2pair(msa, pair, att)
        pair = self.pair2pair(pair)
        msa = self.pair2msa(pair, msa)
        return msa, pair

接下来我们可以细看一下该神经网络模型中四个层分别是如何定义的：

2.3.1 MSA2MSA层

在该模型中，共有四层，其中两层模型为一组，分别是对L和N两个维度构建的attention模型，其中，该模型的输入为MSA提取出的特征(B, N, L, K)四元组，输出为更新后的该四元组：

class MSA2MSA(nn.Module):
    def __init__(self, n_layer=1, n_att_head=8, n_feat=256, r_ff=4, p_drop=0.1,
                 performer_N_opts=None, performer_L_opts=None):
        super(MSA2MSA, self).__init__()
        enc_layer_1 = EncoderLayer(d_model=n_feat, d_ff=n_feat*r_ff,
                                   heads=n_att_head, p_drop=p_drop,
                                   use_tied=True)
                                   
        self.encoder_1 = Encoder(enc_layer_1, n_layer)
        enc_layer_2 = EncoderLayer(d_model=n_feat, d_ff=n_feat*r_ff,
                                   heads=n_att_head, p_drop=p_drop,
                                   performer_opts=performer_N_opts)
        self.encoder_2 = Encoder(enc_layer_2, n_layer)

接下来我们分析一下EncoderLayer层：该层同样是一个自定义的神经网络模型，定义了一些基于transformer - attention机制的模型，该EncoderLayer层使用预先归一化的操作来使得训练过程更加稳定。

另外在网上查找到一篇文章，介绍了预先归一化preLayerNorm对transformer模型训练的作用，因为transformer具体表现有 warm-up 阶段超参数敏感、优化过程收敛速度慢等问题，所以在优化 Transformer 结构时，除了设置初始学习率与它的衰减策略，往往还需要在训练的初始阶段设置一个非常小（接近0）的学习率，让它经过一定的迭代轮数后逐渐增长到初始的学习率，这个过程也被称作 warm-up 阶段。

Warm-up 是原始 Transformer 结构优化时的一个必备学习率调整策略。Transformer 结构对于 warm-up 的超参数（持续轮数、增长方式、初始学习率等）非常敏感，若调整不慎，往往会使得模型无法正常收敛，所以预归一化起到让transformer训练过程加速收敛的作用，

具体的介绍可以查看https://www.msra.cn/zh-cn/news/features/pre-ln-transformer这篇文章

（另外，transformer的相关知识会在之后的博文中进行补充说明。）

EncoderLayer模型与我们之前分析过的AxialEncoderLayer类似，都是基于参数判断是构建自注意力模型还是多注意力模型MultiheadAttention

在EncoderLayer层的具体代码中：

def __init__(self, d_model, d_ff, heads, p_drop=0.1, performer_opts=None, use_tied=False):
    super(EncoderLayer, self).__init__()
    self.use_performer = performer_opts is not None
    self.use_tied = use_tied
    if self.use_performer:
            self.attn = SelfAttention(dim=d_model, heads=heads, dropout=p_drop,
                                  generalized_attention=True, **performer_opts)
    elif use_tied:
        self.attn = SoftTiedMultiheadAttention(d_model, heads, dropout=p_drop)
    else:
        self.attn = MultiheadAttention(d_model, heads, dropout=p_drop)

    self.ff = FeedForwardLayer(d_model, d_ff, p_drop=p_drop)

    self.norm1 = LayerNorm(d_model)
    self.norm2 = LayerNorm(d_model)
    self.dropout1 = nn.Dropout(p_drop, inplace=True)
    self.dropout2 = nn.Dropout(p_drop, inplace=True)

最终构建的是一个自注意力模型，该模型在第八篇博客有分析，在自注意力模型的下一层是一个前馈层，最后是将前馈层输出的数据进行正则归一化。

2.3.2 MSA2Pair层

MSA2Pair模型的输入是MSA提取的四元组特征值以及初始残基对的四元组，首先将msa参数进行归一化，可以用于减少内存的占用量，之后经过一个全连接层将归一化后的数据保存至x_down变量(1)，之后再经过一个正则化层和编码层输出序列的权重参数(2)，之后经过一个池化层，之后将(1) \ (2) 两个矩阵进行拼接操作，再然后将拼接后的矩阵元素利用torch的内置函数展开成一维数据，最后根据一个ResNet残差神经网络通过多次的卷积操作更新Pair的参数，并返回。

class MSA2Pair(nn.Module):
   def forward(self, msa, pair_orig, att):

        B, N, L, _ = msa.shape

        msa = self.norm_1(msa)
        x_down = self.proj_1(msa)

        x_down = self.norm_down(x_down)
        w_seq = self.encoder(x_down).reshape(B, L, 1, N).permute(0,3,1,2)
        feat_1d = w_seq * x_down
        
        pair = self.coevol(x_down, feat_1d)

        feat_1d = feat_1d.sum(1)
        
        # query sequence info
        query = x_down[:,0] # (B,L,K)
        feat_1d = torch.cat((feat_1d, query), dim=-1)
        # tile 1D features
        left = feat_1d.unsqueeze(2).repeat(1, 1, L, 1)
        right = feat_1d.unsqueeze(1).repeat(1, L, 1, 1)

        pair_orig = self.norm_orig(pair_orig)
        pair = self.norm_new(pair)
        pair = torch.cat((pair_orig, pair, left, right, att), -1)
        pair = pair.permute(0,3,1,2).contiguous() # prep for convolution layer
        pair = self.update(pair)
        pair = pair.permute(0,2,3,1).contiguous() # (B, L, L, C)

        return pair

下面我们具体介绍一下该模型中用到的一些新的自定义层：

   def __init__(self, n_feat=64, n_feat_out=128, n_feat_proj=32,
             n_resblock=1, p_drop=0.1, n_att_head=8):
    super(MSA2Pair, self).__init__()
    # project down embedding dimension (n_feat --> n_feat_proj)
    self.norm_1 = LayerNorm(n_feat)
    self.proj_1 = nn.Linear(n_feat, n_feat_proj)
    
    self.encoder = SequenceWeight(n_feat_proj, 1, dropout=p_drop)
    self.coevol = CoevolExtractor(n_feat_proj, n_feat_out)

    # ResNet to update pair features 
    self.norm_down = LayerNorm(n_feat_proj)
    self.norm_orig = LayerNorm(n_feat_out)
    self.norm_new  = LayerNorm(n_feat_out)
    self.update = ResidualNetwork(n_resblock, n_feat_out*2+n_feat_proj*4+n_att_head, n_feat_out, n_feat_out, p_drop=p_drop)

首先是上面代码所提到的SequenceWeight层，该层模型比较简单，主要是定义了两个全连接层和一个Drop层，作用是用于提取序列的权重参数；CoevolExtractor层模型用了两层神经网络模型，正则化-全连接 层，进行父模型参数经过池化层前的操作。

这两个模型的代码如下：

class SequenceWeight(nn.Module):
    def __init__(self, d_model, heads, dropout=0.1):
        super(SequenceWeight, self).__init__()
        self.heads = heads
        self.d_model = d_model
        self.d_k = d_model // heads
        self.scale = 1.0 / math.sqrt(self.d_k)

        self.to_query = nn.Linear(d_model, d_model)
        self.to_key = nn.Linear(d_model, d_model)
        self.dropout = nn.Dropout(dropout, inplace=True)

    def forward(self, msa):
        B, N, L = msa.shape[:3]
        
        msa = msa.permute(0,2,1,3) # (B, L, N, K)
        tar_seq = msa[:,:,0].unsqueeze(2) # (B, L, 1, K)
        
        q = self.to_query(tar_seq).view(B, L, 1, self.heads, self.d_k).permute(0,1,3,2,4).contiguous() # (B, L, h, 1, k)
        k = self.to_key(msa).view(B, L, N, self.heads, self.d_k).permute(0,1,3,4,2).contiguous() # (B, L, h, k, N)
        
        q = q * self.scale
        attn = torch.matmul(q, k) # (B, L, h, 1, N)
        attn = F.softmax(attn, dim=-1)
        return self.dropout(attn)

最后是一个残差卷积网络ResNet，我们在之后的文章中继续分析

3. 总结

本周分析的代码涉及到了不少的新知识，包括注意力模型、transformer等，对于这些内容介绍性博客会在之后额外补上。