BiFormer.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU",
    "gpuClass": "standard"
  },
  "cells": [
    {
      "cell_type": "code",
      "source": [
        "!pip install einops"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "9w8EigkTR7sJ",
        "outputId": "6dba729d-f860-46a1-f416-1d344924cfed"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/\n",
            "Collecting einops\n",
            "  Downloading einops-0.6.1-py3-none-any.whl (42 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m42.2/42.2 kB\u001b[0m \u001b[31m3.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hInstalling collected packages: einops\n",
            "Successfully installed einops-0.6.1\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import tensorflow as tf\n",
        "from tensorflow.keras import Model\n",
        "from tensorflow.keras.layers import Layer\n",
        "from tensorflow.keras import Sequential\n",
        "import tensorflow.keras.layers as nn\n",
        "\n",
        "from tensorflow import einsum\n",
        "from einops import rearrange, repeat\n",
        "from einops.layers.tensorflow import Rearrange\n",
        "import numpy as np"
      ],
      "metadata": {
        "id": "uVXTk49tR4li"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def pair(t):\n",
        "    return t if isinstance(t, tuple) else (t, t)\n",
        "\n",
        "class PreNorm(Layer):\n",
        "    def __init__(self, fn):\n",
        "        super(PreNorm, self).__init__()\n",
        "\n",
        "        self.norm = nn.LayerNormalization()\n",
        "        self.fn = fn\n",
        "\n",
        "    def call(self, x, training=True):\n",
        "        return self.fn(self.norm(x), training=training)\n",
        "\n",
        "class Attentionlepe(tf.keras.layers.Layer):\n",
        "    def __init__(self, d_model, num_heads, dropout_rate):\n",
        "        super(Attentionlepe, self).__init__()\n",
        "        \n",
        "        self.num_heads = num_heads\n",
        "        self.d_model = d_model\n",
        "        self.dropout_rate = dropout_rate\n",
        "        \n",
        "        self.local_att_layer = tf.keras.layers.MultiHeadAttention(num_heads=self.num_heads, key_dim=self.d_model, \n",
        "                                                                  dropout=self.dropout_rate)\n",
        "        self.att_drop = tf.keras.layers.Dropout(self.dropout_rate)\n",
        "        self.local_att_norm = tf.keras.layers.LayerNormalization(epsilon=1e-6)\n",
        "\n",
        "    def call(self, inputs, mask=None):\n",
        "        bsz, qlen = tf.shape(inputs)[0], tf.shape(inputs)[1]\n",
        "        \n",
        "        pos_emb = self.add_positional_embeddings(inputs)\n",
        "        \n",
        "        pos_emb_reshaped = tf.reshape(pos_emb, [bsz, qlen, self.num_heads, self.d_model // self.num_heads])\n",
        "        pos_emb_reshaped = tf.transpose(pos_emb_reshaped, [0, 2, 1, 3])\n",
        "        pos_emb_reshaped = tf.reshape(pos_emb_reshaped, [bsz * self.num_heads, qlen, self.d_model // self.num_heads])\n",
        "        \n",
        "        out = self.local_att_layer(pos_emb_reshaped, pos_emb_reshaped, attention_mask=None, return_attention_scores=False)\n",
        "        \n",
        "        out = tf.reshape(out, [bsz, self.num_heads, qlen, self.d_model // self.num_heads])\n",
        "        out = tf.transpose(out, [0, 2, 1, 3])\n",
        "        out = tf.reshape(out, [bsz, qlen, self.d_model])\n",
        "        \n",
        "        out = self.att_drop(out)\n",
        "        out = self.local_att_norm(inputs + out)\n",
        "        \n",
        "        return out\n",
        "    \n",
        "    \n",
        "\n",
        "    def add_positional_embeddings(self, inputs):\n",
        "        bsz, seq_len = tf.shape(inputs)[0], tf.shape(inputs)[1]\n",
        "        pos_emb = self.get_absolute_position_embeddings(seq_len, self.d_model)\n",
        "        \n",
        "        pos_emb_reshaped = tf.reshape(pos_emb, [1, seq_len, self.d_model])\n",
        "        pos_emb_reshaped = tf.broadcast_to(pos_emb_reshaped, [bsz, seq_len, self.d_model])\n",
        "        \n",
        "        return pos_emb_reshaped\n",
        "\n",
        "    def get_absolute_position_embeddings(self, seq_len, hidden_size):\n",
        "        freqs = tf.range(0, hidden_size, 2, dtype=tf.float32)\n",
        "        freqs = 1 / (10000 ** (freqs / hidden_size))\n",
        "        pos = tf.range(0, seq_len, dtype=tf.float32)\n",
        "        pos = tf.expand_dims(pos, axis=1)\n",
        "        pos = tf.broadcast_to(pos, [seq_len, hidden_size // 2])\n",
        "        pos = tf.cast(pos, tf.float32)\n",
        "        sin_emb = tf.sin(pos * freqs)\n",
        "        cos_emb = tf.cos(pos * freqs)\n",
        "        pos_emb = tf.stack([sin_emb, cos_emb], axis=-1)\n",
        "        pos_emb = tf.reshape(pos_emb, [seq_len, hidden_size])\n",
        "        \n",
        "        return pos_emb\n",
        "\n",
        "\n",
        "class BilevelRoutingAttention(tf.keras.layers.Layer):\n",
        "    def __init__(self, d_model, num_iterations):\n",
        "        super(BilevelRoutingAttention, self).__init__()\n",
        "\n",
        "        self.d_model = d_model\n",
        "        self.num_iterations = num_iterations\n",
        "\n",
        "        self.w1 = self.add_weight(name='w1', shape=(d_model, d_model), initializer='glorot_uniform', trainable=True)\n",
        "        self.w2 = self.add_weight(name='w2', shape=(d_model, d_model), initializer='glorot_uniform', trainable=True)\n",
        "        self.v = self.add_weight(name='v', shape=(d_model, 1), initializer='glorot_uniform', trainable=True)\n",
        "\n",
        "    def call(self, inputs, mask=None):\n",
        "       \n",
        "        u = tf.matmul(inputs, self.w1)  # (batch_size, seq_len, d_model)\n",
        "        sequence_weights = tf.nn.softmax(tf.squeeze(tf.matmul(tf.nn.tanh(u), self.v), axis=-1))  # (batch_size, seq_len)\n",
        "\n",
        "        sequence_output = tf.reduce_sum(tf.expand_dims(sequence_weights, axis=-1) * inputs, axis=1)  # (batch_size, d_model)\n",
        "\n",
        "        if mask is not None:\n",
        "            mask = tf.expand_dims(mask, axis=-1)\n",
        "            inputs *= mask\n",
        "            u = tf.matmul(inputs, self.w2)  # (batch_size, seq_len, d_model)\n",
        "            token_weights = tf.nn.softmax(tf.squeeze(tf.matmul(tf.nn.tanh(u), self.v), axis=-1) + (1 - mask) * -1e9)\n",
        "        else:\n",
        "            u = tf.matmul(inputs, self.w2)  # (batch_size, seq_len, d_model)\n",
        "            token_weights = tf.nn.softmax(tf.squeeze(tf.matmul(tf.nn.tanh(u), self.v), axis=-1))\n",
        "\n",
        "        token_output = tf.reduce_sum(tf.expand_dims(token_weights, axis=-1) * inputs, axis=1)  # (batch_size, d_model)\n",
        "\n",
        "        for i in range(self.num_iterations):\n",
        "            token_output = token_output + sequence_output\n",
        "            u = tf.matmul(tf.expand_dims(token_output, axis=1), self.w2)  # (batch_size, 1, d_model)\n",
        "            token_weights = tf.nn.softmax(tf.squeeze(tf.matmul(tf.nn.tanh(u), self.v), axis=-1) + (1 - mask) * -1e9)\n",
        "            token_output = tf.reduce_sum(tf.expand_dims(token_weights, axis=-1) * inputs, axis=1)  # (batch_size, d_model)\n",
        "\n",
        "        return token_output\n",
        "\n",
        "class MLP(Layer):\n",
        "    def __init__(self, dim, hidden_dim, dropout=0.0):\n",
        "        super(MLP, self).__init__()\n",
        "\n",
        "        def GELU():\n",
        "            def gelu(x, approximate=False):\n",
        "                if approximate:\n",
        "                    coeff = tf.cast(0.044715, x.dtype)\n",
        "                    return 0.5 * x * (1.0 + tf.tanh(0.7978845608028654 * (x + coeff * tf.pow(x, 3))))\n",
        "                else:\n",
        "                    return 0.5 * x * (1.0 + tf.math.erf(x / tf.cast(1.4142135623730951, x.dtype)))\n",
        "\n",
        "            return nn.Activation(gelu)\n",
        "\n",
        "        self.net = Sequential([\n",
        "            nn.Dense(units=hidden_dim),\n",
        "            GELU(),\n",
        "            nn.Dropout(rate=dropout),\n",
        "            nn.Dense(units=dim),\n",
        "            nn.Dropout(rate=dropout)\n",
        "        ])\n",
        "\n",
        "    def call(self, x, training=True):\n",
        "        return self.net(x, training=training)\n",
        "\n",
        "class Attention(Layer):\n",
        "    def __init__(self, dim, heads=8, dim_head=64, dropout=0.0):\n",
        "        super(Attention, self).__init__()\n",
        "        inner_dim = dim_head * heads\n",
        "        project_out = not (heads == 1 and dim_head == dim)\n",
        "\n",
        "        self.heads = heads\n",
        "        self.scale = dim_head ** -0.5\n",
        "\n",
        "        self.attend = nn.Softmax()\n",
        "        self.to_qkv = nn.Dense(units=inner_dim * 3, use_bias=False)\n",
        "\n",
        "        if project_out:\n",
        "            self.to_out = [\n",
        "                nn.Dense(units=dim),\n",
        "                nn.Dropout(rate=dropout)\n",
        "            ]\n",
        "        else:\n",
        "            self.to_out = []\n",
        "\n",
        "        self.to_out = Sequential(self.to_out)\n",
        "\n",
        "    def call(self, x, training=True):\n",
        "        qkv = self.to_qkv(x)\n",
        "        qkv = tf.split(qkv, num_or_size_splits=3, axis=-1)\n",
        "        q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h=self.heads), qkv)\n",
        "\n",
        "        # dots = tf.matmul(q, tf.transpose(k, perm=[0, 1, 3, 2])) * self.scale\n",
        "        dots = einsum('b h i d, b h j d -> b h i j', q, k) * self.scale\n",
        "        attn = self.attend(dots)\n",
        "\n",
        "        # x = tf.matmul(attn, v)\n",
        "        x = einsum('b h i j, b h j d -> b h i d', attn, v)\n",
        "        x = rearrange(x, 'b h n d -> b n (h d)')\n",
        "        x = self.to_out(x, training=training)\n",
        "\n",
        "        return x\n",
        "\n",
        "class block(Layer):\n",
        "    def __init__(self, dim, depth, heads, dim_head, mlp_dim,topk=-1,layer_scale_val=None, dropout=0.0):\n",
        "        super(block, self).__init__()\n",
        "        self.layer_scale_val = layer_scale_val\n",
        "        self.layers = []\n",
        "        if topk>0:\n",
        "          att = BilevelRoutingAttention(dim,dim_head)\n",
        "        if topk== -1:\n",
        "          att = Attention(dim, heads=heads, dim_head=dim_head, dropout=dropout)\n",
        "        if topk==-2:\n",
        "          att = Attentionlepe(dim,dim_head,dropout)\n",
        "\n",
        "        for _ in range(depth):\n",
        "            self.layers.append([\n",
        "                PreNorm(att),\n",
        "                PreNorm(MLP(dim, mlp_dim, dropout=dropout))\n",
        "            ])\n",
        "\n",
        "    def call(self, x, training=True):\n",
        "        for attn, mlp in self.layers:\n",
        "          if self.layer_scale_val:\n",
        "            x = attn(x, training=training)*self.layer_scale_val + x\n",
        "            x = mlp(x, training=training)*self.layer_scale_val + x\n",
        "          else:\n",
        "            x = attn(x, training=training) + x\n",
        "            x = mlp(x, training=training) + x\n",
        "\n",
        "        return x\n",
        "        \n",
        "class BiFormer(Model):\n",
        "    def __init__(self, image_size, embed_size,patch_size, num_classes, dim, depth, heads, mlp_dim,\n",
        "                 pool='cls', dim_head=64, dropout=0.0, emb_dropout=0.0):\n",
        "       \n",
        "        super(BiFormer, self).__init__()\n",
        "\n",
        "        image_height, image_width = pair(image_size)\n",
        "        patch_height, patch_width = pair(patch_size)\n",
        "\n",
        "        assert image_height % patch_height == 0 and image_width % patch_width == 0, 'Image dimensions must be divisible by the patch size.'\n",
        "\n",
        "        num_patches = (image_height // patch_height) * (image_width // patch_width)\n",
        "        assert pool in {'cls', 'mean'}, 'pool type must be either cls (cls token) or mean (mean pooling)'\n",
        "\n",
        "        self.patch_embedding = Sequential([\n",
        "            Rearrange('b (h p1) (w p2) c -> b (h w) (p1 p2 c)', p1=patch_height, p2=patch_width),\n",
        "            nn.Dense(units=dim)\n",
        "        ], name='patch_embedding')\n",
        "        self.downsample = Sequential()\n",
        "        for i in embed_size:\n",
        "          self.downsample.add(nn.Conv2D(i, 3, activation='relu', padding=\"same\"))\n",
        "          self.downsample.add(nn.BatchNormalization())\n",
        "        self.pos_embedding = tf.Variable(initial_value=tf.random.normal([1, num_patches + 1, dim]))\n",
        "        self.cls_token = tf.Variable(initial_value=tf.random.normal([1, 1, dim]))\n",
        "        self.dropout = nn.Dropout(rate=emb_dropout)\n",
        "\n",
        "        self.block = block(dim, depth, heads, dim_head, mlp_dim,layer_scale_val=0.4,dropout= dropout)\n",
        "\n",
        "        self.pool = pool\n",
        "\n",
        "        self.mlp_head = Sequential([\n",
        "            nn.LayerNormalization(),\n",
        "            nn.Dense(units=num_classes)\n",
        "        ], name='mlp_head')\n",
        "\n",
        "    def call(self, img, training=True, **kwargs):\n",
        "        x = self.downsample(img)\n",
        "        x = self.patch_embedding(x)\n",
        "        b, n, d = x.shape\n",
        "\n",
        "        cls_tokens = repeat(self.cls_token, '() n d -> b n d', b=b)\n",
        "        x = tf.concat([cls_tokens, x], axis=1)\n",
        "        x += self.pos_embedding[:, :(n + 1)]\n",
        "        x = self.dropout(x, training=training)\n",
        "\n",
        "        x = self.block(x, training=training)\n",
        "\n",
        "        if self.pool == 'mean':\n",
        "            y = tf.reduce_mean(x, axis=1)\n",
        "        else:\n",
        "            y = x[:, 0]\n",
        "\n",
        "        y = self.mlp_head(y)\n",
        "        return y"
      ],
      "metadata": {
        "id": "IAh7kuPbU1MU"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "v = BiFormer(\n",
        "    image_size = 32,\n",
        "    embed_size=[2,4,6,8],\n",
        "    patch_size = 4,\n",
        "    num_classes = 10,\n",
        "    dim = 128,\n",
        "    depth = 6,\n",
        "    heads = 16,\n",
        "    mlp_dim = 256,\n",
        "    dropout = 0.1,\n",
        "    emb_dropout = 0.1\n",
        ")"
      ],
      "metadata": {
        "id": "y6SvGebpMlxc"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "img = tf.random.normal([2, 32, 32, 3])\n",
        "outa = v(img)\n",
        "print(outa)"
      ],
      "metadata": {
        "id": "GYzOR118YMGW",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "702ca6a8-1862-4f2d-913f-983f9ad73956"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "tf.Tensor(\n",
            "[[ 1.8189895  -2.9456773   3.1443267  -1.1026647   0.23636061  0.88873965\n",
            "  -2.1820972   0.377267   -1.5959678  -1.7842641 ]\n",
            " [ 0.8837092  -2.3297973   1.4872115  -1.6793201   1.268717    0.75414175\n",
            "  -0.925108   -0.1175254  -1.3871202  -0.61367285]], shape=(2, 10), dtype=float32)\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import keras\n",
        "num_classes = 10\n",
        "input_shape = (32, 32, 3)\n",
        "batch_size = 32\n",
        "\n",
        "(x_train, y_train), (x_test, y_test) = keras.datasets.cifar10.load_data()\n",
        "\n",
        "y_train = keras.utils.to_categorical(y_train, num_classes)\n",
        "y_test = keras.utils.to_categorical(y_test, num_classes)\n",
        "\n",
        "print(f\"x_train shape: {x_train.shape} - y_train shape: {y_train.shape}\")\n",
        "print(f\"x_test shape: {x_test.shape} - y_test shape: {y_test.shape}\")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "RYBkNJqCJor_",
        "outputId": "83169d32-ec13-4468-af11-8d96c5a997d6"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Downloading data from https://www.cs.toronto.edu/~kriz/cifar-10-python.tar.gz\n",
            "170498071/170498071 [==============================] - 4s 0us/step\n",
            "x_train shape: (50000, 32, 32, 3) - y_train shape: (50000, 10)\n",
            "x_test shape: (10000, 32, 32, 3) - y_test shape: (10000, 10)\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# norr = tf.keras.layers.Resizing(\n",
        "#     224, 224, interpolation=\"bicubic\", crop_to_aspect_ratio=False)\n",
        "train_dataset = tf.data.Dataset.from_tensor_slices((x_train, y_train))\n",
        "train_dataset = train_dataset.shuffle(buffer_size=1024).batch(batch_size)\n",
        "\n",
        "# Prepare the validation dataset.\n",
        "val_dataset = tf.data.Dataset.from_tensor_slices((x_test, y_test))\n",
        "val_dataset = val_dataset.batch(batch_size)"
      ],
      "metadata": {
        "id": "XuEZkSAIMbbR"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "optimizer = keras.optimizers.Adam()\n",
        "# Instantiate a loss function.\n",
        "loss_fn = keras.losses.CategoricalCrossentropy(from_logits=True)\n",
        "train_acc_metric = keras.metrics.CategoricalAccuracy()\n",
        "val_acc_metric = keras.metrics.CategoricalAccuracy()"
      ],
      "metadata": {
        "id": "7HPPJvGPMjmB"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import numpy as np \n",
        "\n",
        "@tf.function\n",
        "def train_step(x, y):\n",
        "    with tf.GradientTape() as tape:\n",
        "        logits = v(x, training=True)\n",
        "        loss_value = loss_fn(y, logits)\n",
        "        # print(loss_value)\n",
        "    grads = tape.gradient(loss_value, v.trainable_weights)\n",
        "    # print(2)\n",
        "    optimizer.apply_gradients(zip(grads, v.trainable_weights))\n",
        "    # print(3)\n",
        "    logits = v(x,training=False)\n",
        "    train_acc_metric.update_state(y, logits)\n",
        "    return loss_value\n",
        "@tf.function\n",
        "def test_step(x, y):\n",
        "    val_logits = v(x, training=False)\n",
        "    val_acc_metric.update_state(y, val_logits)\n",
        "import time\n",
        "train_acc_list=[]\n",
        "train_loss_list=[]\n",
        "epochs = 10\n",
        "for epoch in range(epochs):\n",
        "    print(\"\\nStart of epoch %d\" % (epoch,))\n",
        "    start_time = time.time()\n",
        "    train_loss = []\n",
        "    for step, (x_batch_train, y_batch_train) in enumerate(train_dataset):\n",
        "        loss_value = train_step(x_batch_train, y_batch_train)\n",
        "        train_loss.append(float(loss_value))\n",
        "        if step % 200 == 0:\n",
        "            print(\n",
        "                \"Training loss (for one batch) at step %d: %.4f\"\n",
        "                % (step, float(loss_value))\n",
        "            )\n",
        "            print(\"Seen so far: %d samples\" % ((step + 1) * batch_size))\n",
        "    train_loss_list.append(np.mean(train_loss))\n",
        "    train_acc = train_acc_metric.result()\n",
        "    print(\"Training acc over epoch: %.4f\" % (float(train_acc),))\n",
        "    train_acc_list.append(float(train_acc))\n",
        "    train_acc_metric.reset_states()\n",
        "    for x_batch_val, y_batch_val in val_dataset:\n",
        "        test_step(x_batch_val, y_batch_val)\n",
        "\n",
        "    val_acc = val_acc_metric.result()\n",
        "    val_acc_metric.reset_states()\n",
        "    print(\"Validation acc: %.4f\" % (float(val_acc),))\n",
        "    print(\"Time taken: %.2fs\" % (time.time() - start_time))"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "TycS5BAkXmGq",
        "outputId": "79930a71-dcf7-4432-c7f0-46117f32243b"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            "Start of epoch 0\n",
            "Training loss (for one batch) at step 0: 3.1575\n",
            "Seen so far: 32 samples\n",
            "Training loss (for one batch) at step 200: 1.8259\n",
            "Seen so far: 6432 samples\n",
            "Training loss (for one batch) at step 400: 1.8718\n",
            "Seen so far: 12832 samples\n",
            "Training loss (for one batch) at step 600: 1.9602\n",
            "Seen so far: 19232 samples\n",
            "Training loss (for one batch) at step 800: 1.7096\n",
            "Seen so far: 25632 samples\n",
            "Training loss (for one batch) at step 1000: 1.5373\n",
            "Seen so far: 32032 samples\n",
            "Training loss (for one batch) at step 1200: 1.2513\n",
            "Seen so far: 38432 samples\n",
            "Training loss (for one batch) at step 1400: 1.2664\n",
            "Seen so far: 44832 samples\n",
            "Training acc over epoch: 0.4150\n",
            "Validation acc: 0.4709\n",
            "Time taken: 110.04s\n",
            "\n",
            "Start of epoch 1\n",
            "Training loss (for one batch) at step 0: 1.5677\n",
            "Seen so far: 32 samples\n",
            "Training loss (for one batch) at step 200: 1.7078\n",
            "Seen so far: 6432 samples\n",
            "Training loss (for one batch) at step 400: 1.8163\n",
            "Seen so far: 12832 samples\n",
            "Training loss (for one batch) at step 600: 1.4782\n",
            "Seen so far: 19232 samples\n",
            "Training loss (for one batch) at step 800: 1.2524\n",
            "Seen so far: 25632 samples\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import matplotlib.pyplot as plt\n",
        "plt.plot(train_loss_list)\n",
        "plt.show()"
      ],
      "metadata": {
        "id": "Evx-fSliBx_v"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import matplotlib.pyplot as plt\n",
        "plt.plot(train_acc_list)\n",
        "plt.show()"
      ],
      "metadata": {
        "id": "XVck4fpS9Ic-"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "BqIn4tSI9KSj"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}