{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.12.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceType":"competition","sourceId":3004,"databundleVersionId":861823}],"dockerImageVersionId":31329,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import torch\nimport torch.nn as nn\nfrom torch.utils.data import DataLoader, Dataset, Subset, random_split\nfrom torchvision import datasets, transforms\nimport matplotlib.pyplot as plt\nfrom typing import Callable\nfrom tqdm import tqdm\n\ndef get_default_device() -> torch.device:\n    \"\"\"CUDA → MPS → CPU\"\"\"\n    if torch.cuda.is_available():\n        return torch.device('cuda')\n    if torch.backends.mps.is_available():\n        return torch.device('mps')\n    return torch.device('cpu')\n\n\ndef load_data(val_fraction: float = 0.1, seed: int = 42) -> tuple[Dataset, Dataset, Dataset]:\n    \"\"\"\n    Возвращает (train, val, test).\n    На train — аугментации, на val и test — чистый ToTensor.\n    Val отрезается от тренировочной части MNIST, test — это отдельная тестовая часть.\n    \"\"\"\n    normalize = transforms.Normalize((0.1307,), (0.3081,))\n    train_transform = transforms.Compose([\n        transforms.RandomAffine(degrees=10, translate=(0.1, 0.1), scale=(0.9, 1.1)),\n        transforms.ToTensor(),\n        normalize,\n    ])\n    plain_transform = transforms.Compose([transforms.ToTensor(), normalize])\n\n    # Два экземпляра тренировочного датасета: один с аугментациями, второй без.\n    # Берём из них одни и те же индексы — так val получает те же картинки, что были бы в train,\n    # но без аугментаций (метрика чище, не шумит от случайных поворотов).\n    train_data_aug = datasets.MNIST(root='./data', train=True,  download=True, transform=train_transform)\n    train_data_plain = datasets.MNIST(root='./data', train=True,  download=True, transform=plain_transform)\n    test_data = datasets.MNIST(root='./data', train=False, download=True, transform=plain_transform)\n\n    n = len(train_data_aug)\n    val_size = int(n * val_fraction)\n    train_size = n - val_size\n\n    generator = torch.Generator().manual_seed(seed)\n    # random_split работает с любым sized-объектом; используем range(n), чтобы получить детерминированные индексы.\n    train_split, val_split = random_split(range(n), [train_size, val_size], generator=generator)\n    train_indices = list(train_split)\n    val_indices = list(val_split)\n\n    train_subset = Subset(train_data_aug, train_indices)\n    val_subset = Subset(train_data_plain, val_indices)\n\n    return train_subset, val_subset, test_data\n\nclass TrainableMixin:\n    @property\n    def device(self) -> torch.device:\n        return next(self.parameters()).device\n\n    def _make_loader(self, data: Dataset, batch_size: int, shuffle: bool) -> DataLoader:\n        pin = self.device.type == 'cuda'\n        num_workers = 4 if pin else 0\n        return DataLoader(\n            data, batch_size=batch_size,\n            shuffle=shuffle, pin_memory=pin,\n            num_workers=num_workers,\n            persistent_workers=num_workers > 0,\n        )\n\n    @torch.no_grad()\n    def _eval_loader(self, loader: DataLoader, loss_fn: nn.Module) -> tuple[float, float]:\n        \"\"\"Считает (avg_loss, accuracy) на любом лоадере. Ничего не печатает.\"\"\"\n        self.eval()\n        device = self.device\n        total_loss = 0.0\n        correct = 0\n        total = 0\n        for images, labels in loader:\n            images, labels = images.to(device), labels.to(device)\n            logits = self(images)\n            loss = loss_fn(logits, labels)\n            total_loss += loss.item() * labels.size(0)\n            predictions = logits.argmax(dim=1)\n            correct += (predictions == labels).sum().item()\n            total += labels.size(0)\n        avg_loss = total_loss / total\n        accuracy = 100 * correct / total\n        return avg_loss, accuracy\n\n    def fit(\n        self,\n        train_data: Dataset,\n        valid_data: Dataset | None = None,\n        epochs: int = 25,\n        lr: float = 1e-3,\n        weight_decay: float = 1e-4,\n        batch_size: int = 64,\n        verbose: bool = True,\n    ) -> dict[str, list[float]]:\n        loss_fn = nn.CrossEntropyLoss(label_smoothing=0.05)\n        optimizer = torch.optim.AdamW(self.parameters(), lr=lr, weight_decay=weight_decay)\n        scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=epochs)\n        device = self.device\n\n        train_loader = self._make_loader(train_data, batch_size=batch_size, shuffle=True)\n        valid_loader = self._make_loader(valid_data, batch_size=batch_size, shuffle=False) if valid_data is not None else None\n\n        history: dict[str, list[float]] = {'train_loss': [], 'val_loss': [], 'val_acc': []}\n\n        for epoch in range(epochs):\n            self.train()\n            running_loss = 0.0\n            seen = 0\n            pbar = tqdm(\n                train_loader,\n                desc=f\"epoch {epoch+1}/{epochs}\",\n                disable=not verbose,\n                leave=True,\n            )\n            for images, labels in pbar:\n                images, labels = images.to(device), labels.to(device)\n                predictions = self(images)\n                loss = loss_fn(predictions, labels)\n\n                optimizer.zero_grad()\n                loss.backward()\n                optimizer.step()\n\n                # Усредняем по числу примеров, а не по числу батчей — корректнее,\n                # если последний батч меньше остальных.\n                running_loss += loss.item() * labels.size(0)\n                seen += labels.size(0)\n                pbar.set_postfix(loss=f\"{loss.item():.4f}\")\n\n            scheduler.step()\n            train_loss = running_loss / seen\n            history['train_loss'].append(train_loss)\n\n            if valid_loader is not None:\n                val_loss, val_acc = self._eval_loader(valid_loader, loss_fn)\n                history['val_loss'].append(val_loss)\n                history['val_acc'].append(val_acc)\n                if verbose:\n                    pbar.write(\n                        f\"Эпоха {epoch+1}/{epochs} — \"\n                        f\"train_loss: {train_loss:.4f} | \"\n                        f\"val_loss: {val_loss:.4f} | \"\n                        f\"val_acc: {val_acc:.2f}%\"\n                    )\n            else:\n                if verbose:\n                    pbar.write(f\"Эпоха {epoch+1}/{epochs} — train_loss: {train_loss:.4f}\")\n\n        self.history = history\n        return history\n\n    @torch.no_grad()\n    def evaluate(self, test_data: Dataset, batch_size: int = 64) -> float:\n        loss_fn = nn.CrossEntropyLoss()\n        test_loader = self._make_loader(test_data, batch_size=batch_size, shuffle=False)\n        _, accuracy = self._eval_loader(test_loader, loss_fn)\n        print(f\"Точность на тесте: {accuracy:.2f}%\")\n        return accuracy\n\n\n    def plot(self, save_path: str | None = None) -> None:\n        \"\"\"\n        Рисует историю обучения. Без аргумента — показывает на экран.\n        С путём — сохраняет в файл.\n        \"\"\"\n        history = getattr(self, 'history', None)\n        if not history or not history.get('train_loss'):\n            print(\"Нечего рисовать: модель ещё не обучалась (или history не сохранён).\")\n            return\n\n        epochs = range(1, len(history['train_loss']) + 1)\n        has_val = len(history.get('val_loss', [])) > 0\n\n        fig, axes = plt.subplots(1, 2, figsize=(12, 4))\n\n        ax = axes[0]\n        ax.plot(epochs, history['train_loss'], label='train_loss', marker='o', markersize=3)\n        if has_val:\n            ax.plot(epochs, history['val_loss'], label='val_loss', marker='o', markersize=3)\n        ax.set_xlabel('эпоха')\n        ax.set_ylabel('loss')\n        ax.set_title('Loss')\n        ax.grid(alpha=0.3)\n        ax.legend()\n\n        ax = axes[1]\n        if has_val:\n            ax.plot(epochs, history['val_acc'], label='val_acc', color='tab:green', marker='o', markersize=3)\n            best_epoch = int(torch.tensor(history['val_acc']).argmax().item()) + 1\n            best_acc = history['val_acc'][best_epoch - 1]\n            ax.axvline(best_epoch, color='gray', linestyle='--', alpha=0.5)\n            ax.set_title(f'Val accuracy (лучшая: {best_acc:.2f}% на эпохе {best_epoch})')\n        ax.set_xlabel('эпоха')\n        ax.set_ylabel('accuracy, %')\n        ax.grid(alpha=0.3)\n        ax.legend()\n\n        plt.tight_layout()\n        if save_path is None:\n            plt.show()\n        else:\n            plt.savefig(save_path, dpi=120)\n            plt.close()\n            print(f\"График сохранён в {save_path}\")\n\n\n\n    @torch.no_grad()\n    def predict(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:\n        self.eval()\n        # Приводим к [B, C, H, W]\n        if x.dim() == 2:           # [H, W]\n            x = x.unsqueeze(0).unsqueeze(0)\n        elif x.dim() == 3:         # [C, H, W]\n            x = x.unsqueeze(0)\n        x = x.to(self.device)\n        logits = self(x)\n        probabilities = torch.softmax(logits, dim=1)\n        digits = logits.argmax(dim=1)\n        confidences = probabilities.max(dim=1).values\n        return digits.cpu(), confidences.cpu()\n\n    def save_model(self, path: str = 'model.pth') -> None:\n        torch.save({\n            'params': self.params,\n            'class_name': type(self).__name__,\n            'state_dict': self.state_dict(),\n            'history': getattr(self, 'history', None),\n        }, path)\n        print(f\"Модель сохранена в {path}\")\n\n    @classmethod\n    def load_model(cls, path: str = 'model.pth', device: torch.device | str | None = None) -> 'TrainableMixin':\n        if device is None:\n            device = get_default_device()\n        device = torch.device(device)\n        load = torch.load(path, map_location=device, weights_only=True)\n        model = cls(**load['params'])\n        model.load_state_dict(load['state_dict'])\n        model.to(device)\n        model.eval()\n        if load.get('history') is not None:\n            model.history = load['history']\n        print(f\"Модель загружена из {path} на {device}\")\n        return model\n\nclass ConvBlock(nn.Module):\n    def __init__(\n        self,\n        in_channels: int,\n        out_channels: int,\n        kernel_size: int = 3,\n        padding: int = 1,\n        stride: int = 1,\n        activation: Callable[[], nn.Module] = nn.ReLU,\n        bn: bool = True,\n        dropout: float = 0.0,\n    ) -> None:\n        super().__init__()\n        self.conv = nn.Conv2d(\n            in_channels, out_channels,\n            kernel_size=kernel_size, padding=padding, bias=not bn,\n            stride=stride,\n        )\n        self.bn: nn.Module = nn.BatchNorm2d(out_channels) if bn else nn.Identity()\n        self.activation = activation()\n        self.dropout: nn.Module = nn.Dropout2d(dropout) if dropout > 0.0 else nn.Identity()\n\n    def forward(self, x: torch.Tensor) -> torch.Tensor:\n        return (self.dropout(self.activation(self.bn(self.conv(x)))))\n\n\nclass Perceptron(nn.Module):\n    def __init__(\n        self,\n        in_features: int,\n        out_features: int,\n        activation: Callable[[], nn.Module] = nn.ReLU,\n        dropout: float = 0.0,\n    ) -> None:\n        super().__init__()\n        self.linear = nn.Linear(in_features, out_features)\n        self.activation = activation()\n        self.dropout: nn.Module = nn.Dropout(dropout) if dropout > 0.0 else nn.Identity()\n\n    def forward(self, x: torch.Tensor) -> torch.Tensor:\n        return self.dropout(self.activation(self.linear(x)))\n\n\nclass MLPClassifier(nn.Module):\n    def __init__(\n        self,\n        in_features: int,\n        fc_sizes: tuple[int, ...],\n        activation: Callable[[], nn.Module] = nn.ReLU,\n        dropout: float = 0.0,\n    ) -> None:\n        super().__init__()\n        sizes = (in_features,) + fc_sizes\n        blocks = [\n            Perceptron(sizes[i], sizes[i + 1], activation=activation, dropout=dropout)\n            for i in range(len(sizes) - 2)\n        ]\n        self.net = nn.Sequential(*blocks, nn.Linear(sizes[-2], sizes[-1]))\n\n    def forward(self, x: torch.Tensor) -> torch.Tensor:\n        return self.net(x)\n\nACTIVATIONS = {'relu': nn.ReLU, 'silu': nn.SiLU, 'gelu': nn.GELU, 'leaky_relu': nn.LeakyReLU}\n\nclass CNN(nn.Module, TrainableMixin):\n    def __init__(\n        self,\n        in_channels: int = 1,\n        conv_channels: tuple[int, ...] = (32, 64, 128),\n        fc_sizes: tuple[int, ...] = (512, 10),\n        activation: str = 'relu',\n        dropout: float = 0.0,\n        input_size: int = 28,\n    ) -> None:\n        super().__init__()\n        activation_cls = ACTIVATIONS[activation]\n        self.params = {\n            'in_channels': in_channels,\n            'conv_channels': conv_channels,\n            'fc_sizes': fc_sizes,\n            'activation': activation,\n            'dropout': dropout,\n            'input_size': input_size,\n        }\n\n        n_blocks = len(conv_channels)\n        strides = [\n            2 if (i == n_blocks - 1) or (conv_channels[i] != conv_channels[i + 1]) else 1\n            for i in range(n_blocks)\n        ]\n\n        channels = (in_channels,) + conv_channels\n        self.conv_blocks = nn.ModuleList([\n            ConvBlock(\n                channels[i],\n                channels[i + 1],\n                stride=strides[i],\n                activation=activation_cls,\n                dropout=dropout if i < len(conv_channels) - 1 else 0.0,\n                )\n            for i in range(len(conv_channels))\n        ])\n\n        spatial = input_size\n        for s in strides:\n            if s == 2:\n                spatial = (spatial + 1) // 2\n\n        self.flatten = nn.Flatten()\n\n        self.classifier = MLPClassifier(\n            in_features=conv_channels[-1] * spatial * spatial,\n            fc_sizes=fc_sizes,\n            activation=activation_cls,\n            dropout=dropout,\n        )\n\n    def forward(self, x: torch.Tensor) -> torch.Tensor:\n        for block in self.conv_blocks:\n            x = block(x)\n        return self.classifier(self.flatten(x))\n\ndef show_prediction(model, image: torch.Tensor, true_label: int, save_path: str, name: str) -> None:\n    \"\"\"Делает предикт, печатает результат и сохраняет картинку.\"\"\"\n    digits, confidences = model.predict(image)\n    pred = digits[0].item()\n    conf = confidences[0].item()\n    correct = \"✓\" if pred == true_label else \"✗\"\n    print(f\"[{name}] истина: {true_label}, предсказание: {pred} ({conf:.1%}) {correct}\")\n\n    plt.imshow(image.squeeze(), cmap='gray')\n    plt.title(f\"{name}: {pred} | истина: {true_label}\")\n    plt.axis('off')\n    plt.savefig(save_path)\n    plt.close()\n\n\ndef main():\n    device = get_default_device()\n    print(f\"Используем device: {device}\")\n\n    # --- Данные ---\n\n    train_data, valid_data, test_data = load_data()\n\n    # --- Модель ---\n    cnn = CNN(in_channels=1, conv_channels=(32, 32, 64, 64, 128, 128), fc_sizes=(256, 10), activation=\"silu\", dropout=0.1).to(device)\n\n    # --- Обучение ---\n    print(\"\\n=== Обучаем CNN ===\")\n    cnn.fit(train_data, valid_data, epochs=50, lr=1e-3)\n    cnn.plot()\n    cnn.evaluate(test_data)\n\n    # --- Сравним обе модели на одной картинке ---\n    print(\"\\n=== Предикт на test_data[1] ===\")\n    image, true_label = test_data[1]\n    show_prediction(cnn, image, true_label, 'pred_cnn.png', 'CNN')\n\n    # --- Сохранение ---\n    cnn.save_model('mnist_cnn.pth')\n\n    # --- Загрузка и проверка, что точность совпадает ---\n    print(\"\\n=== Проверка загрузки ===\")\n    loaded_cnn = CNN.load_model('mnist_cnn.pth', device=device)\n    loaded_cnn.evaluate(test_data)\n\nif __name__ == \"__main__\":\n    main()","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2026-05-13T08:05:17.777069Z","iopub.execute_input":"2026-05-13T08:05:17.777859Z","iopub.status.idle":"2026-05-13T08:14:57.476452Z","shell.execute_reply.started":"2026-05-13T08:05:17.777812Z","shell.execute_reply":"2026-05-13T08:14:57.475534Z"}},"outputs":[],"execution_count":null}]}