A 2D perceiver-resampler network with one cross attention layers by (grid_size**2) learnable queries and 2d sincos pos_emb Outputs: A tensor with the shape of (grid_size**2, embed_dim)
| 90 | |
| 91 | |
| 92 | class Resampler(nn.Module): |
| 93 | """ |
| 94 | A 2D perceiver-resampler network with one cross attention layers by |
| 95 | (grid_size**2) learnable queries and 2d sincos pos_emb |
| 96 | Outputs: |
| 97 | A tensor with the shape of (grid_size**2, embed_dim) |
| 98 | """ |
| 99 | def __init__( |
| 100 | self, |
| 101 | grid_size, |
| 102 | embed_dim, |
| 103 | num_heads, |
| 104 | kv_dim=None, |
| 105 | norm_layer=nn.LayerNorm |
| 106 | ): |
| 107 | super().__init__() |
| 108 | self.num_queries = grid_size ** 2 |
| 109 | self.embed_dim = embed_dim |
| 110 | self.num_heads = num_heads |
| 111 | |
| 112 | self.pos_embed = nn.Parameter( |
| 113 | torch.from_numpy(get_2d_sincos_pos_embed(embed_dim, grid_size)).float() |
| 114 | ).requires_grad_(False) |
| 115 | |
| 116 | self.query = nn.Parameter(torch.zeros(self.num_queries, embed_dim)) |
| 117 | trunc_normal_(self.query, std=.02) |
| 118 | |
| 119 | if kv_dim is not None and kv_dim != embed_dim: |
| 120 | self.kv_proj = nn.Linear(kv_dim, embed_dim, bias=False) |
| 121 | else: |
| 122 | self.kv_proj = nn.Identity() |
| 123 | |
| 124 | self.attn = nn.MultiheadAttention(embed_dim, num_heads) |
| 125 | self.ln_q = norm_layer(embed_dim) |
| 126 | self.ln_kv = norm_layer(embed_dim) |
| 127 | |
| 128 | # self.apply(self._init_weights) |
| 129 | |
| 130 | def _init_weights(self, m): |
| 131 | if isinstance(m, nn.Linear): |
| 132 | trunc_normal_(m.weight, std=.02) |
| 133 | if isinstance(m, nn.Linear) and m.bias is not None: |
| 134 | nn.init.constant_(m.bias, 0) |
| 135 | elif isinstance(m, nn.LayerNorm): |
| 136 | nn.init.constant_(m.bias, 0) |
| 137 | nn.init.constant_(m.weight, 1.0) |
| 138 | |
| 139 | def forward(self, x, attn_mask=None): |
| 140 | |
| 141 | pos_embed = get_abs_pos(self.pos_embed, x.size(1)) |
| 142 | |
| 143 | x = self.kv_proj(x) |
| 144 | x = self.ln_kv(x).permute(1, 0, 2) |
| 145 | |
| 146 | N = x.shape[1] |
| 147 | q = self.ln_q(self.query) |
| 148 | out = self.attn( |
| 149 | self._repeat(q, N) + self.pos_embed.unsqueeze(1), |