Transformer

gistfile1.txt

import torch
import torch.nn as nn
import torch.nn.functional as F

class TransformerEncoderLayer(nn.Module):
    """
    Один слой энкодера трансформера (Post-LN).

    Аргументы:
        d_model (int): Размерность эмбеддингов (например, 512)
        nhead (int): Количество голов внимания (например, 8)
        dim_feedforward (int): Внутренняя размерность FFN (обычно 4*d_model)
        dropout (float): Вероятность дропаута (по умолчанию 0.1)
    """
    def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.1):
        super().__init__()

        # 1. Слой многоголового внимания
        self.self_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout)

        # 2. Двухслойная FFN-сеть
        self.linear1 = nn.Linear(d_model, dim_feedforward)
        self.linear2 = nn.Linear(dim_feedforward, d_model)

        # 3. Нормализация и дропауты
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.dropout = nn.Dropout(dropout)
        self.dropout1 = nn.Dropout(dropout)
        self.dropout2 = nn.Dropout(dropout)

        # 4. Активация для FFN (в оригинале ReLU)
        self.activation = F.relu

    def forward(self, src, src_mask=None, src_key_padding_mask=None):
        """
        Вход:
            src: (S, B, D) - входная последовательность (длина, батч, размерность)
            src_mask: (S, S) - маска внимания (опционально)
            src_key_padding_mask: (B, S) - маска батча (опционально)
        """
        # Шаг 1: Самовнимание с residual connection
        attn_output, _ = self.self_attn(
            src, src, src,
            attn_mask=src_mask,
            key_padding_mask=src_key_padding_mask
        )
        src = src + self.dropout1(attn_output)
        src = self.norm1(src)  # Post-LN: нормализация ПОСЛЕ residual

        # Шаг 2: FFN с residual connection
        ff_output = self.linear1(src)
        ff_output = self.activation(ff_output)
        ff_output = self.dropout(ff_output)
        ff_output = self.linear2(ff_output)
        src = src + self.dropout2(ff_output)
        src = self.norm2(src)

        return src

# Пример использования (без вывода данных)
if __name__ == "__main__":
    # Параметры модели
    d_model = 512    # Размерность эмбеддингов
    nhead = 8        # 8 голов внимания
    seq_len = 100    # Длина последовательности
    batch_size = 32  # Размер батча

    # Создаём слой энкодера
    encoder_layer = TransformerEncoderLayer(d_model, nhead)

    # Генерируем тестовые данные:
    # (S, B, D) = (длина_последовательности, батч, размерность)
    test_input = torch.rand(seq_len, batch_size, d_model)

    # Вызов слоя (без масок)
    output = encoder_layer(test_input)

    # Типичный пайплайн:
    # 1. Добавить позиционные энкодинги к test_input
    # 2. Пропустить через несколько слоёв энкодера
    # 3. Обработать выход модели