(self, x)
| 66 | self.num_heads = num_heads |
| 67 | |
| 68 | def forward(self, x): |
| 69 | x = x.flatten(start_dim=2).permute(2, 0, 1) # NCHW -> (HW)NC |
| 70 | x = torch.cat([x.mean(dim=0, keepdim=True), x], dim=0) # (HW+1)NC |
| 71 | x = x + self.positional_embedding[:, None, :].to(x.dtype) # (HW+1)NC |
| 72 | x, _ = F.multi_head_attention_forward( |
| 73 | query=x[:1], key=x, value=x, |
| 74 | embed_dim_to_check=x.shape[-1], |
| 75 | num_heads=self.num_heads, |
| 76 | q_proj_weight=self.q_proj.weight, |
| 77 | k_proj_weight=self.k_proj.weight, |
| 78 | v_proj_weight=self.v_proj.weight, |
| 79 | in_proj_weight=None, |
| 80 | in_proj_bias=torch.cat([self.q_proj.bias, self.k_proj.bias, self.v_proj.bias]), |
| 81 | bias_k=None, |
| 82 | bias_v=None, |
| 83 | add_zero_attn=False, |
| 84 | dropout_p=0, |
| 85 | out_proj_weight=self.c_proj.weight, |
| 86 | out_proj_bias=self.c_proj.bias, |
| 87 | use_separate_proj_weight=True, |
| 88 | training=self.training, |
| 89 | need_weights=False |
| 90 | ) |
| 91 | return x.squeeze(0) |
| 92 | |
| 93 | |
| 94 | class ModifiedResNet(nn.Module): |
nothing calls this directly
no outgoing calls
no test coverage detected