{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "d6cb560a-707c-4c1e-b50f-2b9c8037a3c7",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "from PIL import Image\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "import torchvision\n",
    "import torchvision.transforms as transforms\n",
    "import matplotlib.pyplot as plt\n",
    "from rbm.model import Model\n",
    "from rbm.entity import Entity, EntityParams, TrainingParams\n",
    "from rbm.matrix import Mat, np, rms_error_accu\n",
    "from rbm.torch import Optimizer\n",
    "import math"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "4065302d-862f-4639-99d7-e396d055434c",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "transform = transforms.Compose([\n",
    "    transforms.ToTensor()\n",
    "])                              "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "ffe48fc6-0f57-48a0-b5e9-aa68cbc8edf1",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "train_data = torchvision.datasets.MNIST(root='./data', train=True, transform=transform, download=True)\n",
    "test_data = torchvision.datasets.MNIST(root='./data', train=False, transform=transform, download=True)\n",
    "\n",
    "train_loader = torch.utils.data.DataLoader(train_data, batch_size=32, shuffle=True, num_workers=2)\n",
    "test_loader = torch.utils.data.DataLoader(test_data, batch_size=32, shuffle=True, num_workers=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "51b90664-e470-4c9a-884d-374a2ece61b2",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "60000\n"
     ]
    }
   ],
   "source": [
    "print(len(train_data))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "33501c0f-6d78-43e9-b2ce-adbe40b6d95a",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[28, 28]\n"
     ]
    }
   ],
   "source": [
    "print(list(train_data[1][0].size()[1:3]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 89,
   "id": "aedf601a-9084-4047-98e8-481f4d41a2d9",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "N = 60000\n",
    "size = 784\n",
    "n_side = 10\n",
    "train_images = np.zeros(shape=(N,size))\n",
    "for i in range(N):\n",
    "    image = train_data[i][0]\n",
    "    flat_array = torch.flatten(image)\n",
    "    train_images[i, :] = np.asarray(flat_array.numpy())\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 90,
   "id": "5ee865c3-2af4-487f-a184-1291f99839b1",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "60000\n"
     ]
    }
   ],
   "source": [
    "print(len(train_images))\n",
    "image_indices = np.random.randint(0, N-1, n_side*n_side)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 91,
   "id": "b4310933-d850-4128-a2dd-27d88c004846",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 143768 characters]",
      "text/plain": [
       "<Figure size 3000x3000 with 100 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig, axes = plt.subplots(n_side, n_side, figsize=(30,30))\n",
    "index = 0             \n",
    "for x in range(n_side):\n",
    "    for y in range(n_side):\n",
    "        inp = train_images[image_indices[index]]\n",
    "        img = np.reshape(inp, (28, 28))\n",
    "        axes[x,y].imshow(np.asnumpy(img))\n",
    "        axes[x,y].axis('off')\n",
    "        index += 1\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 92,
   "id": "08143a53-12a1-4896-a4b2-166dedf0e021",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "class TestModel(Model):\n",
    "\tdef __init__(self, name: str, work_dir: str = '.'):\n",
    "\t\tsuper().__init__(name, work_dir)\n",
    "\t\tself.unit1 = Entity((28*28, 100), EntityParams(do_gaussian_visible=False, do_gaussian_hidden=False), TrainingParams(learning_rate=0.1, momentum=0.9, num_epochs=10000, mini_batch_size=1000, l1_lambda=0.004, do_rao_blackwell=True), True)\n",
    "#\t\tself.unit2 = Entity((500, 64), EntityParams(do_gaussian_visible=False, do_gaussian_hidden=False), TrainingParams(learning_rate=0.04, momentum=0.9, num_epochs=1000, mini_batch_size=100, l2_lambda=0.01, do_rao_blackwell=True), True)\n",
    "\n",
    "\tdef forward(self, x: Mat):\n",
    "\t\tx = self.unit1.forward(x)\n",
    "#\t\tx = self.unit2.forward(x)\n",
    "\t\treturn x\n",
    "\n",
    "\tdef backward(self, x: Mat):\n",
    "#\t\tx = self.unit2.reconstruct(x)\n",
    "\t\tx = self.unit1.reconstruct(x)\n",
    "\t\treturn x\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 93,
   "id": "6b3be9f6-38e1-442d-a7bd-1f3c164ec8d4",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "results/mnist_test-0-state.npz loaded successfully!\n",
      "-------------------------------------------\n",
      "Entity-784x100: progress              : 0%\n",
      "Entity-784x100: err_rms              : 0.014414508244480833\n",
      "Entity-784x100: l2_norm              : 25.413601322380156\n",
      "-------------------------------------------\n",
      "Entity-784x100: progress              : 10%\n",
      "Entity-784x100: err_rms              : 0.016305932137787714\n",
      "Entity-784x100: l2_norm              : 18.362961748495835\n",
      "-------------------------------------------\n",
      "Entity-784x100: progress              : 20%\n",
      "Entity-784x100: err_rms              : 0.01620548503399525\n",
      "Entity-784x100: l2_norm              : 18.48548452029196\n",
      "-------------------------------------------\n",
      "Entity-784x100: progress              : 30%\n",
      "Entity-784x100: err_rms              : 0.016206704015135853\n",
      "Entity-784x100: l2_norm              : 18.482700817069272\n",
      "-------------------------------------------\n",
      "Entity-784x100: progress              : 40%\n",
      "Entity-784x100: err_rms              : 0.016202401468158593\n",
      "Entity-784x100: l2_norm              : 18.4855822431084\n",
      "-------------------------------------------\n",
      "Entity-784x100: progress              : 50%\n",
      "Entity-784x100: err_rms              : 0.01620215643869973\n",
      "Entity-784x100: l2_norm              : 18.486069147676194\n",
      "-------------------------------------------\n",
      "Entity-784x100: progress              : 60%\n",
      "Entity-784x100: err_rms              : 0.016202034096700593\n",
      "Entity-784x100: l2_norm              : 18.486323079421453\n",
      "-------------------------------------------\n",
      "Entity-784x100: progress              : 70%\n",
      "Entity-784x100: err_rms              : 0.016201959457167476\n",
      "Entity-784x100: l2_norm              : 18.486481204151325\n",
      "-------------------------------------------\n",
      "Entity-784x100: progress              : 80%\n",
      "Entity-784x100: err_rms              : 0.01623254730225882\n",
      "Entity-784x100: l2_norm              : 18.47969998768271\n",
      "-------------------------------------------\n",
      "Entity-784x100: progress              : 90%\n",
      "Entity-784x100: err_rms              : 0.01623250724701853\n",
      "Entity-784x100: l2_norm              : 18.47978577224968\n",
      "-------------------------------------------\n",
      "Entity-784x100: progress              : 100%\n",
      "Entity-784x100: err_rms              : 0.016232475519207378\n",
      "Entity-784x100: l2_norm              : 18.479853839350724\n",
      "-------------------------------------------\n",
      "Entity-784x100: progress              : 100%\n",
      "Entity-784x100: err_rms_total              : 0.01627253394868903\n",
      "Entity-784x100: l2_norm              : 18.477539124418225\n",
      "results/mnist_test-0-state.npz saved successfully!\n"
     ]
    }
   ],
   "source": [
    "work_dir = \"results\"\n",
    "prj_name = \"mnist_test\"\n",
    "prj_root = \"/home/jens/work/repos/Rbm\"\n",
    "\n",
    "# Create model\n",
    "model = TestModel(prj_name, \"results\")\n",
    "\n",
    "# Init state\n",
    "model.init(0.01)\n",
    "\n",
    "# load state\n",
    "model.load()\n",
    "\n",
    "# Train\n",
    "model.train(train_images)\n",
    "\n",
    "# save state\n",
    "model.save()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 94,
   "id": "d036b106-77a3-4993-bc2a-5fb164b8ec46",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 259888 characters]",
      "text/plain": [
       "<Figure size 3000x3000 with 100 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Plot reconstructions\n",
    "fig, axes = plt.subplots(n_side, n_side, figsize=(30,30))\n",
    "index = 0             \n",
    "for x in range(n_side):\n",
    "    for y in range(n_side):\n",
    "        inp = train_images[image_indices[index]]\n",
    "        recon = model.backward(model.forward(inp))\n",
    "        recon -= np.min(recon)\n",
    "        recon = recon / np.max(recon)\n",
    "        img = np.reshape(recon, (28, 28))\n",
    "        axes[x,y].imshow(np.asnumpy(img))\n",
    "        axes[x,y].axis('off')\n",
    "        index += 1\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 95,
   "id": "9adce292-cae6-455c-9409-c57ae71a36d5",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(28, 28, 100)\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 281284 characters]",
      "text/plain": [
       "<Figure size 3000x3000 with 100 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Plot weights\n",
    "weights = model.unit1.state.w_hv\n",
    "n, n_hid = weights.shape\n",
    "weight_indices = np.random.randint(0, N-1, n_hid)\n",
    "w = np.reshape(weights, (28, 28, n_hid))\n",
    "print(w.shape)\n",
    "fig, axes = plt.subplots(n_side, n_side, figsize=(30,30))\n",
    "\n",
    "index = 0\n",
    "for x in range(n_side):\n",
    "    for y in range(n_side):\n",
    "        inp = np.asnumpy(w[:,:, weight_indices[index]])\n",
    "        axes[x,y].imshow(inp)\n",
    "        axes[x,y].axis('off')\n",
    "        index += 1\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2a44660c-6ce7-4363-b345-5745f9b6f2d0",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c2beab97-44c5-4a5c-bf4e-56f01753a4ee",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
