A module which performs QKV attention and splits in a different order.
| 380 | |
| 381 | |
| 382 | class QKVAttention(nn.Module): |
| 383 | """ |
| 384 | A module which performs QKV attention and splits in a different order. |
| 385 | """ |
| 386 | |
| 387 | def __init__(self, n_heads): |
| 388 | super().__init__() |
| 389 | self.n_heads = n_heads |
| 390 | |
| 391 | def forward(self, qkv): |
| 392 | """ |
| 393 | Apply QKV attention. |
| 394 | :param qkv: an [N x (3 * H * C) x T] tensor of Qs, Ks, and Vs. |
| 395 | :return: an [N x (H * C) x T] tensor after attention. |
| 396 | """ |
| 397 | bs, width, length = qkv.shape |
| 398 | assert width % (3 * self.n_heads) == 0 |
| 399 | ch = width // (3 * self.n_heads) |
| 400 | q, k, v = qkv.chunk(3, dim=1) |
| 401 | scale = 1 / math.sqrt(math.sqrt(ch)) |
| 402 | weight = th.einsum( |
| 403 | "bct,bcs->bts", |
| 404 | (q * scale).view(bs * self.n_heads, ch, length), |
| 405 | (k * scale).view(bs * self.n_heads, ch, length), |
| 406 | ) # More stable with f16 than dividing afterwards |
| 407 | weight = th.softmax(weight.float(), dim=-1).type(weight.dtype) |
| 408 | a = th.einsum("bts,bcs->bct", weight, v.reshape(bs * self.n_heads, ch, length)) |
| 409 | return a.reshape(bs, -1, length) |
| 410 | |
| 411 | @staticmethod |
| 412 | def count_flops(model, _x, y): |
| 413 | return count_flops_attn(model, _x, y) |
| 414 | |
| 415 | |
| 416 | class Timestep(nn.Module): |