(
self,
grid_size,
embed_dim,
num_heads,
kv_dim=None,
norm_layer=nn.LayerNorm
)
| 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): |
no test coverage detected