| 293 | return new_tokens, hit_map.reshape(b, -1, h, w) |
| 294 | |
| 295 | class Decoder(nn.Module): |
| 296 | def __init__(self, args) -> None: |
| 297 | super().__init__() |
| 298 | ''' |
| 299 | c1 :128, 120, 120 |
| 300 | c2 :256, 60, 60 |
| 301 | c3 :512, 30, 30 |
| 302 | c4 :1024, 15 ,15 |
| 303 | ''' |
| 304 | token_dim = args.token_dim |
| 305 | self.tokens = nn.Embedding(args.num_token, token_dim) |
| 306 | trunc_normal_(self.tokens.weight, std=0.02) |
| 307 | |
| 308 | dims = [1024, 512, 256, 128] |
| 309 | pe_shapes = [30, 60, 120] |
| 310 | |
| 311 | self.layers = [] |
| 312 | for pe_shape in pe_shapes: |
| 313 | self.layers.append(LoadLayer(token_dim, drop=.1, bias=False, pe_shape=pe_shape)) |
| 314 | self.cgattention1 = CGAttention(token_dim=token_dim, |
| 315 | vis_dim=token_dim, |
| 316 | hidden_dim=token_dim, |
| 317 | drop=.1, |
| 318 | bias=True) |
| 319 | self.cgattention2 = CGAttention(token_dim=token_dim, |
| 320 | vis_dim=token_dim, |
| 321 | hidden_dim=token_dim, |
| 322 | drop=.1, |
| 323 | bias=True) |
| 324 | self.layers = nn.ModuleList(self.layers) |
| 325 | self.fuses = [] |
| 326 | for dim in [dims[0], dims[2], dims[3]]: |
| 327 | self.fuses.append(Fusion(dim, token_dim, token_dim, bias=True)) |
| 328 | self.fuses = nn.ModuleList(self.fuses) |
| 329 | self.proj = DProjector(text_dim=token_dim, in_dim=token_dim) |
| 330 | |
| 331 | def forward(self, vis, text, pad_mask): |
| 332 | x_c4, x_c3, x_c2, x_c1 = vis |
| 333 | tokens = self.tokens.weight[None,...].expand(x_c1.shape[0], -1, -1) |
| 334 | maps = [] |
| 335 | v = x_c4 |
| 336 | for load, layer, fuse, v_ in zip(self.layers,[self.cgattention1,self.cgattention2,self.cgattention2], self.fuses, [x_c3, x_c2, x_c1]): |
| 337 | v = fuse(v, v_) |
| 338 | tokens, pe = load(tokens, text, pad_mask) |
| 339 | tokens, hitmap = layer(tokens, v, pe=pe) |
| 340 | maps.append(hitmap) |
| 341 | out = self.proj(v, tokens[:,-1]) |
| 342 | return out, maps |