| 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 | |