vae.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/Zoybzo/generative-models-from-scratch/blob/main/vae.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install loguru"
      ],
      "metadata": {
        "id": "9v0cs76-Ze5N",
        "outputId": "55299b57-11aa-41e6-e6fb-bf2582a3447d",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "id": "9v0cs76-Ze5N",
      "execution_count": 1,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Collecting loguru\n",
            "  Downloading loguru-0.7.3-py3-none-any.whl.metadata (22 kB)\n",
            "Downloading loguru-0.7.3-py3-none-any.whl (61 kB)\n",
            "\u001b[?25l   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m0.0/61.6 kB\u001b[0m \u001b[31m?\u001b[0m eta \u001b[36m-:--:--\u001b[0m\r\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m61.6/61.6 kB\u001b[0m \u001b[31m5.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hInstalling collected packages: loguru\n",
            "Successfully installed loguru-0.7.3\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "id": "initial_id",
      "metadata": {
        "collapsed": true,
        "ExecuteTime": {
          "end_time": "2024-12-15T03:15:34.116139Z",
          "start_time": "2024-12-15T03:15:34.000771Z"
        },
        "id": "initial_id"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "import math\n",
        "import matplotlib as plt\n",
        "from tqdm import trange\n",
        "from PIL import Image\n",
        "from loguru import logger\n",
        "\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import torchvision\n",
        "from torchvision import transforms\n",
        "from torch.utils.data import DataLoader"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "torch.cuda.is_available()"
      ],
      "metadata": {
        "id": "zG86QVHAZnwY",
        "outputId": "7f882f2a-edfb-4f42-b7bc-f15e29bbf783",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "id": "zG86QVHAZnwY",
      "execution_count": 3,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "True"
            ]
          },
          "metadata": {},
          "execution_count": 3
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# New Section"
      ],
      "metadata": {
        "id": "zQjwVkXJZJAI"
      },
      "id": "zQjwVkXJZJAI"
    },
    {
      "cell_type": "code",
      "outputs": [],
      "source": [
        "train_batch_size = 32\n",
        "test_batch_size = 32\n",
        "epochs = 500\n",
        "input_dim = 784\n",
        "latent_dim = 64\n",
        "output_dim = 784\n",
        "lr = 1e-5\n",
        "eval_num_images = 1\n",
        "save_dir = \"./generate/vae/\"\n",
        "if not os.path.exists(save_dir):\n",
        "    os.makedirs(save_dir)\n",
        "device = \"cuda\" if torch.cuda.is_available() else \"cpu\""
      ],
      "metadata": {
        "ExecuteTime": {
          "end_time": "2024-12-14T10:53:45.313174Z",
          "start_time": "2024-12-14T10:53:45.309931Z"
        },
        "id": "ca9ceae55bae9824"
      },
      "id": "ca9ceae55bae9824",
      "execution_count": 12
    },
    {
      "cell_type": "code",
      "outputs": [],
      "source": [
        "class VAE(nn.Module):\n",
        "    def __init__(self, input_dim, latent_dim, output_dim):\n",
        "        super().__init__()\n",
        "        # encoder\n",
        "        self.encoder_mean = nn.Linear(input_dim, latent_dim)\n",
        "        # 方差必须是非负数，但是MLP的输出有正有负，直接输出方差会导致计算KL损失时出现nan（ln(-1)）\n",
        "        # 因此我们在取对数方差，输出ln(sigma)，ln可正可负，且保证了sigma为非负数\n",
        "        # self.encoder_std = nn.Linear(input_dim, latent_dim)\n",
        "        self.encoder_ln_std = nn.Linear(input_dim, latent_dim)\n",
        "        self.decoder = nn.Linear(latent_dim, output_dim)\n",
        "\n",
        "    def forward(self, x):\n",
        "        mean, ln_std = self.encoder_mean(x), self.encoder_ln_std(x)\n",
        "        epsilon = torch.randn_like(mean)\n",
        "        z = mean + torch.exp(ln_std) * epsilon\n",
        "        x = self.decoder(z)\n",
        "        return x, mean, ln_std\n",
        "\n",
        "    def calculate_loss(self, x):\n",
        "        x, mean, ln_std = self.forward(x)\n",
        "        mse_loss = nn.MSELoss()(x, feature)\n",
        "        # kl = 0.5 * (mean ^ 2 + std ^ 2 - 2ln(std) - 1)\n",
        "        kl_loss = torch.mean(0.5 * torch.sum(mean ** 2 + torch.exp(ln_std) ** 2 - 2 * ln_std - 1))\n",
        "        loss = mse_loss + kl_loss\n",
        "        return loss"
      ],
      "metadata": {
        "ExecuteTime": {
          "end_time": "2024-12-15T03:31:42.837977Z",
          "start_time": "2024-12-15T03:31:42.827018Z"
        },
        "id": "3c7edfa3912e770b"
      },
      "id": "3c7edfa3912e770b",
      "execution_count": 5
    },
    {
      "cell_type": "code",
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Downloading http://fashion-mnist.s3-website.eu-central-1.amazonaws.com/train-images-idx3-ubyte.gz\n",
            "Downloading http://fashion-mnist.s3-website.eu-central-1.amazonaws.com/train-images-idx3-ubyte.gz to ./data/FashionMNIST/raw/train-images-idx3-ubyte.gz\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 26.4M/26.4M [00:02<00:00, 10.7MB/s]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Extracting ./data/FashionMNIST/raw/train-images-idx3-ubyte.gz to ./data/FashionMNIST/raw\n",
            "\n",
            "Downloading http://fashion-mnist.s3-website.eu-central-1.amazonaws.com/train-labels-idx1-ubyte.gz\n",
            "Downloading http://fashion-mnist.s3-website.eu-central-1.amazonaws.com/train-labels-idx1-ubyte.gz to ./data/FashionMNIST/raw/train-labels-idx1-ubyte.gz\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 29.5k/29.5k [00:00<00:00, 168kB/s]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Extracting ./data/FashionMNIST/raw/train-labels-idx1-ubyte.gz to ./data/FashionMNIST/raw\n",
            "\n",
            "Downloading http://fashion-mnist.s3-website.eu-central-1.amazonaws.com/t10k-images-idx3-ubyte.gz\n",
            "Downloading http://fashion-mnist.s3-website.eu-central-1.amazonaws.com/t10k-images-idx3-ubyte.gz to ./data/FashionMNIST/raw/t10k-images-idx3-ubyte.gz\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 4.42M/4.42M [00:01<00:00, 3.08MB/s]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Extracting ./data/FashionMNIST/raw/t10k-images-idx3-ubyte.gz to ./data/FashionMNIST/raw\n",
            "\n",
            "Downloading http://fashion-mnist.s3-website.eu-central-1.amazonaws.com/t10k-labels-idx1-ubyte.gz\n",
            "Downloading http://fashion-mnist.s3-website.eu-central-1.amazonaws.com/t10k-labels-idx1-ubyte.gz to ./data/FashionMNIST/raw/t10k-labels-idx1-ubyte.gz\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 5.15k/5.15k [00:00<00:00, 6.25MB/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Extracting ./data/FashionMNIST/raw/t10k-labels-idx1-ubyte.gz to ./data/FashionMNIST/raw\n",
            "\n",
            "Train: 60000; Test: 10000)\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\n"
          ]
        }
      ],
      "source": [
        "# Dataset\n",
        "transform = transforms.ToTensor()\n",
        "fashion_mnist_train = torchvision.datasets.FashionMNIST(root='./data', train=True, download=True, transform=transform)\n",
        "fashion_mnist_test = torchvision.datasets.FashionMNIST(root='./data', train=False, download=True, transform=transform)\n",
        "print(f\"Train: {len(fashion_mnist_train)}; Test: {len(fashion_mnist_test)})\")"
      ],
      "metadata": {
        "collapsed": true,
        "ExecuteTime": {
          "end_time": "2024-12-14T10:53:45.367993Z",
          "start_time": "2024-12-14T10:53:45.319833Z"
        },
        "id": "509496815c2d18d7",
        "outputId": "7bdae786-721e-488b-c9f6-8092f358d92d",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "id": "509496815c2d18d7",
      "execution_count": 6
    },
    {
      "cell_type": "code",
      "outputs": [],
      "source": [
        "# Loader\n",
        "train_loader = DataLoader(fashion_mnist_train, batch_size=train_batch_size, shuffle=True, drop_last=True)\n",
        "test_loader = DataLoader(fashion_mnist_test, batch_size=test_batch_size, shuffle=True, drop_last=True)"
      ],
      "metadata": {
        "ExecuteTime": {
          "end_time": "2024-12-14T10:53:45.373856Z",
          "start_time": "2024-12-14T10:53:45.369826Z"
        },
        "id": "5e3220ac2cc9ecdc"
      },
      "id": "5e3220ac2cc9ecdc",
      "execution_count": 7
    },
    {
      "cell_type": "code",
      "outputs": [
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r  0%|          | 0/500 [00:00<?, ?it/s]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 1, Train Loss: 1.1339\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r  0%|          | 1/500 [00:12<1:42:54, 12.37s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 1, Test Loss:  0.5711\n",
            "Epoch 2, Train Loss: 0.4634\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r  0%|          | 2/500 [00:25<1:45:36, 12.72s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 2, Test Loss:  0.3770\n",
            "Epoch 3, Train Loss: 0.3231\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r  1%|          | 3/500 [00:37<1:43:31, 12.50s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 3, Test Loss:  0.2763\n",
            "Epoch 4, Train Loss: 0.2427\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r  1%|          | 4/500 [00:49<1:42:09, 12.36s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 4, Test Loss:  0.2134\n",
            "Epoch 5, Train Loss: 0.1906\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r  1%|          | 5/500 [01:01<1:41:14, 12.27s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 5, Test Loss:  0.1706\n",
            "Epoch 6, Train Loss: 0.1541\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r  1%|          | 6/500 [01:13<1:40:43, 12.23s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 6, Test Loss:  0.1401\n",
            "Epoch 7, Train Loss: 0.1271\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r  1%|▏         | 7/500 [01:26<1:40:17, 12.21s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 7, Test Loss:  0.1166\n",
            "Epoch 8, Train Loss: 0.1064\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r  2%|▏         | 8/500 [01:38<1:39:46, 12.17s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 8, Test Loss:  0.0983\n",
            "Epoch 9, Train Loss: 0.0900\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r  2%|▏         | 9/500 [01:50<1:39:33, 12.17s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 9, Test Loss:  0.0837\n",
            "Epoch 10, Train Loss: 0.0768\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r  2%|▏         | 10/500 [02:02<1:39:00, 12.12s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 10, Test Loss:  0.0719\n",
            "Epoch 11, Train Loss: 0.0660\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r  2%|▏         | 11/500 [02:14<1:38:30, 12.09s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 11, Test Loss:  0.0621\n",
            "Epoch 12, Train Loss: 0.0570\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r  2%|▏         | 12/500 [02:26<1:38:21, 12.09s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 12, Test Loss:  0.0538\n",
            "Epoch 13, Train Loss: 0.0495\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r  3%|▎         | 13/500 [02:38<1:38:01, 12.08s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 13, Test Loss:  0.0469\n",
            "Epoch 14, Train Loss: 0.0431\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r  3%|▎         | 14/500 [02:50<1:37:59, 12.10s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 14, Test Loss:  0.0410\n",
            "Epoch 15, Train Loss: 0.0377\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r  3%|▎         | 15/500 [03:03<1:38:25, 12.18s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 15, Test Loss:  0.0359\n",
            "Epoch 16, Train Loss: 0.0330\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r  3%|▎         | 16/500 [03:15<1:38:21, 12.19s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 16, Test Loss:  0.0317\n",
            "Epoch 17, Train Loss: 0.0290\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r  3%|▎         | 17/500 [03:27<1:38:22, 12.22s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 17, Test Loss:  0.0278\n",
            "Epoch 18, Train Loss: 0.0255\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r  4%|▎         | 18/500 [03:39<1:37:59, 12.20s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 18, Test Loss:  0.0245\n",
            "Epoch 19, Train Loss: 0.0226\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r  4%|▍         | 19/500 [03:51<1:37:55, 12.22s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 19, Test Loss:  0.0217\n",
            "Epoch 20, Train Loss: 0.0200\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r  4%|▍         | 20/500 [04:04<1:37:30, 12.19s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 20, Test Loss:  0.0192\n",
            "Epoch 21, Train Loss: 0.0177\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r  4%|▍         | 21/500 [04:16<1:37:25, 12.20s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 21, Test Loss:  0.0172\n",
            "Epoch 22, Train Loss: 0.0157\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r  4%|▍         | 22/500 [04:28<1:37:19, 12.22s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 22, Test Loss:  0.0152\n",
            "Epoch 23, Train Loss: 0.0140\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r  5%|▍         | 23/500 [04:40<1:37:31, 12.27s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 23, Test Loss:  0.0136\n",
            "Epoch 24, Train Loss: 0.0125\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r  5%|▍         | 24/500 [04:53<1:37:08, 12.25s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 24, Test Loss:  0.0121\n",
            "Epoch 25, Train Loss: 0.0112\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r  5%|▌         | 25/500 [05:05<1:36:44, 12.22s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 25, Test Loss:  0.0109\n",
            "Epoch 26, Train Loss: 0.0100\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r  5%|▌         | 26/500 [05:17<1:36:17, 12.19s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 26, Test Loss:  0.0098\n",
            "Epoch 27, Train Loss: 0.0090\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r  5%|▌         | 27/500 [05:29<1:36:05, 12.19s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 27, Test Loss:  0.0088\n",
            "Epoch 28, Train Loss: 0.0082\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r  6%|▌         | 28/500 [05:41<1:35:34, 12.15s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 28, Test Loss:  0.0080\n",
            "Epoch 29, Train Loss: 0.0074\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r  6%|▌         | 29/500 [05:53<1:35:07, 12.12s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 29, Test Loss:  0.0072\n",
            "Epoch 30, Train Loss: 0.0067\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r  6%|▌         | 30/500 [06:05<1:34:53, 12.11s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 30, Test Loss:  0.0066\n",
            "Epoch 31, Train Loss: 0.0062\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r  6%|▌         | 31/500 [06:18<1:35:02, 12.16s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 31, Test Loss:  0.0060\n",
            "Epoch 32, Train Loss: 0.0057\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r  6%|▋         | 32/500 [06:30<1:35:03, 12.19s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 32, Test Loss:  0.0055\n",
            "Epoch 33, Train Loss: 0.0052\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r  7%|▋         | 33/500 [06:42<1:34:39, 12.16s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 33, Test Loss:  0.0051\n",
            "Epoch 34, Train Loss: 0.0048\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r  7%|▋         | 34/500 [06:54<1:34:23, 12.15s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 34, Test Loss:  0.0047\n",
            "Epoch 35, Train Loss: 0.0045\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r  7%|▋         | 35/500 [07:06<1:34:01, 12.13s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 35, Test Loss:  0.0044\n",
            "Epoch 36, Train Loss: 0.0042\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r  7%|▋         | 36/500 [07:18<1:33:46, 12.13s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 36, Test Loss:  0.0042\n",
            "Epoch 37, Train Loss: 0.0040\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r  7%|▋         | 37/500 [07:30<1:33:31, 12.12s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 37, Test Loss:  0.0039\n",
            "Epoch 38, Train Loss: 0.0038\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r  8%|▊         | 38/500 [07:42<1:33:05, 12.09s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 38, Test Loss:  0.0037\n",
            "Epoch 39, Train Loss: 0.0036\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r  8%|▊         | 39/500 [07:54<1:32:44, 12.07s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 39, Test Loss:  0.0035\n",
            "Epoch 40, Train Loss: 0.0035\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r  8%|▊         | 40/500 [08:07<1:32:43, 12.09s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 40, Test Loss:  0.0034\n",
            "Epoch 41, Train Loss: 0.0033\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r  8%|▊         | 41/500 [08:19<1:32:31, 12.10s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 41, Test Loss:  0.0033\n",
            "Epoch 42, Train Loss: 0.0032\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r  8%|▊         | 42/500 [08:31<1:32:09, 12.07s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 42, Test Loss:  0.0032\n",
            "Epoch 43, Train Loss: 0.0031\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r  9%|▊         | 43/500 [08:43<1:31:52, 12.06s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 43, Test Loss:  0.0031\n",
            "Epoch 44, Train Loss: 0.0030\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r  9%|▉         | 44/500 [08:55<1:32:05, 12.12s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 44, Test Loss:  0.0030\n",
            "Epoch 45, Train Loss: 0.0030\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r  9%|▉         | 45/500 [09:07<1:32:04, 12.14s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 45, Test Loss:  0.0029\n",
            "Epoch 46, Train Loss: 0.0029\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r  9%|▉         | 46/500 [09:19<1:31:52, 12.14s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 46, Test Loss:  0.0029\n",
            "Epoch 47, Train Loss: 0.0029\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r  9%|▉         | 47/500 [09:31<1:31:27, 12.11s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 47, Test Loss:  0.0029\n",
            "Epoch 48, Train Loss: 0.0028\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r 10%|▉         | 48/500 [09:44<1:31:24, 12.13s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 48, Test Loss:  0.0028\n",
            "Epoch 49, Train Loss: 0.0028\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r 10%|▉         | 49/500 [09:56<1:31:10, 12.13s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 49, Test Loss:  0.0028\n",
            "Epoch 50, Train Loss: 0.0028\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r 10%|█         | 50/500 [10:08<1:30:58, 12.13s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 50, Test Loss:  0.0028\n",
            "Epoch 51, Train Loss: 0.0028\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r 10%|█         | 51/500 [10:20<1:30:42, 12.12s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 51, Test Loss:  0.0028\n",
            "Epoch 52, Train Loss: 0.0028\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r 10%|█         | 52/500 [10:32<1:30:35, 12.13s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 52, Test Loss:  0.0027\n",
            "Epoch 53, Train Loss: 0.0028\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r 11%|█         | 53/500 [10:44<1:30:28, 12.14s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 53, Test Loss:  0.0027\n",
            "Epoch 54, Train Loss: 0.0027\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r 11%|█         | 54/500 [10:56<1:30:22, 12.16s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 54, Test Loss:  0.0027\n",
            "Epoch 55, Train Loss: 0.0027\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r 11%|█         | 55/500 [11:09<1:30:12, 12.16s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 55, Test Loss:  0.0027\n",
            "Epoch 56, Train Loss: 0.0027\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r 11%|█         | 56/500 [11:21<1:29:51, 12.14s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 56, Test Loss:  0.0027\n",
            "Epoch 57, Train Loss: 0.0027\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r 11%|█▏        | 57/500 [11:33<1:29:22, 12.11s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 57, Test Loss:  0.0027\n",
            "Epoch 58, Train Loss: 0.0027\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r 12%|█▏        | 58/500 [11:45<1:29:17, 12.12s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 58, Test Loss:  0.0027\n",
            "Epoch 59, Train Loss: 0.0027\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r 12%|█▏        | 59/500 [11:57<1:29:19, 12.15s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 59, Test Loss:  0.0027\n",
            "Epoch 60, Train Loss: 0.0027\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r 12%|█▏        | 60/500 [12:09<1:29:05, 12.15s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 60, Test Loss:  0.0027\n",
            "Epoch 61, Train Loss: 0.0027\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r 12%|█▏        | 61/500 [12:21<1:28:43, 12.13s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 61, Test Loss:  0.0027\n",
            "Epoch 62, Train Loss: 0.0027\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r 12%|█▏        | 62/500 [12:33<1:28:30, 12.13s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 62, Test Loss:  0.0027\n",
            "Epoch 63, Train Loss: 0.0027\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r 13%|█▎        | 63/500 [12:46<1:28:16, 12.12s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 63, Test Loss:  0.0027\n",
            "Epoch 64, Train Loss: 0.0027\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r 13%|█▎        | 64/500 [12:58<1:28:02, 12.12s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 64, Test Loss:  0.0027\n",
            "Epoch 65, Train Loss: 0.0027\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r 13%|█▎        | 65/500 [13:10<1:27:55, 12.13s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 65, Test Loss:  0.0027\n",
            "Epoch 66, Train Loss: 0.0027\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r 13%|█▎        | 66/500 [13:22<1:27:53, 12.15s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 66, Test Loss:  0.0027\n",
            "Epoch 67, Train Loss: 0.0027\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r 13%|█▎        | 67/500 [13:34<1:27:29, 12.12s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 67, Test Loss:  0.0027\n",
            "Epoch 68, Train Loss: 0.0027\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r 14%|█▎        | 68/500 [13:46<1:26:53, 12.07s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 68, Test Loss:  0.0027\n",
            "Epoch 69, Train Loss: 0.0027\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r 14%|█▍        | 69/500 [13:58<1:26:35, 12.05s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 69, Test Loss:  0.0027\n",
            "Epoch 70, Train Loss: 0.0027\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r 14%|█▍        | 70/500 [14:10<1:26:20, 12.05s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 70, Test Loss:  0.0027\n",
            "Epoch 71, Train Loss: 0.0027\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r 14%|█▍        | 71/500 [14:22<1:26:12, 12.06s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 71, Test Loss:  0.0027\n",
            "Epoch 72, Train Loss: 0.0027\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r 14%|█▍        | 72/500 [14:34<1:25:59, 12.06s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 72, Test Loss:  0.0027\n",
            "Epoch 73, Train Loss: 0.0027\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r 15%|█▍        | 73/500 [14:46<1:25:27, 12.01s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 73, Test Loss:  0.0027\n",
            "Epoch 74, Train Loss: 0.0027\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r 15%|█▍        | 74/500 [14:58<1:25:17, 12.01s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 74, Test Loss:  0.0027\n",
            "Epoch 75, Train Loss: 0.0027\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r 15%|█▌        | 75/500 [15:10<1:25:03, 12.01s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 75, Test Loss:  0.0027\n",
            "Epoch 76, Train Loss: 0.0027\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r 15%|█▌        | 76/500 [15:22<1:25:21, 12.08s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 76, Test Loss:  0.0027\n",
            "Epoch 77, Train Loss: 0.0027\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r 15%|█▌        | 77/500 [15:35<1:26:54, 12.33s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 77, Test Loss:  0.0027\n",
            "Epoch 78, Train Loss: 0.0027\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r 16%|█▌        | 78/500 [15:49<1:28:44, 12.62s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 78, Test Loss:  0.0027\n",
            "Epoch 79, Train Loss: 0.0027\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r 16%|█▌        | 79/500 [16:01<1:29:08, 12.70s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 79, Test Loss:  0.0027\n",
            "Epoch 80, Train Loss: 0.0027\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r 16%|█▌        | 80/500 [16:14<1:27:55, 12.56s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 80, Test Loss:  0.0027\n",
            "Epoch 81, Train Loss: 0.0027\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r 16%|█▌        | 81/500 [16:26<1:26:35, 12.40s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 81, Test Loss:  0.0027\n",
            "Epoch 82, Train Loss: 0.0027\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r 16%|█▋        | 82/500 [16:38<1:25:44, 12.31s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 82, Test Loss:  0.0027\n",
            "Epoch 83, Train Loss: 0.0027\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r 17%|█▋        | 83/500 [16:50<1:25:06, 12.25s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 83, Test Loss:  0.0027\n",
            "Epoch 84, Train Loss: 0.0027\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r 17%|█▋        | 84/500 [17:02<1:24:35, 12.20s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 84, Test Loss:  0.0027\n",
            "Epoch 85, Train Loss: 0.0027\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r 17%|█▋        | 85/500 [17:14<1:24:21, 12.20s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 85, Test Loss:  0.0027\n",
            "Epoch 86, Train Loss: 0.0027\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r 17%|█▋        | 86/500 [17:26<1:23:41, 12.13s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 86, Test Loss:  0.0027\n",
            "Epoch 87, Train Loss: 0.0027\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r 17%|█▋        | 87/500 [17:38<1:23:18, 12.10s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 87, Test Loss:  0.0027\n",
            "Epoch 88, Train Loss: 0.0027\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r 18%|█▊        | 88/500 [17:50<1:23:08, 12.11s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 88, Test Loss:  0.0027\n",
            "Epoch 89, Train Loss: 0.0027\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r 18%|█▊        | 89/500 [18:03<1:23:02, 12.12s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 89, Test Loss:  0.0027\n",
            "Epoch 90, Train Loss: 0.0027\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r 18%|█▊        | 90/500 [18:15<1:22:50, 12.12s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 90, Test Loss:  0.0027\n",
            "Epoch 91, Train Loss: 0.0027\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r 18%|█▊        | 91/500 [18:27<1:22:14, 12.07s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 91, Test Loss:  0.0027\n",
            "Epoch 92, Train Loss: 0.0027\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r 18%|█▊        | 92/500 [18:39<1:22:04, 12.07s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 92, Test Loss:  0.0027\n",
            "Epoch 93, Train Loss: 0.0027\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r 19%|█▊        | 93/500 [18:51<1:21:45, 12.05s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 93, Test Loss:  0.0027\n",
            "Epoch 94, Train Loss: 0.0027\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r 19%|█▉        | 94/500 [19:03<1:21:46, 12.09s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 94, Test Loss:  0.0027\n",
            "Epoch 95, Train Loss: 0.0027\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r 19%|█▉        | 95/500 [19:15<1:21:38, 12.09s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 95, Test Loss:  0.0027\n",
            "Epoch 96, Train Loss: 0.0027\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r 19%|█▉        | 96/500 [19:27<1:21:28, 12.10s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 96, Test Loss:  0.0027\n",
            "Epoch 97, Train Loss: 0.0027\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r 19%|█▉        | 97/500 [19:39<1:21:14, 12.10s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 97, Test Loss:  0.0027\n",
            "Epoch 98, Train Loss: 0.0027\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r 20%|█▉        | 98/500 [19:51<1:21:13, 12.12s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 98, Test Loss:  0.0027\n",
            "Epoch 99, Train Loss: 0.0027\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r 20%|█▉        | 99/500 [20:04<1:21:11, 12.15s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 99, Test Loss:  0.0027\n",
            "Epoch 100, Train Loss: 0.0027\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r 20%|██        | 100/500 [20:16<1:20:54, 12.14s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 100, Test Loss:  0.0027\n",
            "Epoch 101, Train Loss: 0.0027\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r 20%|██        | 101/500 [20:28<1:20:43, 12.14s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 101, Test Loss:  0.0027\n",
            "Epoch 102, Train Loss: 0.0027\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r 20%|██        | 102/500 [20:40<1:20:36, 12.15s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 102, Test Loss:  0.0027\n",
            "Epoch 103, Train Loss: 0.0027\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r 21%|██        | 103/500 [20:52<1:20:17, 12.13s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 103, Test Loss:  0.0027\n",
            "Epoch 104, Train Loss: 0.0027\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r 21%|██        | 104/500 [21:04<1:20:15, 12.16s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 104, Test Loss:  0.0027\n",
            "Epoch 105, Train Loss: 0.0027\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r 21%|██        | 105/500 [21:16<1:19:49, 12.13s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 105, Test Loss:  0.0027\n",
            "Epoch 106, Train Loss: 0.0027\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r 21%|██        | 106/500 [21:28<1:19:23, 12.09s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 106, Test Loss:  0.0027\n",
            "Epoch 107, Train Loss: 0.0027\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r 21%|██▏       | 107/500 [21:40<1:19:01, 12.06s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 107, Test Loss:  0.0027\n",
            "Epoch 108, Train Loss: 0.0027\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r 22%|██▏       | 108/500 [21:52<1:18:30, 12.02s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 108, Test Loss:  0.0027\n",
            "Epoch 109, Train Loss: 0.0027\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r 22%|██▏       | 109/500 [22:04<1:18:26, 12.04s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 109, Test Loss:  0.0027\n",
            "Epoch 110, Train Loss: 0.0027\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r 22%|██▏       | 110/500 [22:16<1:18:23, 12.06s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 110, Test Loss:  0.0027\n",
            "Epoch 111, Train Loss: 0.0027\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r 22%|██▏       | 111/500 [22:28<1:18:09, 12.06s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 111, Test Loss:  0.0027\n",
            "Epoch 112, Train Loss: 0.0027\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r 22%|██▏       | 112/500 [22:41<1:17:57, 12.06s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 112, Test Loss:  0.0027\n",
            "Epoch 113, Train Loss: 0.0027\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r 23%|██▎       | 113/500 [22:53<1:17:35, 12.03s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 113, Test Loss:  0.0027\n",
            "Epoch 114, Train Loss: 0.0027\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r 23%|██▎       | 114/500 [23:05<1:17:21, 12.02s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 114, Test Loss:  0.0027\n",
            "Epoch 115, Train Loss: 0.0027\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r 23%|██▎       | 115/500 [23:16<1:17:03, 12.01s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 115, Test Loss:  0.0027\n",
            "Epoch 116, Train Loss: 0.0027\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r 23%|██▎       | 116/500 [23:28<1:16:45, 11.99s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 116, Test Loss:  0.0027\n",
            "Epoch 117, Train Loss: 0.0027\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r 23%|██▎       | 117/500 [23:40<1:16:39, 12.01s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 117, Test Loss:  0.0027\n",
            "Epoch 118, Train Loss: 0.0027\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r 24%|██▎       | 118/500 [23:53<1:16:30, 12.02s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 118, Test Loss:  0.0027\n",
            "Epoch 119, Train Loss: 0.0027\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r 24%|██▍       | 119/500 [24:05<1:16:28, 12.04s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 119, Test Loss:  0.0027\n",
            "Epoch 120, Train Loss: 0.0027\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r 24%|██▍       | 120/500 [24:17<1:16:28, 12.07s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 120, Test Loss:  0.0027\n",
            "Epoch 121, Train Loss: 0.0027\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r 24%|██▍       | 121/500 [24:29<1:16:12, 12.07s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 121, Test Loss:  0.0027\n",
            "Epoch 122, Train Loss: 0.0027\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r 24%|██▍       | 122/500 [24:41<1:15:57, 12.06s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 122, Test Loss:  0.0027\n",
            "Epoch 123, Train Loss: 0.0027\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r 25%|██▍       | 123/500 [24:53<1:15:35, 12.03s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 123, Test Loss:  0.0027\n",
            "Epoch 124, Train Loss: 0.0027\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r 25%|██▍       | 124/500 [25:05<1:15:16, 12.01s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 124, Test Loss:  0.0027\n",
            "Epoch 125, Train Loss: 0.0027\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r 25%|██▌       | 125/500 [25:17<1:15:00, 12.00s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 125, Test Loss:  0.0027\n",
            "Epoch 126, Train Loss: 0.0027\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r 25%|██▌       | 126/500 [25:29<1:14:37, 11.97s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 126, Test Loss:  0.0027\n",
            "Epoch 127, Train Loss: 0.0027\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r 25%|██▌       | 127/500 [25:41<1:14:24, 11.97s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 127, Test Loss:  0.0027\n",
            "Epoch 128, Train Loss: 0.0027\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r 26%|██▌       | 128/500 [25:53<1:14:14, 11.98s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 128, Test Loss:  0.0027\n",
            "Epoch 129, Train Loss: 0.0027\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r 26%|██▌       | 129/500 [26:05<1:14:35, 12.06s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 129, Test Loss:  0.0027\n",
            "Epoch 130, Train Loss: 0.0027\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r 26%|██▌       | 130/500 [26:17<1:14:32, 12.09s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 130, Test Loss:  0.0027\n",
            "Epoch 131, Train Loss: 0.0027\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r 26%|██▌       | 131/500 [26:29<1:14:14, 12.07s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 131, Test Loss:  0.0027\n",
            "Epoch 132, Train Loss: 0.0027\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r 26%|██▋       | 132/500 [26:41<1:14:08, 12.09s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 132, Test Loss:  0.0027\n",
            "Epoch 133, Train Loss: 0.0027\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r 27%|██▋       | 133/500 [26:53<1:13:45, 12.06s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 133, Test Loss:  0.0027\n",
            "Epoch 134, Train Loss: 0.0027\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r 27%|██▋       | 134/500 [27:05<1:13:32, 12.05s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 134, Test Loss:  0.0027\n",
            "Epoch 135, Train Loss: 0.0027\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r 27%|██▋       | 135/500 [27:17<1:13:25, 12.07s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 135, Test Loss:  0.0027\n",
            "Epoch 136, Train Loss: 0.0027\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r 27%|██▋       | 136/500 [27:29<1:13:02, 12.04s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 136, Test Loss:  0.0027\n",
            "Epoch 137, Train Loss: 0.0027\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r 27%|██▋       | 137/500 [27:41<1:12:44, 12.02s/it]"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 137, Test Loss:  0.0027\n",
            "Epoch 138, Train Loss: 0.0027\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r 28%|██▊       | 138/500 [27:53<1:12:20, 11.99s/it]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Epoch 138, Test Loss:  0.0027\n",
            "Epoch 139, Train Loss: 0.0027\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 28%|██▊       | 139/500 [28:05<1:12:12, 12.00s/it]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Epoch 139, Test Loss:  0.0027\n",
            "Epoch 140, Train Loss: 0.0027\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 28%|██▊       | 140/500 [28:17<1:12:05, 12.02s/it]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Epoch 140, Test Loss:  0.0027\n",
            "Epoch 141, Train Loss: 0.0027\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 28%|██▊       | 141/500 [28:29<1:11:42, 11.99s/it]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Epoch 141, Test Loss:  0.0027\n",
            "Epoch 142, Train Loss: 0.0027\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 28%|██▊       | 142/500 [28:41<1:11:26, 11.97s/it]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Epoch 142, Test Loss:  0.0027\n",
            "Epoch 143, Train Loss: 0.0027\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 29%|██▊       | 143/500 [28:53<1:11:01, 11.94s/it]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Epoch 143, Test Loss:  0.0027\n",
            "Epoch 144, Train Loss: 0.0027\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 29%|██▉       | 144/500 [29:05<1:10:48, 11.94s/it]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Epoch 144, Test Loss:  0.0027\n",
            "Epoch 145, Train Loss: 0.0027\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 29%|██▉       | 145/500 [29:17<1:10:33, 11.93s/it]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Epoch 145, Test Loss:  0.0027\n",
            "Epoch 146, Train Loss: 0.0027\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 29%|██▉       | 146/500 [29:29<1:10:25, 11.94s/it]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Epoch 146, Test Loss:  0.0027\n"
          ]
        }
      ],
      "source": [
        "# Create VAE\n",
        "model = VAE(input_dim, latent_dim, output_dim).to(device)\n",
        "optim = torch.optim.AdamW(model.parameters(), lr=lr)\n",
        "# Train\n",
        "for epoch in trange(epochs):\n",
        "    train_loss = 0.0\n",
        "    test_loss = 0.0\n",
        "    # train\n",
        "    model.train()\n",
        "    for idx, (feature, _) in enumerate(train_loader):\n",
        "        optim.zero_grad()\n",
        "        feature = feature.reshape(train_batch_size, -1).to(device)\n",
        "        # loss\n",
        "        loss = model.calculate_loss(feature)\n",
        "        if torch.isnan(loss).any():\n",
        "            logger.error(\"Loss is NaN.\")\n",
        "        loss.backward()\n",
        "        optim.step()\n",
        "        train_loss += loss.item()\n",
        "    # 打印每个 epoch 的平均损失\n",
        "    avg_loss = train_loss / len(train_loader.dataset)\n",
        "    print(f\"Epoch {epoch + 1}, Train Loss: {avg_loss:.4f}\")\n",
        "    # evaluation\n",
        "    model.eval()\n",
        "    with torch.no_grad():\n",
        "        for idx, (feature, _) in enumerate(test_loader):\n",
        "            feature = feature.reshape(test_batch_size, -1).to(device)\n",
        "            loss = model.calculate_loss(feature)\n",
        "            test_loss += loss.item()\n",
        "        avg_test_loss = test_loss / len(test_loader.dataset)\n",
        "        print(f\"Epoch {epoch + 1}, Test Loss: {avg_test_loss: .4f}\")\n",
        "    # generate image\n",
        "    with torch.no_grad():\n",
        "        z = torch.randn(eval_num_images, latent_dim).to(device)\n",
        "        images = model.decoder(z).cpu().reshape(-1, 1, 28, 28)\n",
        "    # save image\n",
        "    for i in range(eval_num_images):\n",
        "        # 获取单张图像，去掉多余的维度\n",
        "        img = images[i, 0].cpu().numpy()\n",
        "        # 将 numpy 数组转换为 PIL 图像\n",
        "        pil_img = Image.fromarray((img * 255).astype('uint8'))  # 转换为 [0, 255] 范围\n",
        "        img_path = os.path.join(save_dir, f\"{epoch}_{i + 1}.png\")\n",
        "        # 保存图像\n",
        "        pil_img.save(img_path)\n"
      ],
      "metadata": {
        "ExecuteTime": {
          "end_time": "2024-12-15T03:32:20.155394Z",
          "start_time": "2024-12-15T03:31:45.498795Z"
        },
        "id": "7a2dd5949b3b0cc9",
        "outputId": "d3d945cc-ce44-436f-c834-1abdc850961e",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "id": "7a2dd5949b3b0cc9",
      "execution_count": null
    },
    {
      "cell_type": "code",
      "outputs": [],
      "source": [],
      "metadata": {
        "id": "91919eb44b05a565"
      },
      "id": "91919eb44b05a565",
      "execution_count": null
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 2
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython2",
      "version": "2.7.6"
    },
    "colab": {
      "provenance": [],
      "toc_visible": true,
      "gpuType": "T4",
      "include_colab_link": true
    },
    "accelerator": "GPU"
  },
  "nbformat": 4,
  "nbformat_minor": 5
}