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