| 159 | nn.init.constant_(self.mlp[-1].bias, 0.) |
| 160 | |
| 161 | def attention(self, x, y): |
| 162 | d_model = self.ln_1.weight.size(0) |
| 163 | q = (x @ self.attn.in_proj_weight[:d_model].T) + self.attn.in_proj_bias[:d_model] |
| 164 | |
| 165 | k = (y @ self.attn.in_proj_weight[d_model:-d_model].T) + self.attn.in_proj_bias[d_model:-d_model] |
| 166 | v = (y @ self.attn.in_proj_weight[-d_model:].T) + self.attn.in_proj_bias[-d_model:] |
| 167 | Tx, Ty, N = q.size(0), k.size(0), q.size(1) |
| 168 | q = q.view(Tx, N, self.attn.num_heads, self.attn.head_dim).permute(1, 2, 0, 3) |
| 169 | k = k.view(Ty, N, self.attn.num_heads, self.attn.head_dim).permute(1, 2, 0, 3) |
| 170 | v = v.view(Ty, N, self.attn.num_heads, self.attn.head_dim).permute(1, 2, 0, 3) |
| 171 | aff = (q @ k.transpose(-2, -1) / (self.attn.head_dim ** 0.5)) |
| 172 | |
| 173 | aff = aff.softmax(dim=-1) |
| 174 | out = aff @ v |
| 175 | out = out.permute(2, 0, 1, 3).flatten(2) |
| 176 | out = self.attn.out_proj(out) |
| 177 | return out |
| 178 | |
| 179 | def forward(self, x, y): |
| 180 | x = x + self.drop_path(self.attention(self.ln_1(x), self.ln_3(y))) |