An attention block that allows spatial positions to attend to each other. Originally ported from here, but adapted to the N-d case. https://github.com/hojonathanho/diffusion/blob/1e0dceb3b3495bbe19116a5e1b3596cd0706c543/diffusion_tf/models/unet.py#L66.
| 328 | |
| 329 | |
| 330 | class AttentionBlock(nn.Module): |
| 331 | """ |
| 332 | An attention block that allows spatial positions to attend to each other. |
| 333 | |
| 334 | Originally ported from here, but adapted to the N-d case. |
| 335 | https://github.com/hojonathanho/diffusion/blob/1e0dceb3b3495bbe19116a5e1b3596cd0706c543/diffusion_tf/models/unet.py#L66. |
| 336 | """ |
| 337 | |
| 338 | def __init__( |
| 339 | self, |
| 340 | channels, |
| 341 | num_heads=1, |
| 342 | num_head_channels=-1, |
| 343 | use_checkpoint=False, |
| 344 | use_new_attention_order=False, |
| 345 | ): |
| 346 | super().__init__() |
| 347 | self.channels = channels |
| 348 | if num_head_channels == -1: |
| 349 | self.num_heads = num_heads |
| 350 | else: |
| 351 | assert ( |
| 352 | channels % num_head_channels == 0 |
| 353 | ), f"q,k,v channels {channels} is not divisible by num_head_channels {num_head_channels}" |
| 354 | self.num_heads = channels // num_head_channels |
| 355 | self.use_checkpoint = use_checkpoint |
| 356 | self.norm = normalization(channels) |
| 357 | self.qkv = conv_nd(1, channels, channels * 3, 1) |
| 358 | if use_new_attention_order: |
| 359 | # split qkv before split heads |
| 360 | self.attention = QKVAttention(self.num_heads) |
| 361 | else: |
| 362 | # split heads before split qkv |
| 363 | self.attention = QKVAttentionLegacy(self.num_heads) |
| 364 | |
| 365 | self.proj_out = zero_module(conv_nd(1, channels, channels, 1)) |
| 366 | |
| 367 | def forward(self, x): |
| 368 | return checkpoint(self._forward, (x,), self.parameters(), True) |
| 369 | |
| 370 | def _forward(self, x): |
| 371 | b, c, *spatial = x.shape |
| 372 | x = x.reshape(b, c, -1) |
| 373 | qkv = self.qkv(self.norm(x)) |
| 374 | h = self.attention(qkv) |
| 375 | h = self.proj_out(h) |
| 376 | return (x + h).reshape(b, c, *spatial) |
| 377 | |
| 378 | |
| 379 | def count_flops_attn(model, _x, y): |