| 271 | |
| 272 | |
| 273 | class DecoderLayer(nn.Module): |
| 274 | def __init__(self, cfg, use_lora=False): |
| 275 | super(DecoderLayer,self).__init__() |
| 276 | self.local_ch = cfg.hrnet.local_ch |
| 277 | self.joint_ch = cfg.hyponet.joint_ch + self.local_ch |
| 278 | self.dropout_rate = 0.1 |
| 279 | self.joints = cfg.hyponet.num_joints |
| 280 | self.edges = cfg.hyponet.num_twists |
| 281 | self.num_item = self.edges+self.joints |
| 282 | self.use_lora = use_lora |
| 283 | feedforward_dim = self.joint_ch * 4 |
| 284 | |
| 285 | self.norm1 = nn.LayerNorm(self.joint_ch) |
| 286 | if self.use_lora: |
| 287 | self.self_attn = LoRAMultiheadAttention(embed_dim=self.joint_ch, num_heads=cfg.hyponet.heads, lora_rank=4, batch_first=True) |
| 288 | else: |
| 289 | self.self_attn = nn.MultiheadAttention(embed_dim=self.joint_ch, num_heads=cfg.hyponet.heads, batch_first = True) |
| 290 | self.dropout1 = nn.Dropout(p=self.dropout_rate) |
| 291 | |
| 292 | self.norm2 = nn.LayerNorm(self.joint_ch) |
| 293 | if self.use_lora: |
| 294 | self.multihead_attn = LoRAMultiheadAttention(embed_dim=self.joint_ch, num_heads=cfg.hyponet.heads, lora_rank=4, batch_first = True) |
| 295 | else: |
| 296 | self.multihead_attn = nn.MultiheadAttention(embed_dim=self.joint_ch, num_heads=cfg.hyponet.heads,batch_first = True) |
| 297 | self.dropout2 = nn.Dropout(p=self.dropout_rate) |
| 298 | |
| 299 | self.linear1 = nn.Linear(self.joint_ch, feedforward_dim) |
| 300 | self.dropout = nn.Dropout(p=self.dropout_rate) |
| 301 | self.linear2 = nn.Linear(feedforward_dim, self.joint_ch) |
| 302 | self.norm3 = nn.LayerNorm(self.joint_ch) |
| 303 | self.dropout3 = nn.Dropout(p=self.dropout_rate) |
| 304 | self.activation = nn.ReLU() |
| 305 | |
| 306 | def with_pos_embed(self, tensor, pos): |
| 307 | return tensor + pos |
| 308 | |
| 309 | def forward(self, tgt, memory, mask= None, mask_ctx = None, pos= None, pos_ctx=None, gen_multi=False): |
| 310 | tgt2 = self.norm1(tgt) |
| 311 | q = k = self.with_pos_embed(tgt2, pos) |
| 312 | tgt2 = self.self_attn(q, k, value=tgt2, attn_mask=mask)[0] |
| 313 | tgt = tgt + self.dropout1(tgt2) |
| 314 | |
| 315 | if gen_multi: |
| 316 | bs = memory.shape[0] |
| 317 | multi_n = tgt.shape[0] // bs |
| 318 | pos = pos.repeat(1,multi_n,1).view(-1,self.num_item*multi_n,self.joint_ch) |
| 319 | tgt2 = self.norm2(tgt).view(bs,-1,self.joint_ch) |
| 320 | tgt2 = self.multihead_attn(query=self.with_pos_embed(tgt2, pos), |
| 321 | key=self.with_pos_embed(memory, pos_ctx), |
| 322 | value=memory, attn_mask=None)[0].contiguous().view(bs*multi_n,-1,self.joint_ch) |
| 323 | tgt = tgt + self.dropout2(tgt2) |
| 324 | else: |
| 325 | tgt2 = self.norm2(tgt) |
| 326 | tgt2 = self.multihead_attn(query=self.with_pos_embed(tgt2, pos), |
| 327 | key=self.with_pos_embed(memory, pos_ctx), |
| 328 | value=memory, attn_mask=None)[0] |
| 329 | tgt = tgt + self.dropout2(tgt2) |
| 330 | |