{
 "cells": [
  {
   "cell_type": "code",
   "id": "d6cb560a-707c-4c1e-b50f-2b9c8037a3c7",
   "metadata": {
    "editable": true,
    "execution": {
     "iopub.execute_input": "2026-06-01T18:05:07.380499Z",
     "iopub.status.busy": "2026-06-01T18:05:07.380324Z",
     "iopub.status.idle": "2026-06-01T18:05:11.030587Z",
     "shell.execute_reply": "2026-06-01T18:05:11.029614Z"
    },
    "slideshow": {
     "slide_type": ""
    },
    "tags": [],
    "ExecuteTime": {
     "end_time": "2026-06-02T05:38:03.403264644Z",
     "start_time": "2026-06-02T05:37:58.800587439Z"
    }
   },
   "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",
    "from image.sub_image import SubImage, normalize\n",
    "import math\n",
    "import random"
   ],
   "outputs": [],
   "execution_count": 1
  },
  {
   "cell_type": "code",
   "id": "4065302d-862f-4639-99d7-e396d055434c",
   "metadata": {
    "editable": true,
    "execution": {
     "iopub.execute_input": "2026-06-01T18:05:11.032871Z",
     "iopub.status.busy": "2026-06-01T18:05:11.032546Z",
     "iopub.status.idle": "2026-06-01T18:05:11.036109Z",
     "shell.execute_reply": "2026-06-01T18:05:11.035245Z"
    },
    "slideshow": {
     "slide_type": ""
    },
    "tags": [],
    "ExecuteTime": {
     "end_time": "2026-06-02T05:38:03.455631319Z",
     "start_time": "2026-06-02T05:38:03.406111606Z"
    }
   },
   "source": [
    "transform = transforms.Compose([\n",
    "    transforms.ToTensor(),\n",
    "    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))                            \n",
    "])                              "
   ],
   "outputs": [],
   "execution_count": 2
  },
  {
   "cell_type": "code",
   "id": "ffe48fc6-0f57-48a0-b5e9-aa68cbc8edf1",
   "metadata": {
    "editable": true,
    "execution": {
     "iopub.execute_input": "2026-06-01T18:05:11.037686Z",
     "iopub.status.busy": "2026-06-01T18:05:11.037526Z",
     "iopub.status.idle": "2026-06-01T18:05:12.466249Z",
     "shell.execute_reply": "2026-06-01T18:05:12.465235Z"
    },
    "slideshow": {
     "slide_type": ""
    },
    "tags": [],
    "ExecuteTime": {
     "end_time": "2026-06-02T05:38:05.021664189Z",
     "start_time": "2026-06-02T05:38:03.458100201Z"
    }
   },
   "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)"
   ],
   "outputs": [],
   "execution_count": 3
  },
  {
   "cell_type": "code",
   "id": "33501c0f-6d78-43e9-b2ce-adbe40b6d95a",
   "metadata": {
    "editable": true,
    "execution": {
     "iopub.execute_input": "2026-06-01T18:05:12.468794Z",
     "iopub.status.busy": "2026-06-01T18:05:12.468559Z",
     "iopub.status.idle": "2026-06-01T18:05:12.472474Z",
     "shell.execute_reply": "2026-06-01T18:05:12.471545Z"
    },
    "slideshow": {
     "slide_type": ""
    },
    "tags": [],
    "ExecuteTime": {
     "end_time": "2026-06-02T05:38:05.078679479Z",
     "start_time": "2026-06-02T05:38:05.023219613Z"
    }
   },
   "source": [
    "print(train_data)"
   ],
   "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"
     ]
    }
   ],
   "execution_count": 4
  },
  {
   "cell_type": "code",
   "id": "aedf601a-9084-4047-98e8-481f4d41a2d9",
   "metadata": {
    "editable": true,
    "execution": {
     "iopub.execute_input": "2026-06-01T18:05:12.474580Z",
     "iopub.status.busy": "2026-06-01T18:05:12.474238Z",
     "iopub.status.idle": "2026-06-01T18:05:12.748650Z",
     "shell.execute_reply": "2026-06-01T18:05:12.747714Z"
    },
    "slideshow": {
     "slide_type": ""
    },
    "tags": [],
    "ExecuteTime": {
     "end_time": "2026-06-02T05:38:05.423423054Z",
     "start_time": "2026-06-02T05:38:05.089269098Z"
    }
   },
   "source": [
    "N = 500\n",
    "n_side = 10\n",
    "PATCH  = 8\n",
    "STRIDE = 1\n",
    "N_VIS  = 3 * PATCH * PATCH\n",
    "\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",
    "print(image.shape)"
   ],
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([3, 32, 32])\n"
     ]
    }
   ],
   "execution_count": 5
  },
  {
   "cell_type": "code",
   "id": "5ee865c3-2af4-487f-a184-1291f99839b1",
   "metadata": {
    "editable": true,
    "execution": {
     "iopub.execute_input": "2026-06-01T18:05:12.750649Z",
     "iopub.status.busy": "2026-06-01T18:05:12.750448Z",
     "iopub.status.idle": "2026-06-01T18:05:14.002882Z",
     "shell.execute_reply": "2026-06-01T18:05:14.002007Z"
    },
    "slideshow": {
     "slide_type": ""
    },
    "tags": [],
    "ExecuteTime": {
     "end_time": "2026-06-02T05:38:31.997395744Z",
     "start_time": "2026-06-02T05:38:05.427260889Z"
    }
   },
   "source": [
    "image_indices = np.array(random.sample(range(min(N, n_side*n_side)), n_side*n_side))\n",
    "print(f\"train_images.shape:{train_images.shape}\")\n",
    "mean_training_batch = np.reshape(np.repeat(np.mean(train_images, axis=1), 3072, axis=0), train_images.shape)\n",
    "var_training_batch  = np.reshape(np.repeat(np.std(train_images,  axis=1), 3072, axis=0), train_images.shape)\n",
    "train_images_norm = (train_images - mean_training_batch) / var_training_batch\n",
    "sub_generator = SubImage(PATCH, PATCH, STRIDE, STRIDE)\n",
    "train_subimages = sub_generator(np.reshape(train_images, (N, 3, 32, 32)), 3)\n",
    "print(train_subimages.shape)\n",
    "train_subimages = np.reshape(train_subimages, (train_subimages.shape[0], N_VIS))\n",
    "print(train_subimages.shape)"
   ],
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "train_images.shape:(500, 3072)\n",
      "(312500, 3, 8, 8)\n",
      "(312500, 192)\n"
     ]
    }
   ],
   "execution_count": 6
  },
  {
   "cell_type": "code",
   "id": "08143a53-12a1-4896-a4b2-166dedf0e021",
   "metadata": {
    "editable": true,
    "execution": {
     "iopub.execute_input": "2026-06-01T18:05:15.452497Z",
     "iopub.status.busy": "2026-06-01T18:05:15.452325Z",
     "iopub.status.idle": "2026-06-01T18:05:15.456017Z",
     "shell.execute_reply": "2026-06-01T18:05:15.455418Z"
    },
    "slideshow": {
     "slide_type": ""
    },
    "tags": [],
    "ExecuteTime": {
     "end_time": "2026-06-02T05:38:32.049085581Z",
     "start_time": "2026-06-02T05:38:31.999246521Z"
    }
   },
   "source": [
    "class TestModel(Model):\n",
    "    def __init__(self, name: str, work_dir: str = '.'):\n",
    "        super().__init__(name, work_dir)\n",
    "        self.unit1 = Entity((N_VIS, 128), 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.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"
   ],
   "outputs": [],
   "execution_count": 7
  },
  {
   "cell_type": "code",
   "id": "6b3be9f6-38e1-442d-a7bd-1f3c164ec8d4",
   "metadata": {
    "editable": true,
    "execution": {
     "iopub.execute_input": "2026-06-01T18:05:15.457849Z",
     "iopub.status.busy": "2026-06-01T18:05:15.457672Z",
     "iopub.status.idle": "2026-06-01T18:05:40.787667Z",
     "shell.execute_reply": "2026-06-01T18:05:40.786713Z"
    },
    "slideshow": {
     "slide_type": ""
    },
    "tags": [],
    "ExecuteTime": {
     "end_time": "2026-06-02T05:47:23.451633187Z",
     "start_time": "2026-06-02T05:38:32.052539866Z"
    }
   },
   "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(normalize(train_subimages))\n",
    "model.train(train_subimages)\n",
    "\n",
    "# save state\n",
    "model.save()"
   ],
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "-------------------------------------------\n",
      "Entity-192x128: progress              : 0%\n",
      "Entity-192x128: err_rms              : 0.02982608561419449\n",
      "Entity-192x128: l2_norm              : 0.5154587953841656\n",
      "-------------------------------------------\n",
      "Entity-192x128: progress              : 10%\n",
      "Entity-192x128: err_rms              : 0.028228882569056157\n",
      "Entity-192x128: l2_norm              : 0.5495957617469555\n",
      "-------------------------------------------\n",
      "Entity-192x128: progress              : 20%\n",
      "Entity-192x128: err_rms              : 0.027550804062706425\n",
      "Entity-192x128: l2_norm              : 0.5878319340382664\n",
      "-------------------------------------------\n",
      "Entity-192x128: progress              : 30%\n",
      "Entity-192x128: err_rms              : 0.02740561917531193\n",
      "Entity-192x128: l2_norm              : 0.5976701067086865\n",
      "-------------------------------------------\n",
      "Entity-192x128: progress              : 40%\n",
      "Entity-192x128: err_rms              : 0.027248082195068083\n",
      "Entity-192x128: l2_norm              : 0.5991757429205762\n",
      "-------------------------------------------\n",
      "Entity-192x128: progress              : 50%\n",
      "Entity-192x128: err_rms              : 0.026824203916998864\n",
      "Entity-192x128: l2_norm              : 0.6027710231695954\n",
      "-------------------------------------------\n",
      "Entity-192x128: progress              : 60%\n",
      "Entity-192x128: err_rms              : 0.02655947728844083\n",
      "Entity-192x128: l2_norm              : 0.6041359064549188\n",
      "-------------------------------------------\n",
      "Entity-192x128: progress              : 70%\n",
      "Entity-192x128: err_rms              : 0.026933387882988033\n",
      "Entity-192x128: l2_norm              : 0.603607961995845\n",
      "-------------------------------------------\n",
      "Entity-192x128: progress              : 80%\n",
      "Entity-192x128: err_rms              : 0.02728050450716867\n",
      "Entity-192x128: l2_norm              : 0.6036734185897112\n",
      "-------------------------------------------\n",
      "Entity-192x128: progress              : 90%\n",
      "Entity-192x128: err_rms              : 0.027154601608006385\n",
      "Entity-192x128: l2_norm              : 0.6046113788661706\n",
      "-------------------------------------------\n",
      "Entity-192x128: progress              : 100%\n",
      "Entity-192x128: err_rms              : 0.027059586883727227\n",
      "Entity-192x128: l2_norm              : 0.6035971257113926\n",
      "-------------------------------------------\n",
      "Entity-192x128: progress              : 100%\n",
      "Entity-192x128: err_rms_total              : 0.02683939309309671\n",
      "Entity-192x128: l2_norm              : 0.604365570485915\n"
     ]
    }
   ],
   "execution_count": 8
  },
  {
   "cell_type": "code",
   "id": "41430b5f",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-06-01T18:05:42.169881Z",
     "iopub.status.busy": "2026-06-01T18:05:42.169722Z",
     "iopub.status.idle": "2026-06-01T18:05:42.691947Z",
     "shell.execute_reply": "2026-06-01T18:05:42.691290Z"
    },
    "ExecuteTime": {
     "end_time": "2026-06-02T05:47:26.463266352Z",
     "start_time": "2026-06-02T05:47:23.454053601Z"
    }
   },
   "source": [
    "import numpy\n",
    "\n",
    "# Reconstruct full 32x32 images by assembling reconstructed sub-images (overlap-averaged)\n",
    "nx_steps = (32 - PATCH) // STRIDE + 1\n",
    "ny_steps = (32 - PATCH) // STRIDE + 1\n",
    "n_images = n_side\n",
    "\n",
    "fig, axes = plt.subplots(n_images, 2, figsize=(6, n_images * 3))\n",
    "\n",
    "for img_idx in range(n_images):\n",
    "    orig = np.reshape(train_images[img_idx], (3, 32, 32))\n",
    "    orig_display = np.asnumpy((orig - np.min(orig)) / (np.max(orig) - np.min(orig)))\n",
    "\n",
    "    accum = numpy.zeros((3, 32, 32), dtype=numpy.float64)\n",
    "    count = numpy.zeros((1, 32, 32), dtype=numpy.float64)\n",
    "    for ix in range(nx_steps):\n",
    "        for iy in range(ny_steps):\n",
    "            sub_idx = img_idx * nx_steps * ny_steps + ix * ny_steps + iy\n",
    "            inp = train_subimages[sub_idx]\n",
    "            recon = model.backward(model.forward(inp))\n",
    "            arr = np.asnumpy(np.reshape(recon, (3, PATCH, PATCH)))\n",
    "            x0, y0 = ix * STRIDE, iy * STRIDE\n",
    "            accum[:, x0:x0+PATCH, y0:y0+PATCH] += arr\n",
    "            count[0, x0:x0+PATCH, y0:y0+PATCH] += 1.0\n",
    "\n",
    "    recon_full = (accum / numpy.maximum(count, 1)).transpose(1, 2, 0)\n",
    "    lo, hi = recon_full.min(), recon_full.max()\n",
    "    recon_full = numpy.clip((recon_full - lo) / (hi - lo + 1e-8), 0, 1)\n",
    "\n",
    "    axes[img_idx, 0].imshow(orig_display.transpose(1, 2, 0))\n",
    "    axes[img_idx, 0].axis('off')\n",
    "    if img_idx == 0:\n",
    "        axes[img_idx, 0].set_title('Original')\n",
    "\n",
    "    axes[img_idx, 1].imshow(recon_full)\n",
    "    axes[img_idx, 1].axis('off')\n",
    "    if img_idx == 0:\n",
    "        axes[img_idx, 1].set_title('Reconstructed')\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<Figure size 600x3000 with 20 Axes>"
      ],
      "image/png": "[encoded data, 135320 characters]"
     },
     "metadata": {},
     "output_type": "display_data",
     "jetTransient": {
      "display_id": null
     }
    }
   ],
   "execution_count": 9
  },
  {
   "cell_type": "code",
   "id": "9adce292-cae6-455c-9409-c57ae71a36d5",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-06-01T18:05:42.694095Z",
     "iopub.status.busy": "2026-06-01T18:05:42.693875Z",
     "iopub.status.idle": "2026-06-01T18:05:44.224681Z",
     "shell.execute_reply": "2026-06-01T18:05:44.223717Z"
    },
    "scrolled": true,
    "ExecuteTime": {
     "end_time": "2026-06-02T05:47:27.823225543Z",
     "start_time": "2026-06-02T05:47:26.465459324Z"
    }
   },
   "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, PATCH, PATCH, 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()"
   ],
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(3, 8, 8, 128)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<Figure size 3000x3000 with 100 Axes>"
      ],
      "image/png": "[encoded data, 68408 characters]"
     },
     "metadata": {},
     "output_type": "display_data",
     "jetTransient": {
      "display_id": null
     }
    }
   ],
   "execution_count": 10
  },
  {
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-06-01T18:05:14.004890Z",
     "iopub.status.busy": "2026-06-01T18:05:14.004617Z",
     "iopub.status.idle": "2026-06-01T18:05:15.450743Z",
     "shell.execute_reply": "2026-06-01T18:05:15.449902Z"
    },
    "ExecuteTime": {
     "end_time": "2026-06-02T05:47:29.498973694Z",
     "start_time": "2026-06-02T05:47:27.824789556Z"
    }
   },
   "cell_type": "code",
   "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_subimages[image_indices[index]]\n",
    "        img = (1+np.reshape(inp, (3, PATCH, PATCH)))/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()"
   ],
   "id": "deddee55057a193c",
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<Figure size 3000x3000 with 100 Axes>"
      ],
      "image/png": "[encoded data, 78000 characters]"
     },
     "metadata": {},
     "output_type": "display_data",
     "jetTransient": {
      "display_id": null
     }
    }
   ],
   "execution_count": 11
  },
  {
   "cell_type": "code",
   "id": "c976fed5-162c-4a70-8219-fc6c68c1fb66",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2026-06-02T05:47:29.546682653Z",
     "start_time": "2026-06-02T05:47:29.500214253Z"
    }
   },
   "source": [],
   "outputs": [],
   "execution_count": 11
  },
  {
   "cell_type": "code",
   "id": "661dfc9d-e58a-4994-a5fa-c377d7f743e5",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2026-06-02T05:47:29.594759213Z",
     "start_time": "2026-06-02T05:47:29.548216752Z"
    }
   },
   "source": [],
   "outputs": [],
   "execution_count": 11
  },
  {
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-06-01T18:05:40.789519Z",
     "iopub.status.busy": "2026-06-01T18:05:40.789280Z",
     "iopub.status.idle": "2026-06-01T18:05:42.168103Z",
     "shell.execute_reply": "2026-06-01T18:05:42.167459Z"
    },
    "ExecuteTime": {
     "end_time": "2026-06-02T05:47:31.761626427Z",
     "start_time": "2026-06-02T05:47:29.596363431Z"
    }
   },
   "cell_type": "code",
   "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_subimages[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, PATCH, PATCH))\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()"
   ],
   "id": "bb342d25fa0ab4a4",
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<Figure size 3000x3000 with 100 Axes>"
      ],
      "image/png": "[encoded data, 79716 characters]"
     },
     "metadata": {},
     "output_type": "display_data",
     "jetTransient": {
      "display_id": null
     }
    }
   ],
   "execution_count": 12
  },
  {
   "cell_type": "code",
   "id": "c1a67046-17e4-4924-8067-da23481eb6d3",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2026-06-02T05:47:31.811677110Z",
     "start_time": "2026-06-02T05:47:31.764068701Z"
    }
   },
   "source": [],
   "outputs": [],
   "execution_count": 12
  }
 ],
 "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
}