| 54 | |
| 55 | |
| 56 | class AttentionPool2d(nn.Module): |
| 57 | def __init__(self, spacial_dim: int, embed_dim: int, num_heads: int, output_dim: int = None): |
| 58 | super().__init__() |
| 59 | self.positional_embedding = nn.Parameter(torch.randn(spacial_dim ** 2 + 1, embed_dim) / embed_dim ** 0.5) |
| 60 | self.k_proj = nn.Linear(embed_dim, embed_dim) |
| 61 | self.q_proj = nn.Linear(embed_dim, embed_dim) |
| 62 | self.v_proj = nn.Linear(embed_dim, embed_dim) |
| 63 | self.c_proj = nn.Linear(embed_dim, output_dim or embed_dim) |
| 64 | self.num_heads = num_heads |
| 65 | |
| 66 | def forward(self, x): |
| 67 | x = x.reshape(x.shape[0], x.shape[1], x.shape[2] * x.shape[3]).permute(2, 0, 1) # NCHW -> (HW)NC |
| 68 | x = torch.cat([x.mean(dim=0, keepdim=True), x], dim=0) # (HW+1)NC |
| 69 | x = x + self.positional_embedding[:, None, :].to(x.dtype) # (HW+1)NC |
| 70 | x, _ = multi_head_attention_forward( |
| 71 | query=x, key=x, value=x, |
| 72 | embed_dim_to_check=x.shape[-1], |
| 73 | num_heads=self.num_heads, |
| 74 | q_proj_weight=self.q_proj.weight, |
| 75 | k_proj_weight=self.k_proj.weight, |
| 76 | v_proj_weight=self.v_proj.weight, |
| 77 | in_proj_weight=None, |
| 78 | in_proj_bias=torch.cat([self.q_proj.bias, self.k_proj.bias, self.v_proj.bias]), |
| 79 | bias_k=None, |
| 80 | bias_v=None, |
| 81 | add_zero_attn=False, |
| 82 | dropout_p=0, |
| 83 | out_proj_weight=self.c_proj.weight, |
| 84 | out_proj_bias=self.c_proj.bias, |
| 85 | use_separate_proj_weight=True, |
| 86 | training=self.training, |
| 87 | need_weights=False |
| 88 | ) |
| 89 | |
| 90 | return x[0] |
| 91 | |
| 92 | |
| 93 | class ModifiedResNet(nn.Module): |