| 74 | layer.weight.data.mul_(norm / layer.weight.data.norm(2).item()) |
| 75 | |
| 76 | class DecoderLayer(nn.Module): |
| 77 | def __init__(self, cfg,dim): |
| 78 | super(DecoderLayer,self).__init__() |
| 79 | self.joint_ch = dim |
| 80 | self.self_attn = nn.MultiheadAttention(embed_dim=self.joint_ch, num_heads=cfg.scorenet.heads,batch_first = True) |
| 81 | self.multihead_attn = nn.MultiheadAttention(embed_dim=self.joint_ch, num_heads=cfg.scorenet.heads,batch_first = True) |
| 82 | self.dropout_rate = 0.1 |
| 83 | feedforward_dim = self.joint_ch * 4 |
| 84 | # MLP |
| 85 | self.linear1 = nn.Linear(self.joint_ch, feedforward_dim) |
| 86 | self.dropout = nn.Dropout(p=self.dropout_rate) |
| 87 | self.linear2 = nn.Linear(feedforward_dim, self.joint_ch) |
| 88 | |
| 89 | # Layer Normalization & Dropout |
| 90 | self.norm1 = nn.LayerNorm(self.joint_ch) |
| 91 | self.norm2 = nn.LayerNorm(self.joint_ch) |
| 92 | self.norm3 = nn.LayerNorm(self.joint_ch) |
| 93 | self.dropout1 = nn.Dropout(p=self.dropout_rate) |
| 94 | self.dropout2 = nn.Dropout(p=self.dropout_rate) |
| 95 | self.dropout3 = nn.Dropout(p=self.dropout_rate) |
| 96 | self.activation = nn.ReLU() |
| 97 | |
| 98 | def with_pos_embed(self, tensor, pos): |
| 99 | return tensor + pos |
| 100 | |
| 101 | def forward(self, tgt, memory, mask=None, mask_ctx=None, pos=None, pos_ctx=None): |
| 102 | tgt2 = self.norm1(tgt) |
| 103 | q = k = self.with_pos_embed(tgt2, pos[0]) |
| 104 | tgt2 = self.self_attn(q, k, value=tgt2, attn_mask=mask)[0] |
| 105 | tgt = tgt + self.dropout1(tgt2) |
| 106 | tgt2 = self.norm2(tgt) |
| 107 | tgt2 = tgt2.view(memory.shape[0], -1, self.joint_ch) |
| 108 | tgt2 = self.multihead_attn(query=self.with_pos_embed(tgt2, pos[1]), |
| 109 | key=self.with_pos_embed(memory, pos_ctx), |
| 110 | value=memory)[0] |
| 111 | tgt2 = tgt2.contiguous().view(tgt.shape[0], tgt.shape[1], tgt.shape[2]) |
| 112 | tgt = tgt + self.dropout2(tgt2) |
| 113 | tgt2 = self.norm3(tgt) |
| 114 | tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt2)))) |
| 115 | tgt = tgt + self.dropout3(tgt2) |
| 116 | return tgt |
| 117 | |
| 118 | |
| 119 | class ScoreNet(nn.Module): |