| 315 | |
| 316 | |
| 317 | class SelfAttention1d(nn.Module): |
| 318 | def __init__(self, in_channels: int, n_head: int = 1, dropout_rate: float = 0.0): |
| 319 | super().__init__() |
| 320 | self.channels = in_channels |
| 321 | self.group_norm = nn.GroupNorm(1, num_channels=in_channels) |
| 322 | self.num_heads = n_head |
| 323 | |
| 324 | self.query = nn.Linear(self.channels, self.channels) |
| 325 | self.key = nn.Linear(self.channels, self.channels) |
| 326 | self.value = nn.Linear(self.channels, self.channels) |
| 327 | |
| 328 | self.proj_attn = nn.Linear(self.channels, self.channels, bias=True) |
| 329 | |
| 330 | self.dropout = nn.Dropout(dropout_rate, inplace=True) |
| 331 | |
| 332 | def transpose_for_scores(self, projection: torch.Tensor) -> torch.Tensor: |
| 333 | new_projection_shape = projection.size()[:-1] + (self.num_heads, -1) |
| 334 | # move heads to 2nd position (B, T, H * D) -> (B, T, H, D) -> (B, H, T, D) |
| 335 | new_projection = projection.view(new_projection_shape).permute(0, 2, 1, 3) |
| 336 | return new_projection |
| 337 | |
| 338 | def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: |
| 339 | residual = hidden_states |
| 340 | batch, channel_dim, seq = hidden_states.shape |
| 341 | |
| 342 | hidden_states = self.group_norm(hidden_states) |
| 343 | hidden_states = hidden_states.transpose(1, 2) |
| 344 | |
| 345 | query_proj = self.query(hidden_states) |
| 346 | key_proj = self.key(hidden_states) |
| 347 | value_proj = self.value(hidden_states) |
| 348 | |
| 349 | query_states = self.transpose_for_scores(query_proj) |
| 350 | key_states = self.transpose_for_scores(key_proj) |
| 351 | value_states = self.transpose_for_scores(value_proj) |
| 352 | |
| 353 | scale = 1 / math.sqrt(math.sqrt(key_states.shape[-1])) |
| 354 | |
| 355 | attention_scores = torch.matmul(query_states * scale, key_states.transpose(-1, -2) * scale) |
| 356 | attention_probs = torch.softmax(attention_scores, dim=-1) |
| 357 | |
| 358 | # compute attention output |
| 359 | hidden_states = torch.matmul(attention_probs, value_states) |
| 360 | |
| 361 | hidden_states = hidden_states.permute(0, 2, 1, 3).contiguous() |
| 362 | new_hidden_states_shape = hidden_states.size()[:-2] + (self.channels,) |
| 363 | hidden_states = hidden_states.view(new_hidden_states_shape) |
| 364 | |
| 365 | # compute next hidden_states |
| 366 | hidden_states = self.proj_attn(hidden_states) |
| 367 | hidden_states = hidden_states.transpose(1, 2) |
| 368 | hidden_states = self.dropout(hidden_states) |
| 369 | |
| 370 | output = hidden_states + residual |
| 371 | |
| 372 | return output |
| 373 | |
| 374 | |