Causal self-attention with a single head.
| 234 | |
| 235 | |
| 236 | class AttentionBlock(nn.Module): |
| 237 | """ |
| 238 | Causal self-attention with a single head. |
| 239 | """ |
| 240 | |
| 241 | def __init__(self, dim): |
| 242 | super().__init__() |
| 243 | self.dim = dim |
| 244 | |
| 245 | # layers |
| 246 | self.norm = RMS_norm(dim) |
| 247 | self.to_qkv = nn.Conv2d(dim, dim * 3, 1) |
| 248 | self.proj = nn.Conv2d(dim, dim, 1) |
| 249 | |
| 250 | # zero out the last layer params |
| 251 | nn.init.zeros_(self.proj.weight) |
| 252 | |
| 253 | def forward(self, x): |
| 254 | identity = x |
| 255 | b, c, t, h, w = x.size() |
| 256 | x = rearrange(x, 'b c t h w -> (b t) c h w') |
| 257 | x = self.norm(x) |
| 258 | # compute query, key, value |
| 259 | q, k, v = self.to_qkv(x).reshape(b * t, 1, c * 3, -1).permute( |
| 260 | 0, 1, 3, 2).contiguous().chunk(3, dim=-1) |
| 261 | |
| 262 | # apply attention |
| 263 | x = F.scaled_dot_product_attention( |
| 264 | q, |
| 265 | k, |
| 266 | v, |
| 267 | #attn_mask=block_causal_mask(q, block_size=h * w) |
| 268 | ) |
| 269 | x = x.squeeze(1).permute(0, 2, 1).reshape(b * t, c, h, w) |
| 270 | |
| 271 | # output |
| 272 | x = self.proj(x) |
| 273 | x = rearrange(x, '(b t) c h w-> b c t h w', t=t) |
| 274 | return x + identity |
| 275 | |
| 276 | |
| 277 | class Encoder3d(nn.Module): |