Полная реализация FeedForward

- Основной класс FeedForward
- Тесты для всех функций
- Пример использования с визуализацией
- Документация с блок-схемой

example/feed_forward_example.py

@@ -0,0 +1,74 @@
+"""
+Пример использования FeedForward слоя из архитектуры Transformer
+
+Демонстрирует:
+1. Базовое применение
+2. Разницу между режимами train/eval
+3. Визуализацию изменений внутри сети
+"""
+
+import torch
+import matplotlib.pyplot as plt
+import os
+from simple_llm.transformer.feed_forward import FeedForward
+
+def plot_layer_outputs(outputs, titles, filename):
+    """Визуализация выходов разных слоев"""
+    plt.figure(figsize=(15, 5))
+    for i, (out, title) in enumerate(zip(outputs, titles)):
+        plt.subplot(1, len(outputs), i+1)
+        plt.imshow(out[0].detach().numpy(), cmap='viridis', aspect='auto')
+        plt.title(title)
+        plt.colorbar()
+    plt.tight_layout()
+    
+    # Создаем папку если нет
+    os.makedirs('example_output', exist_ok=True)
+    plt.savefig(f'example_output/{filename}')
+    plt.close()
+
+def main():
+    # Конфигурация
+    emb_size = 128
+    dropout = 0.1
+    
+    # Инициализация
+    ff = FeedForward(emb_size, dropout)
+    print(f"Архитектура сети:\n{ff.net}")
+
+    # Тестовые данные
+    x = torch.randn(1, 20, emb_size)  # [batch, seq_len, emb_size]
+    
+    # 1. Базовый forward pass
+    output = ff(x)
+    print(f"\nФорма входа: {x.shape} -> Форма выхода: {output.shape}")
+
+    # 2. Сравнение режимов train/eval
+    ff.train()
+    train_out = ff(x)
+    ff.eval()
+    eval_out = ff(x)
+    diff = torch.abs(train_out - eval_out).max().item()
+    print(f"\nМаксимальное расхождение (train vs eval): {diff:.6f}")
+
+    # 3. Визуализация преобразований
+    with torch.no_grad():
+        # Получаем выходы каждого слоя
+        layer1_out = ff.net[0](x)
+        relu_out = ff.net[1](layer1_out)
+        layer2_out = ff.net[2](relu_out)
+        
+        plot_layer_outputs(
+            outputs = [x, layer1_out, relu_out, layer2_out],
+            titles = [
+                'Входные данные', 
+                'После первого Linear', 
+                'После ReLU', 
+                'После второго Linear'
+            ],
+            filename = 'feed_forward_layers.png'
+        )
+
+if __name__ == "__main__":
+    main()
+    print("\nГотово! Результаты сохранены в example_output/feed_forward_layers.png")

simple_llm/transformer/feed_forward.py

@@ -0,0 +1,68 @@
+from torch import nn
+import torch
+
+class FeedForward(nn.Module):
+    """
+    Слой прямой связи (Feed Forward Network) для архитектуры трансформеров.
+    
+    Этот слой состоит из двух линейных преобразований с расширением внутренней размерности
+    в 4 раза и механизмом dropout для регуляризации. Между линейными слоями применяется
+    активация ReLU.
+
+    Алгоритм работы:
+    1. Входной тензор x (размерность: [batch_size, seq_len, emb_size])
+    2. Линейное преобразование: emb_size -> 4*emb_size
+    3. Активация ReLU
+    4. Линейное преобразование: 4*emb_size -> emb_size
+    5. Применение dropout
+    6. Возврат результата (размерность: [batch_size, seq_len, emb_size])
+
+    Предназначение:
+    - Добавляет нелинейность в архитектуру трансформера
+    - Обеспечивает взаимодействие между различными размерностями эмбеддингов
+    - Работает независимо для каждого токена в последовательности
+
+    Примеры использования:
+    
+    >>> # Инициализация слоя
+    >>> ff = FeedForward(emb_size=512, dropout=0.1)
+    >>>
+    >>> # Прямой проход
+    >>> x = torch.randn(32, 10, 512)  # [batch_size, seq_len, emb_size]
+    >>> output = ff(x)
+    >>> print(output.shape)  # torch.Size([32, 10, 512])
+    >>>
+    >>> # Работа с разными типами данных
+    >>> x_double = torch.randn(32, 10, 512, dtype=torch.float64)
+    >>> output_double = ff(x_double)
+    >>> print(output_double.dtype)  # torch.float64
+    """
+    def __init__(self, emb_size: int, dropout: float = 0.1):
+        """
+        Инициализация слоя Feed Forward Network.
+        
+        Args:
+            emb_size: Размерность входных эмбеддингов
+            dropout: Вероятность dropout для регуляризации (по умолчанию: 0.1)
+        """
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Linear(emb_size, 4 * emb_size),
+            nn.ReLU(),
+            nn.Linear(4 * emb_size, emb_size),
+            nn.Dropout(dropout)
+        )
+
+    def forward(self, x: torch.Tensor):
+        """
+        Прямой проход через слой Feed Forward Network.
+        
+        Args:
+            x: Входной тензор размерности [batch_size, seq_len, emb_size]
+            
+        Returns:
+            Тензор той же размерности, что и входной
+        """
+        # Приводим все параметры сети к типу входного тензора
+        self.net = self.net.to(x.dtype)
+        return self.net(x)

tests/test_feed_forward.py

@@ -0,0 +1,56 @@
+import torch
+import pytest
+from simple_llm.transformer.feed_forward import FeedForward
+
+class TestFeedForward:
+    @pytest.fixture
+    def ff_layer(self):
+        return FeedForward(emb_size=512)
+
+    def test_initialization(self, ff_layer):
+        assert isinstance(ff_layer.net, torch.nn.Sequential)
+        assert len(ff_layer.net) == 4
+        assert isinstance(ff_layer.net[0], torch.nn.Linear)
+        assert isinstance(ff_layer.net[1], torch.nn.ReLU)
+        assert isinstance(ff_layer.net[2], torch.nn.Linear)
+        assert isinstance(ff_layer.net[3], torch.nn.Dropout)
+        
+        assert ff_layer.net[0].in_features == 512
+        assert ff_layer.net[0].out_features == 2048
+        assert ff_layer.net[2].in_features == 2048
+        assert ff_layer.net[2].out_features == 512
+
+    def test_forward_pass_shape(self, ff_layer):
+        batch_size = 4
+        seq_len = 10
+        x = torch.randn(batch_size, seq_len, 512)
+        output = ff_layer(x)
+        
+        assert output.shape == (batch_size, seq_len, 512)
+
+    def test_dropout_training(self):
+        ff_layer = FeedForward(512, dropout=0.5)
+        ff_layer.train()
+        x = torch.randn(2, 5, 512)
+        output = ff_layer(x)
+        
+        # Проверяем, что dropout действительно работает в режиме обучения
+        layers = ff_layer.net
+        no_dropout = layers[2](layers[1](layers[0](x)))
+        assert not torch.allclose(output, no_dropout)
+
+    def test_dropout_eval(self):
+        ff_layer = FeedForward(512, dropout=0.5)
+        ff_layer.eval()
+        x = torch.randn(2, 5, 512)
+        output = ff_layer(x)
+        
+        # В eval режиме dropout не должен работать
+        layers = ff_layer.net
+        expected = layers[2](layers[1](layers[0](x)))
+        assert torch.allclose(output, expected)
+
+    def test_dtype_preservation(self, ff_layer):
+        x = torch.randn(2, 5, 512, dtype=torch.float64)
+        output = ff_layer(x)
+        assert output.dtype == torch.float64