| 110 | class MSA_emb(nn.Module): |
| 111 | # Get initial seed MSA embedding |
| 112 | def __init__(self, d_msa=256, d_pair=128, d_state=32, d_init=22+22+2+2, |
| 113 | minpos=-32, maxpos=32, p_drop=0.1, input_seq_onehot=False): |
| 114 | super(MSA_emb, self).__init__() |
| 115 | self.emb = nn.Linear(d_init, d_msa) # embedding for general MSA |
| 116 | self.emb_q = nn.Embedding(22, d_msa) # embedding for query sequence -- used for MSA embedding |
| 117 | self.emb_left = nn.Embedding(22, d_pair) # embedding for query sequence -- used for pair embedding |
| 118 | self.emb_right = nn.Embedding(22, d_pair) # embedding for query sequence -- used for pair embedding |
| 119 | self.emb_state = nn.Embedding(22, d_state) |
| 120 | self.drop = nn.Dropout(p_drop) |
| 121 | self.pos = PositionalEncoding2D(d_pair, minpos=minpos, maxpos=maxpos, p_drop=p_drop) |
| 122 | |
| 123 | self.input_seq_onehot=input_seq_onehot |
| 124 | |
| 125 | self.reset_parameter() |
| 126 | |
| 127 | def reset_parameter(self): |
| 128 | self.emb = init_lecun_normal(self.emb) |