diff --git a/example/feed_forward_example.py b/example/feed_forward_example.py new file mode 100644 index 0000000..4262b2c --- /dev/null +++ b/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") diff --git a/example_output/feed_forward_layers.png b/example_output/feed_forward_layers.png new file mode 100644 index 0000000..186b168 Binary files /dev/null and b/example_output/feed_forward_layers.png differ diff --git a/example_output/multi_head_attention.png b/example_output/multi_head_attention.png new file mode 100644 index 0000000..104b2e7 Binary files /dev/null and b/example_output/multi_head_attention.png differ diff --git a/simple_llm/transformer/feed_forward.py b/simple_llm/transformer/feed_forward.py new file mode 100644 index 0000000..81e7558 --- /dev/null +++ b/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) \ No newline at end of file diff --git a/tests/test_feed_forward.py b/tests/test_feed_forward.py new file mode 100644 index 0000000..ea0fb94 --- /dev/null +++ b/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