| import torch |
| import torch.nn |
| import torch.nn.functional as F |
| from .multi_head_attention import MultiHeadAttention, AttentionMask |
| from typing import Optional, Callable, Dict |
| from dataclasses import dataclass |
| |
|
|
| ActivationFunction = Callable[[torch.Tensor], torch.Tensor] |
|
|
|
|
| class TransformerEncoderLayer(torch.nn.Module): |
| def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.1, activation: ActivationFunction = F.relu, |
| attention_dropout=0): |
| super(TransformerEncoderLayer, self).__init__() |
| self.self_attn = MultiHeadAttention(d_model, nhead, dropout=attention_dropout) |
| self.linear1 = torch.nn.Linear(d_model, dim_feedforward) |
| self.dropout = torch.nn.Dropout(dropout) |
| self.linear2 = torch.nn.Linear(dim_feedforward, d_model) |
|
|
| self.norm1 = torch.nn.LayerNorm(d_model) |
| self.norm2 = torch.nn.LayerNorm(d_model) |
| self.dropout1 = torch.nn.Dropout(dropout) |
| self.dropout2 = torch.nn.Dropout(dropout) |
|
|
| self.activation = activation |
| self.reset_parameters() |
|
|
| def forward(self, src: torch.Tensor, mask: Optional[AttentionMask] = None) -> torch.Tensor: |
| src2 = self.self_attn(src, src, mask) |
| src = src + self.dropout1(src2) |
| src = self.norm1(src) |
| src2 = self.linear2(self.dropout(self.activation(self.linear1(src)))) |
| src = src + self.dropout2(src2) |
| src = self.norm2(src) |
| return src |
|
|
| def reset_parameters(self): |
| torch.nn.init.xavier_uniform_(self.linear1.weight, gain=torch.nn.init.calculate_gain('relu') \ |
| if self.activation is F.relu else 1.0) |
| torch.nn.init.xavier_uniform_(self.linear2.weight) |
|
|
|
|
| class TransformerDecoderLayer(torch.nn.Module): |
| def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.1, activation: ActivationFunction = F.relu, |
| attention_dropout=0): |
| super(TransformerDecoderLayer, self).__init__() |
|
|
| self.self_attn = MultiHeadAttention(d_model, nhead, dropout=attention_dropout) |
| self.multihead_attn = MultiHeadAttention(d_model, nhead, dropout=attention_dropout) |
| |
| self.linear1 = torch.nn.Linear(d_model, dim_feedforward) |
| self.dropout = torch.nn.Dropout(dropout) |
| self.linear2 = torch.nn.Linear(dim_feedforward, d_model) |
|
|
| self.norm1 = torch.nn.LayerNorm(d_model) |
| self.norm2 = torch.nn.LayerNorm(d_model) |
| self.norm3 = torch.nn.LayerNorm(d_model) |
| self.dropout1 = torch.nn.Dropout(dropout) |
| self.dropout2 = torch.nn.Dropout(dropout) |
| self.dropout3 = torch.nn.Dropout(dropout) |
|
|
| self.activation = activation |
| self.reset_parameters() |
|
|
| def forward(self, tgt: torch.Tensor, memory: torch.Tensor, tgt_mask: Optional[torch.Tensor] = None, |
| memory_key_padding_mask: Optional[torch.Tensor] = None, |
| full_target: Optional[torch.Tensor] = None, pos_offset: int = 0) -> torch.Tensor: |
| |
| assert pos_offset == 0 or tgt_mask is None |
| tgt2 = self.self_attn(tgt, tgt if full_target is None else full_target, mask=AttentionMask(None, tgt_mask)) |
| tgt = tgt + self.dropout1(tgt2) |
| tgt = self.norm1(tgt) |
| tgt2 = self.multihead_attn(tgt, memory, mask=AttentionMask(memory_key_padding_mask, None)) |
| tgt = tgt + self.dropout2(tgt2) |
| tgt = self.norm2(tgt) |
| tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt)))) |
| tgt = tgt + self.dropout3(tgt2) |
| tgt = self.norm3(tgt) |
| return tgt |
|
|
| def reset_parameters(self): |
| torch.nn.init.xavier_uniform_(self.linear1.weight, gain=torch.nn.init.calculate_gain('relu') \ |
| if self.activation is F.relu else 1.0) |
| torch.nn.init.xavier_uniform_(self.linear2.weight) |
|
|
|
|
| class TransformerDecoderBase(torch.nn.Module): |
| @dataclass |
| class State: |
| step: int |
| state: Dict[int, torch.Tensor] |
|
|
| def __init__(self, d_model: int): |
| super().__init__() |
| self.d_model = d_model |
|
|
| def create_state(self, batch_size: int, max_length: int, device: torch.device) -> State: |
| return self.State(0, {i: torch.empty([batch_size, max_length, self.d_model], device=device) |
| for i in range(len(self.layers))}) |
|
|
| def one_step_forward(self, state: State, data: torch.Tensor, *args, **kwargs): |
| assert data.shape[1] == 1, f"For one-step forward should have one timesteps, but shape is {data.shape}" |
| assert state.step < state.state[0].shape[1] |
|
|
| for i, l in enumerate(self.layers): |
| state.state[i][:, state.step:state.step + 1] = data |
| data = l(data, *args, **kwargs, full_target=state.state[i][:, :state.step + 1], |
| pos_offset=state.step) |
|
|
| state.step += 1 |
| return data |
|
|
|
|
| class TransformerEncoder(torch.nn.Module): |
| def __init__(self, layer, n_layers: int, *args, **kwargs): |
| super().__init__() |
| self.layers = torch.nn.ModuleList([layer(*args, **kwargs) for _ in range(n_layers)]) |
|
|
| def forward(self, data: torch.Tensor, *args, **kwargs): |
| for l in self.layers: |
| data = l(data, *args, **kwargs) |
| return data |
|
|
|
|
| class TransformerDecoder(TransformerDecoderBase): |
| def __init__(self, layer, n_layers: int, d_model: int, *args, **kwargs): |
| super().__init__(d_model) |
| self.layers = torch.nn.ModuleList([layer(d_model, *args, **kwargs) for _ in range(n_layers)]) |
|
|
| def forward(self, data: torch.Tensor, *args, **kwargs): |
| for l in self.layers: |
| data = l(data, *args, **kwargs) |
| return data |
|
|
|
|
| def TransformerEncoderWithLayer(layer = TransformerEncoder): |
| return lambda *args, **kwargs: TransformerEncoder(layer, *args, **kwargs) |
|
|
|
|
| def TransformerDecoderWithLayer(layer = TransformerDecoder): |
| return lambda *args, **kwargs: TransformerDecoder(layer, *args, **kwargs) |
|
|
|
|
| class Transformer(torch.nn.Module): |
| def __init__(self, d_model: int = 512, nhead: int = 8, num_encoder_layers: int = 6, |
| num_decoder_layers: int = 6, dim_feedforward: int = 2048, dropout: float = 0.1, |
| activation: ActivationFunction = F.relu, encoder_layer=TransformerEncoderWithLayer(), |
| decoder_layer=TransformerDecoderWithLayer(), attention_dropout: float = 0): |
| super().__init__() |
|
|
| self.encoder = encoder_layer(num_encoder_layers, d_model, nhead, dim_feedforward, |
| dropout, activation, attention_dropout) |
| self.decoder = decoder_layer(num_decoder_layers, d_model, nhead, dim_feedforward, |
| dropout, activation, attention_dropout) |
|
|
| def forward(self, src: torch.Tensor, tgt: torch.Tensor, tgt_mask: Optional[torch.Tensor] = None, |
| src_mask: Optional[AttentionMask] = None): |
|
|
| memory = self.encoder(src, src_mask) |
| return self.decoder(tgt, memory, tgt_mask, src_mask.src_length_mask if src_mask is not None else None) |
|
|
| @staticmethod |
| def generate_square_subsequent_mask(sz: int, device: torch.device) -> torch.Tensor: |
| return torch.triu(torch.ones(sz, sz, dtype=torch.bool, device=device), diagonal=1) |
|
|
