(
self,
channels,
num_heads=1,
num_head_channels=-1,
use_checkpoint=False,
use_new_attention_order=False,
)
| 468 | """ |
| 469 | |
| 470 | def __init__( |
| 471 | self, |
| 472 | channels, |
| 473 | num_heads=1, |
| 474 | num_head_channels=-1, |
| 475 | use_checkpoint=False, |
| 476 | use_new_attention_order=False, |
| 477 | ): |
| 478 | super().__init__() |
| 479 | self.channels = channels |
| 480 | if num_head_channels == -1: |
| 481 | self.num_heads = num_heads |
| 482 | else: |
| 483 | assert ( |
| 484 | channels % num_head_channels == 0 |
| 485 | ), f"q,k,v channels {channels} is not divisible by num_head_channels {num_head_channels}" |
| 486 | self.num_heads = channels // num_head_channels |
| 487 | self.use_checkpoint = use_checkpoint |
| 488 | self.norm = normalization(channels) |
| 489 | self.qkv = conv_nd(1, channels, channels * 3, 1) |
| 490 | if use_new_attention_order: |
| 491 | # split qkv before split heads |
| 492 | self.attention = QKVAttention(self.num_heads) |
| 493 | else: |
| 494 | # split heads before split qkv |
| 495 | self.attention = QKVAttentionLegacy(self.num_heads) |
| 496 | |
| 497 | self.proj_out = zero_module(conv_nd(1, channels, channels, 1)) |
| 498 | |
| 499 | def forward(self, x): |
| 500 | return checkpoint(self._forward, (x,), self.parameters(), True) # TODO: check checkpoint usage, is True # TODO: fix the .half call!!! |
nothing calls this directly
no test coverage detected