{
 "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, sample_gaussian\n",
    "from rbm.torch import Optimizer\n",
    "import math\n",
    "import random"
   ]
  },
  {
   "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",
    "    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))                            \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.CIFAR10(root='./data', train=True, transform=transform, download=True)\n",
    "test_data = torchvision.datasets.CIFAR10(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": "33501c0f-6d78-43e9-b2ce-adbe40b6d95a",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Dataset CIFAR10\n",
      "    Number of datapoints: 50000\n",
      "    Root location: ./data\n",
      "    Split: Train\n",
      "    StandardTransform\n",
      "Transform: Compose(\n",
      "               ToTensor()\n",
      "               Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5))\n",
      "           )\n"
     ]
    }
   ],
   "source": [
    "print(train_data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 110,
   "id": "aedf601a-9084-4047-98e8-481f4d41a2d9",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "N = 5000\n",
    "n_side = 10\n",
    "N_VIS = 3072\n",
    "image, label = train_data[0]\n",
    "size = len(torch.flatten(image))\n",
    "train_images = np.zeros(shape=(N, size))\n",
    "train_images_tensor = torch.Tensor(N, size)\n",
    "for i in range(N):\n",
    "    image, label = train_data[i]\n",
    "    flat_array = torch.flatten(image)\n",
    "    train_images[i, :] = np.asarray(flat_array.numpy())\n",
    "    train_images_tensor[i, :] = flat_array\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 111,
   "id": "5ee865c3-2af4-487f-a184-1291f99839b1",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "len(train_images):5000\n"
     ]
    }
   ],
   "source": [
    "image_indices = np.array(random.sample(range(min(N, n_side*n_side)), n_side*n_side))\n",
    "print(f\"len(train_images):{len(train_images)}\")\n",
    "mean_training_batch = np.reshape(np.repeat(np.mean(train_images, axis=1), N_VIS, axis=0), train_images.shape)\n",
    "var_training_batch = np.reshape(np.repeat(np.std(train_images, axis=1), N_VIS, axis=0), train_images.shape)\n",
    "#print(f\"mean_training_batch:{mean_training_batch}\")\n",
    "#print(f\"var_training_batch:{var_training_batch}\")\n",
    "train_images_norm = (train_images - mean_training_batch) / var_training_batch\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 112,
   "id": "b4310933-d850-4128-a2dd-27d88c004846",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 587304 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 = (1+np.reshape(inp, (3, 32, 32)))/2\n",
    "        img = img.transpose((1,2,0))\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": 113,
   "id": "08143a53-12a1-4896-a4b2-166dedf0e021",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "class TestModel(Model):\n",
    "    def __init__(self, name: str, work_dir: str = '.'):\n",
    "        super().__init__(name, work_dir)\n",
    "        #self.unit1 = Entity((3072, 125), EntityParams(do_gaussian_visible=True, do_gaussian_hidden=True), TrainingParams(learning_rate=0.0001, momentum=0.9, num_epochs=1000, mini_batch_size=0, weight_decay=0.0, l2_lambda=0.0))\n",
    "        self.unit1 = Entity((3072, 400), EntityParams(do_gaussian_visible=True, do_gaussian_hidden=False), TrainingParams(learning_rate=0.005, momentum=0.9, num_epochs=1000, mini_batch_size=1000, weight_decay=0, l1_lambda=0.02))\n",
    "    \n",
    "    def forward(self, x: Mat):\n",
    "        x = self.unit1.forward(x)\n",
    "        return x\n",
    "    \n",
    "    def backward(self, x: Mat):\n",
    "        x = self.unit1.reconstruct(x)\n",
    "        return x\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 114,
   "id": "6b3be9f6-38e1-442d-a7bd-1f3c164ec8d4",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "results/cifar_test-0-state.npz loaded successfully!\n",
      "-------------------------------------------\n",
      "Entity-3072x400: progress              : 0%\n",
      "Entity-3072x400: err_rms              : 0.06896937980428916\n",
      "Entity-3072x400: l2_norm              : 0.5225304597575441\n",
      "-------------------------------------------\n",
      "Entity-3072x400: progress              : 10%\n",
      "Entity-3072x400: err_rms              : 0.06491332057085623\n",
      "Entity-3072x400: l2_norm              : 0.5345670271350904\n",
      "-------------------------------------------\n",
      "Entity-3072x400: progress              : 20%\n",
      "Entity-3072x400: err_rms              : 0.06351822782056743\n",
      "Entity-3072x400: l2_norm              : 0.5493135828521942\n",
      "-------------------------------------------\n",
      "Entity-3072x400: progress              : 30%\n",
      "Entity-3072x400: err_rms              : 0.06272928581494916\n",
      "Entity-3072x400: l2_norm              : 0.5584482708950883\n",
      "-------------------------------------------\n",
      "Entity-3072x400: progress              : 40%\n",
      "Entity-3072x400: err_rms              : 0.06217302887206066\n",
      "Entity-3072x400: l2_norm              : 0.564909103255255\n",
      "-------------------------------------------\n",
      "Entity-3072x400: progress              : 50%\n",
      "Entity-3072x400: err_rms              : 0.061807095817242357\n",
      "Entity-3072x400: l2_norm              : 0.5693171228562531\n",
      "-------------------------------------------\n",
      "Entity-3072x400: progress              : 60%\n",
      "Entity-3072x400: err_rms              : 0.061533665224563217\n",
      "Entity-3072x400: l2_norm              : 0.5725814743457811\n",
      "-------------------------------------------\n",
      "Entity-3072x400: progress              : 70%\n",
      "Entity-3072x400: err_rms              : 0.06131226857443747\n",
      "Entity-3072x400: l2_norm              : 0.5752021306731918\n",
      "-------------------------------------------\n",
      "Entity-3072x400: progress              : 80%\n",
      "Entity-3072x400: err_rms              : 0.061144399319335754\n",
      "Entity-3072x400: l2_norm              : 0.577123812977092\n",
      "-------------------------------------------\n",
      "Entity-3072x400: progress              : 90%\n",
      "Entity-3072x400: err_rms              : 0.06104187402269581\n",
      "Entity-3072x400: l2_norm              : 0.5785267226575117\n",
      "-------------------------------------------\n",
      "Entity-3072x400: progress              : 100%\n",
      "Entity-3072x400: err_rms              : 0.06091846055902186\n",
      "Entity-3072x400: l2_norm              : 0.5797975264092875\n",
      "-------------------------------------------\n",
      "Entity-3072x400: progress              : 100%\n",
      "Entity-3072x400: err_rms_total              : 0.060889724352856944\n",
      "Entity-3072x400: l2_norm              : 0.580201157147108\n",
      "results/cifar_test-0-state.npz saved successfully!\n"
     ]
    }
   ],
   "source": [
    "work_dir = \"results\"\n",
    "prj_name = \"cifar_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",
    "# save state\n",
    "model.save()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 115,
   "id": "d036b106-77a3-4993-bc2a-5fb164b8ec46",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 473912 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, (3, 32, 32))\n",
    "        img = img.transpose((1,2,0))\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": 116,
   "id": "9adce292-cae6-455c-9409-c57ae71a36d5",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(3, 32, 32, 400)\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 520912 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.array(random.sample(range(n_hid), n_side*n_side))\n",
    "w = np.reshape(weights, (3, 32, 32, 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",
    "        inp -= np.min(inp)\n",
    "        inp = inp / np.max(inp)\n",
    "        img = inp.transpose((1,2,0))\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": null,
   "id": "c976fed5-162c-4a70-8219-fc6c68c1fb66",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "661dfc9d-e58a-4994-a5fa-c377d7f743e5",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "df88dced-1e66-4f7d-a744-0f3f4f5dd271",
   "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
}
