Provide examples to Github and checkpoints to get help from edge impulse forums

.gitattributes

@@ -32,3 +32,7 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+ckpt_e09.pth.tar filter=lfs diff=lfs merge=lfs -text
+ckpt_e10.pth.tar filter=lfs diff=lfs merge=lfs -text
+ckpt_e49.pth.tar filter=lfs diff=lfs merge=lfs -text
+ckpt_pytorch_1_11_e00.pth.tar filter=lfs diff=lfs merge=lfs -text

Test_ONNX_Convert.ipynb

@@ -0,0 +1,259 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "id": "d0cd0fce",
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2023-05-15T16:14:15.749744Z",
+     "start_time": "2023-05-15T16:14:15.540642Z"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "import onnx\n",
+    "\n",
+    "\n",
+    "onnx_model = onnx.load(\"ckpt/model.onnx\")\n",
+    "onnx.checker.check_model(onnx_model)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "id": "27f06a8c",
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2023-05-15T16:14:15.751689Z",
+     "start_time": "2023-05-15T16:14:15.748975Z"
+    }
+   },
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 10,
+   "id": "f9167299",
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2023-05-15T16:14:15.777873Z",
+     "start_time": "2023-05-15T16:14:15.753825Z"
+    }
+   },
+   "outputs": [
+    {
+     "ename": "ModuleNotFoundError",
+     "evalue": "No module named 'onnx_tf'",
+     "output_type": "error",
+     "traceback": [
+      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
+      "\u001b[0;31mModuleNotFoundError\u001b[0m                       Traceback (most recent call last)",
+      "Cell \u001b[0;32mIn[10], line 8\u001b[0m\n\u001b[1;32m      6\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mtorch\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mautograd\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m Variable\n\u001b[1;32m      7\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01monnx\u001b[39;00m\n\u001b[0;32m----> 8\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01monnx_tf\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mbackend\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m prepare\n\u001b[1;32m     10\u001b[0m model \u001b[38;5;241m=\u001b[39m onnx\u001b[38;5;241m.\u001b[39mload(\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mckpt/model_5x.onnx\u001b[39m\u001b[38;5;124m'\u001b[39m)\n\u001b[1;32m     11\u001b[0m tf_rep \u001b[38;5;241m=\u001b[39m prepare(model)\n",
+      "\u001b[0;31mModuleNotFoundError\u001b[0m: No module named 'onnx_tf'"
+     ]
+    }
+   ],
+   "source": [
+    "import torch\n",
+    "import torch.nn as nn\n",
+    "import torch.nn.functional as F\n",
+    "import torch.optim as optim\n",
+    "from torchvision import datasets, transforms\n",
+    "from torch.autograd import Variable\n",
+    "import onnx\n",
+    "from onnx_tf.backend import prepare\n",
+    "\n",
+    "model = onnx.load('ckpt/model_5x.onnx')\n",
+    "tf_rep = prepare(model)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "2d1db936",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "f4eb7d23",
+   "metadata": {
+    "scrolled": false
+   },
+   "outputs": [],
+   "source": [
+    "import torch\n",
+    "import torch.onnx as torch.onnx\n",
+    "import onnx\n",
+    "\n",
+    "\n",
+    "import torch.nn as nn\n",
+    "\n",
+    "from models.model import STBVMM\n",
+    "\n",
+    "# # Initialize model with checkpointing enabled\n",
+    "# model = STBVMM(img_size=384, patch_size=1, in_chans=3,\n",
+    "#                  embed_dim=48, depths=[6, 6, 6, 6], num_heads=[6, 6, 6, 6],\n",
+    "#                  window_size=8, mlp_ratio=2., qkv_bias=True, qk_scale=None,\n",
+    "#                  drop_rate=0., attn_drop_rate=0., drop_path_rate=0.1,\n",
+    "#                  norm_layer=nn.LayerNorm, ape=False, patch_norm=True,\n",
+    "#                  use_checkpoint=True, img_range=1., resi_connection='1conv',\n",
+    "#                  manipulator_num_resblk = 1)\n",
+    "\n",
+    "model = STBVMM(img_size=384, patch_size=1, in_chans=3,\n",
+    "                embed_dim=192, depths=[6, 6, 6, 6, 6, 6], num_heads=[6, 6, 6, 6, 6, 6],\n",
+    "                window_size=8, mlp_ratio=2., qkv_bias=True, qk_scale=None,\n",
+    "                drop_rate=0., attn_drop_rate=0., drop_path_rate=0.1,\n",
+    "                norm_layer=nn.LayerNorm, ape=False, patch_norm=True,\n",
+    "                use_checkpoint=False, img_range=1., resi_connection='1conv',\n",
+    "                manipulator_num_resblk=1)\n",
+    "\n",
+    "# Load pretrained weights from checkpoint\n",
+    "checkpoint = torch.load('ckpt/ckpt_e49.pth.tar')\n",
+    "# print(checkpoint.keys())\n",
+    "\n",
+    "# print(checkpoint['state_dict'])\n",
+    "\n",
+    "model.load_state_dict(checkpoint['state_dict'], strict= False)\n",
+    "\n",
+    "# Set the model to eval mode\n",
+    "model.eval()\n",
+    "\n",
+    "# Export model to ONNX\n",
+    "inputs = (torch.randn(1, 3, 384, 384), torch.randn(1, 3, 384, 384), 5)\n",
+    "input_names = [\"a\", \"b\", \"amp\"]\n",
+    "output_names = [\"output\"]\n",
+    "dynamic_axes = {\"a\": {0: \"batch_size\", 2: \"height\", 3: \"width\"},\n",
+    "                \"b\": {0: \"batch_size\", 2: \"height\", 3: \"width\"},\n",
+    "                \"output\": {0: \"batch_size\", 2: \"height\", 3: \"width\"}}\n",
+    "onnx.export(model, inputs, \"model_checkpoint_5x_50ep.onnx\", input_names=input_names, output_names=output_names,\n",
+    "            dynamic_axes=dynamic_axes, opset_version=11)\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "id": "8d70c0e9",
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2023-05-16T04:50:57.341575Z",
+     "start_time": "2023-05-16T04:50:57.144003Z"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "import onnx\n",
+    "\n",
+    "onnx_model = onnx.load(\"ckpt/model_checkpoint_5x_50ep.onnx\")\n",
+    "onnx.checker.check_model(onnx_model)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "ec9bacfd",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import onnxruntime as ort\n",
+    "import numpy as np\n",
+    "import cv2\n",
+    "x, y = test_data[0][0], test_data[0][1]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "id": "1bea7afb",
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2023-05-16T05:23:43.230764Z",
+     "start_time": "2023-05-16T05:23:29.950857Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Using device: cpu\r\n",
+      "demo_video/STB-VMM_Freezer_x20_mag\r\n",
+      "processing sample: 0\r\n",
+      "Traceback (most recent call last):\r\n",
+      "  File \"/Users/raoulritter/STB-VMM/onnxrun.py\", line 112, in <module>\r\n",
+      "    main(args)\r\n",
+      "  File \"/Users/raoulritter/STB-VMM/onnxrun.py\", line 53, in main\r\n",
+      "    ort_outs = ort_session.run(None, ort_inputs)\r\n",
+      "  File \"/opt/anaconda3/envs/afstudeer/lib/python3.10/site-packages/onnxruntime/capi/onnxruntime_inference_collection.py\", line 200, in run\r\n",
+      "    return self._sess.run(output_names, input_feed, run_options)\r\n",
+      "RuntimeError: Input must be a list of dictionaries or a single numpy array for input 'a'.\r\n"
+     ]
+    }
+   ],
+   "source": [
+    "!python onnxrun.py -j4 -b1 --load_ckpt ckpt/model_checkpoint_5x_50ep.onnx --save_dir demo_video/STB-VMM_Freezer_x20_mag -m 5 --video_path demo_video/STB-VMM_Freezer_x20_original/frame --num_data 6644 --mode static\n",
+    "\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "c77312e7",
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2023-05-16T05:14:51.342430Z",
+     "start_time": "2023-05-16T05:14:48.496955Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "1.11.0\n",
+      "0.12.0\n"
+     ]
+    }
+   ],
+   "source": [
+    "import torch\n",
+    "import torchvision\n",
+    "print(torch.__version__)\n",
+    "print(torchvision.__version__)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "702d9d85",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python [conda env:afstudeer]",
+   "language": "python",
+   "name": "conda-env-afstudeer-py"
+  },
+  "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.10.11"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

ckpt_e09.pth.tar

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4b503d280322ad5a257fd760878447c3e13e80800baa06b8288ee37bd79173ce
+size 149374251

ckpt_e10.pth.tar

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b63b6fdfbd487d5482c0e4b821040df85e315407a2205e755b426cf9c94492ce
+size 149368279

ckpt_e49.pth.tar

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1f1df7bebba895be14728293138812826c1affeb4777f76be960e8eb100ed362
+size 149368983

ckpt_pytorch_1_11_e00.pth.tar

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:449bb57e7b0c3a17580f5512cde77397abf6178ea30f28a726d700eac2343920
+size 149368087

model.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:de4e1e8b51f1cf371159c53a9efdc93cdd879c8f4406941e20301d05d3718c67
+size 137258146

model.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:32c2cea3c5ef96e308f5c8b8a0b6418d8e00cacf1c6c5b3e388e796f00ccb079
+size 149306703

model_2.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e72d64cec630d95ba0bc4e49ee7c7e2f1a4a71dcef83ab3b2b1e24ab75fd7a9c
+size 136977868

model_5x.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:de4e1e8b51f1cf371159c53a9efdc93cdd879c8f4406941e20301d05d3718c67
+size 137258146

model_checkpoint_5x_50ep.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7681377d1db055c8b4f4a345052830beb20a027526235fcc6c4c2232edb876bb
+size 136977913

onnxrun.py

@@ -0,0 +1,112 @@
+
+import argparse
+import os
+import numpy as np
+import torch
+import torch.utils.data as data
+from PIL import Image
+from utils.data_loader import ImageFromFolderTest
+import onnxruntime as ort
+
+
+def main(args):
+    # Device choice (auto)
+    if args.device == 'auto':
+        device = 'cuda' if torch.cuda.is_available() else 'cpu'
+    else:
+        device = args.device
+
+    print(f'Using device: {device}')
+
+    # Create ONNX Inference Session
+    ort_session = ort.InferenceSession(args.load_ckpt)
+
+    # Check saving directory
+    save_dir = args.save_dir
+    if not os.path.exists(save_dir):
+        os.makedirs(save_dir)
+    print(save_dir)
+
+    # Data loader
+    dataset_mag = ImageFromFolderTest(
+        args.video_path, mag=args.mag, mode=args.mode, num_data=args.num_data, preprocessing=False)
+    data_loader = data.DataLoader(dataset_mag,
+                                  batch_size=args.batch_size,
+                                  shuffle=False,
+                                  num_workers=args.workers,
+                                  pin_memory=False)
+
+    # Magnification
+    for i, (xa, xb, mag_factor) in enumerate(data_loader):
+        if i % args.print_freq == 0:
+            print('processing sample: %d' % i)
+
+        xa = xa.to(device)
+        xb = xb.to(device)
+
+        # Infer using ONNX model
+        mag_factor = torch.tensor([[args.mag]]).to(device)  # Create a constant tensor for the magnification factor
+        ort_inputs = {ort_session.get_inputs()[0].name: xa,
+                      ort_session.get_inputs()[1].name: xb,
+                      ort_session.get_inputs()[2].name: mag_factor}
+        #y_hat, _, _, _ = ort_session.run(ort_inputs)
+        ort_outs = ort_session.run(None, ort_inputs)
+        y_hat = ort_outs[0]
+
+        # ort_inputs = {ort_session.get_inputs()[0].name: xa,
+        #               ort_session.get_inputs()[1].name: xb}
+        # y_hat, _, _, _ = ort_session.run(None, ort_inputs)
+
+        if i == 0:
+            # Back to image scale (0-255)
+            tmp = xa.permute(0, 2, 3, 1).cpu().detach().numpy()
+            tmp = np.clip(tmp, -1.0, 1.0)
+            tmp = ((tmp + 1.0) * 127.5).astype(np.uint8)
+
+            # Save first frame
+            fn = os.path.join(save_dir, 'STBVMM_%s_%06d.png' % (args.mode, i))
+            im = Image.fromarray(np.concatenate(tmp, 0))
+            im.save(fn)
+
+        # back to image scale (0-255)
+        y_hat = y_hat.permute(0, 2, 3, 1).cpu().detach().numpy()
+        y_hat = np.clip(y_hat, -1.0, 1.0)
+        y_hat = ((y_hat + 1.0) * 127.5).astype(np.uint8)
+
+        # Save frames
+        fn = os.path.join(save_dir, 'STBVMM_%s_%06d.png' % (args.mode, i+1))
+        im = Image.fromarray(np.concatenate(y_hat, 0))
+        im.save(fn)
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(
+        description='Swin Transformer Based Video Motion Magnification')
+
+    # Application parameters
+    parser.add_argument('-i', '--video_path', type=str, metavar='PATH', required=True,
+                        help='path to video input frames')
+    parser.add_argument('-c', '--load_ckpt', type=str, metavar='PATH', required=True,
+                        help='path to load ONNX model')
+    parser.add_argument('-o', '--save_dir', default='demo', type=str, metavar='PATH',
+                        help='path to save generated frames (default: demo)')
+    parser.add_argument('-m', '--mag', metavar='N', default=20.0, type=float,
+                        help='magnification factor (default: 20.0)')
+    parser.add_argument('--mode', default='static', type=str, choices=['static', 'dynamic'],
+                        help='magnification mode (static, dynamic)')
+    parser.add_argument('-n', '--num_data', type=int, metavar='N', required=True,
+                        help='number of frames')
+
+    # Execute parameters
+    parser.add_argument('-j', '--workers', default=16, type=int, metavar='N',
+                        help='number of data loading workers (default: 16)')
+    parser.add_argument('-b', '--batch_size', default=1, type=int,
+                        metavar='N', help='batch size (default: 1)')
+    parser.add_argument('-p', '--print_freq', default=100, type=int,
+                        metavar='N', help='print frequency (default: 100)')
+
+    # Device
+    parser.add_argument('--device', type=str, default='auto',
+                        choices=['auto', 'cpu', 'cuda', 'mps', 'xla'],
+                        help='select device [auto/cpu/cuda] (default: auto)')
+    args = parser.parse_args()
+    main(args)

requirements.txt

@@ -0,0 +1,5 @@
+torch==2.0
+Pillow==9.3
+torchvision
+torchaudio
+numpy

run.py

@@ -0,0 +1,133 @@
+import argparse
+import os
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.utils.data as data
+import torchvision.datasets as datasets
+from PIL import Image
+
+from utils.data_loader import ImageFromFolderTest
+from models.model import STBVMM
+
+
+def main(args):
+    # Device choice (auto)
+    if args.device == 'auto':
+        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    else:
+        device = args.device
+
+    print(f'Using device: {device}')
+
+    # Create model
+    model = STBVMM(img_size=384, patch_size=1, in_chans=3,
+                   embed_dim=192, depths=[6, 6, 6, 6, 6, 6], num_heads=[6, 6, 6, 6, 6, 6],
+                   window_size=8, mlp_ratio=2., qkv_bias=True, qk_scale=None,
+                   drop_rate=0., attn_drop_rate=0., drop_path_rate=0.1,
+                   norm_layer=nn.LayerNorm, ape=False, patch_norm=True,
+                   use_checkpoint=False, img_range=1., resi_connection='1conv',
+                   manipulator_num_resblk=1).to(device)
+
+    # Load checkpoint
+    if os.path.isfile(args.load_ckpt):
+        print("=> loading checkpoint '{}'".format(args.load_ckpt))
+        checkpoint = torch.load(args.load_ckpt)
+        args.start_epoch = checkpoint['epoch']
+
+        model.load_state_dict(checkpoint['state_dict'])
+
+        print("=> loaded checkpoint '{}' (epoch {})"
+              .format(args.load_ckpt, checkpoint['epoch']))
+    else:
+        print("=> no checkpoint found at '{}'".format(args.load_ckpt))
+        assert (False)
+
+    # Check saving directory
+    save_dir = args.save_dir
+    if not os.path.exists(save_dir):
+        os.makedirs(save_dir)
+    print(save_dir)
+
+    # Data loader
+    dataset_mag = ImageFromFolderTest(
+        args.video_path, mag=args.mag, mode=args.mode, num_data=args.num_data, preprocessing=False)
+    data_loader = data.DataLoader(dataset_mag,
+                                  batch_size=args.batch_size,
+                                  shuffle=False,
+                                  num_workers=args.workers,
+                                  pin_memory=False)
+
+    # Generate frames
+    model.eval()
+
+    # Magnification
+    for i, (xa, xb, mag_factor) in enumerate(data_loader):
+        if i % args.print_freq == 0:
+            print('processing sample: %d' % i)
+
+        mag_factor = mag_factor.unsqueeze(1).unsqueeze(1).unsqueeze(1)
+
+        xa = xa.to(device)
+        xb = xb.to(device)
+        mag_factor = mag_factor.to(device)
+
+        y_hat, _, _, _ = model(xa, xb, mag_factor)
+
+        if i == 0:
+            # Back to image scale (0-255)
+            tmp = xa.permute(0, 2, 3, 1).cpu().detach().numpy()
+            tmp = np.clip(tmp, -1.0, 1.0)
+            tmp = ((tmp + 1.0) * 127.5).astype(np.uint8)
+
+            # Save first frame
+            fn = os.path.join(save_dir, 'STBVMM_%s_%06d.png' % (args.mode, i))
+            im = Image.fromarray(np.concatenate(tmp, 0))
+            im.save(fn)
+
+        # back to image scale (0-255)
+        y_hat = y_hat.permute(0, 2, 3, 1).cpu().detach().numpy()
+        y_hat = np.clip(y_hat, -1.0, 1.0)
+        y_hat = ((y_hat + 1.0) * 127.5).astype(np.uint8)
+
+        # Save frames
+        fn = os.path.join(save_dir, 'STBVMM_%s_%06d.png' % (args.mode, i+1))
+        im = Image.fromarray(np.concatenate(y_hat, 0))
+        im.save(fn)
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(
+        description='Swin Transformer Based Video Motion Magnification')
+
+    # Application parameters
+    parser.add_argument('-i', '--video_path', type=str, metavar='PATH', required=True,
+                        help='path to video input frames')
+    parser.add_argument('-c', '--load_ckpt', type=str, metavar='PATH', required=True,
+                        help='path to load checkpoint')
+    parser.add_argument('-o', '--save_dir', default='demo', type=str, metavar='PATH',
+                        help='path to save generated frames (default: demo)')
+    parser.add_argument('-m', '--mag', metavar='N', default=20.0, type=float,
+                        help='magnification factor (default: 20.0)')
+    parser.add_argument('--mode', default='static', type=str, choices=['static', 'dynamic'],
+                        help='magnification mode (static, dynamic)')
+    parser.add_argument('-n', '--num_data', type=int, metavar='N', required=True,
+                        help='number of frames')
+
+    # Execute parameters
+    parser.add_argument('-j', '--workers', default=16, type=int, metavar='N',
+                        help='number of data loading workers (default: 16)')
+    parser.add_argument('-b', '--batch_size', default=1, type=int,
+                        metavar='N', help='batch size (default: 1)')
+    parser.add_argument('-p', '--print_freq', default=100, type=int,
+                        metavar='N', help='print frequency (default: 100)')
+
+    # Device
+    parser.add_argument('--device', type=str, default='auto',
+                        choices=['auto', 'cpu', 'cuda', 'mps', 'xla'],
+                        help='select device [auto/cpu/cuda] (default: auto)')
+
+    args = parser.parse_args()
+
+    main(args)

test_convert.py

@@ -0,0 +1,36 @@
+import torch
+
+from models.model import STBVMM
+
+
+model = STBVMM(img_size=384, patch_size=1, in_chans=3,
+                embed_dim=192, depths=[6, 6, 6, 6, 6, 6], num_heads=[6, 6, 6, 6, 6, 6],
+                window_size=8, mlp_ratio=2., qkv_bias=True, qk_scale=None,
+                drop_rate=0., attn_drop_rate=0., drop_path_rate=0.1,
+                norm_layer=nn.LayerNorm, ape=False, patch_norm=True,
+                use_checkpoint=False, img_range=1., resi_connection='1conv',
+                manipulator_num_resblk=1).to("cpu")
+
+
+
+checkpoint = torch.load('ckpt/ckpt_e10.pth.tar')
+# print(checkpoint.keys())
+
+print(checkpoint['state_dict'])
+
+model.load_state_dict(checkpoint['state_dict'], strict= False)
+# Get the keys in the checkpoint's state_dict
+checkpoint_keys = set(checkpoint['state_dict'].keys())
+
+# Get the keys in the current model's state_dict
+model_keys = set(model.state_dict().keys())
+
+# Find the difference between the keys
+keys_only_in_checkpoint = checkpoint_keys - model_keys
+keys_only_in_model = model_keys - checkpoint_keys
+
+# Print the results
+print("Keys only in the checkpoint's state_dict:")
+print(keys_only_in_checkpoint)
+
+