{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "Implementing Batch Normalization using TensorFlow"
      ],
      "metadata": {
        "id": "bzdp0lSSbpWm"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Import the Required Library\n",
        "import tensorflow as tf\n",
        "\n",
        "# Load and Preprocess the Dataset\n",
        "(x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data()\n",
        "\n",
        "x_train = x_train.astype(\"float32\") / 255.0\n",
        "x_test = x_test.astype(\"float32\") / 255.0\n",
        "\n",
        "x_train = x_train.reshape(-1, 784)\n",
        "x_test = x_test.reshape(-1, 784)\n",
        "\n",
        "x_train = x_train[:5000]\n",
        "y_train = y_train[:5000]\n",
        "\n",
        "# Create the Neural Network\n",
        "model = tf.keras.Sequential([\n",
        "    tf.keras.Input(shape=(784,)),\n",
        "    tf.keras.layers.Dense(64),\n",
        "    tf.keras.layers.BatchNormalization(),\n",
        "    tf.keras.layers.Activation(\"relu\"),\n",
        "    tf.keras.layers.Dense(10, activation=\"softmax\")\n",
        "])\n",
        "\n",
        "\n",
        "# Compile the Model\n",
        "model.compile(\n",
        "    optimizer=\"adam\",\n",
        "    loss=\"sparse_categorical_crossentropy\",\n",
        "    metrics=[\"accuracy\"]\n",
        ")\n",
        "\n",
        "# Train the Model\n",
        "model.fit(\n",
        "    x_train,\n",
        "    y_train,\n",
        "    epochs=2,\n",
        "    batch_size=32\n",
        ")"
      ],
      "metadata": {
        "id": "6l2P8jc4bqHl"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Implementing Batch Normalization using PyTorch"
      ],
      "metadata": {
        "id": "hc2dhPxXeHNm"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Import the Required Libraries\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.optim as optim\n",
        "from torchvision import datasets, transforms\n",
        "from torch.utils.data import DataLoader, Subset\n",
        "\n",
        "# Load and Prepare the Dataset\n",
        "transform = transforms.ToTensor()\n",
        "\n",
        "train_dataset = datasets.MNIST(\n",
        "    root=\"./data\",\n",
        "    train=True,\n",
        "    download=True,\n",
        "    transform=transform\n",
        ")\n",
        "\n",
        "\n",
        "train_dataset = Subset(train_dataset, range(5000))\n",
        "\n",
        "\n",
        "train_loader = DataLoader(\n",
        "    train_dataset,\n",
        "    batch_size=32,\n",
        "    shuffle=True\n",
        ")\n",
        "\n",
        "# Define the Neural Network with Batch Normalization\n",
        "class Model(nn.Module):\n",
        "    def __init__(self):\n",
        "        super().__init__()\n",
        "\n",
        "        self.flatten = nn.Flatten()\n",
        "        self.fc1 = nn.Linear(784, 64)\n",
        "        self.bn = nn.BatchNorm1d(64)\n",
        "        self.relu = nn.ReLU()\n",
        "        self.fc2 = nn.Linear(64, 10)\n",
        "\n",
        "    def forward(self, x):\n",
        "        x = self.flatten(x)\n",
        "        x = self.fc1(x)\n",
        "        x = self.bn(x)\n",
        "        x = self.relu(x)\n",
        "        x = self.fc2(x)\n",
        "        return x\n",
        "\n",
        "# Create the Model and Define the Loss Function and Optimizer\n",
        "model = Model()\n",
        "\n",
        "criterion = nn.CrossEntropyLoss()\n",
        "\n",
        "optimizer = optim.Adam(\n",
        "    model.parameters(),\n",
        "    lr=0.001\n",
        ")\n",
        "\n",
        "# Train the Model\n",
        "for epoch in range(2):\n",
        "    running_loss = 0.0\n",
        "\n",
        "    for images, labels in train_loader:\n",
        "\n",
        "        optimizer.zero_grad()\n",
        "\n",
        "        outputs = model(images)\n",
        "\n",
        "        loss = criterion(outputs, labels)\n",
        "\n",
        "        loss.backward()\n",
        "\n",
        "        optimizer.step()\n",
        "\n",
        "        running_loss += loss.item()\n",
        "\n",
        "    print(f\"Epoch {epoch+1}, Loss: {running_loss / len(train_loader):.4f}\")"
      ],
      "metadata": {
        "id": "d4iTiQWpeSiO"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}