Test_ONNX_Convert.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "62ee8afc",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/opt/homebrew/Caskroom/miniconda/base/envs/afstudeer/lib/python3.10/site-packages/torch/functional.py:504: UserWarning: torch.meshgrid: in an upcoming release, it will be required to pass the indexing argument. (Triggered internally at /Users/runner/work/pytorch/pytorch/pytorch/aten/src/ATen/native/TensorShape.cpp:3484.)\n",
      "  return _VF.meshgrid(tensors, **kwargs)  # type: ignore[attr-defined]\n",
      "/Users/raoulritter/STB-VMM/models/modelpnnx20x.py:988: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!\n",
      "  if a.shape != b.shape:\n",
      "/Users/raoulritter/STB-VMM/models/modelpnnx20x.py:371: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!\n",
      "  if self.input_resolution == x_size:\n",
      "/Users/raoulritter/STB-VMM/models/modelpnnx20x.py:169: TracerWarning: Converting a tensor to a Python integer might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!\n",
      "  B = int(windows.shape[0] / (H * W / window_size / window_size))\n"
     ]
    }
   ],
   "source": [
    "#CONVERT MODEL TO TORCH SCRIPT\n",
    "\n",
    "\n",
    "import torch\n",
    "# import torchvision.models as models\n",
    "from models.modelpnnx20x import STBVMM\n",
    "\n",
    "# net = models.resnet18(pretrained=True)\n",
    "\n",
    "a_dummy = torch.randn(1, 3, 384, 384)\n",
    "b_dummy = torch.randn(1, 3, 384, 384)\n",
    "# c_dummy = torch.randn(1, 3, 384, 384)  # optional during training\n",
    "\n",
    "# For the amplitude parameter amp\n",
    "# amp_dummy = torch.tensor(5)  # adjusted to be a scalar value\n",
    "\n",
    "# Instantiate your model\n",
    "model = STBVMM()\n",
    "\n",
    "# Load the weights from the checkpoint\n",
    "checkpoint = torch.load('ckpt/ckpt_e49.pth.tar')\n",
    "model.load_state_dict(checkpoint['state_dict'], strict= False)\n",
    "# Make sure the model is in evaluation mode\n",
    "# model.eval()\n",
    "model = model.eval()\n",
    "\n",
    "# x = torch.rand(1, 3, 224, 224)\n",
    "\n",
    "# You could try disabling checking when tracing raises error\n",
    "# mod = torch.jit.trace(net, x, check_trace=False)\n",
    "mod = torch.jit.trace(model, (a_dummy,b_dummy) )\n",
    "\n",
    "mod.save(\"20x/modelpnnx20x.pt\")\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "4f8337c7",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from models.model_jit_noC import STBVMM\n",
    "\n",
    "# Set up dummy input data\n",
    "a_dummy = torch.randn(1, 3, 384, 384)\n",
    "b_dummy = torch.randn(1, 3, 384, 384)\n",
    "amp_dummy = torch.tensor(20)\n",
    "\n",
    "# Instantiate your model\n",
    "model = STBVMM()\n",
    "\n",
    "# Load the weights from the checkpoint\n",
    "checkpoint = torch.load('ckpt/ckpt_e49.pth.tar')\n",
    "model.load_state_dict(checkpoint['state_dict'], strict= False)\n",
    "\n",
    "# Make sure the model is in evaluation mode\n",
    "model = model.eval()\n",
    "\n",
    "\n",
    "# Trace the model\n",
    "mod = torch.jit.trace(model, (a_dummy, b_dummy, amp_dummy))\n",
    "\n",
    "# Save the scripted model\n",
    "mod.save(\"20x/modelpnnx20x.pt\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "18c63535",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from PIL import Image\n",
    "from torchvision import transforms\n",
    "\n",
    "def main():\n",
    "    # Load the TorchScript model\n",
    "    model = torch.jit.load('20x/modelpnnx20x.pt')\n",
    "\n",
    "    # Load the images from disk\n",
    "    image_a = Image.open('demo_video/STB-VMM_demo_x20_static_original/frame_000007.png')\n",
    "    image_b = Image.open('demo_video/STB-VMM_demo_x20_static_original/frame_000008.png')\n",
    "\n",
    "    # Resize the images if necessary\n",
    "    desired_size = (384, 384)\n",
    "    image_a = image_a.resize(desired_size)\n",
    "    image_b = image_b.resize(desired_size)\n",
    "\n",
    "    # Convert the images to tensors\n",
    "    transform = transforms.ToTensor()\n",
    "    a_tensor = transform(image_a).unsqueeze(0)\n",
    "    b_tensor = transform(image_b).unsqueeze(0)\n",
    "\n",
    "    # Set up amp tensor and unsqueeze to match the shape expected by the model\n",
    "    amp_tensor = torch.tensor([20]).unsqueeze(1).unsqueeze(1).unsqueeze(1)\n",
    "\n",
    "    # Run inference\n",
    "    output_tuple = model(a_tensor, b_tensor, amp_tensor)\n",
    "    output = output_tuple[0]  # Access the desired tensor from the tuple\n",
    "\n",
    "    # Reshape the tensor (remove batch dimension)\n",
    "    output = output.squeeze(0)\n",
    "\n",
    "    # Convert the output tensor to a PIL image\n",
    "    output_image = transforms.ToPILImage()(output.cpu())\n",
    "\n",
    "    # Save the combined image\n",
    "    output_image.save('combined_image.png')\n",
    "\n",
    "\n",
    "if __name__ == '__main__':\n",
    "    main()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "32f9a8d0",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from PIL import Image\n",
    "from torchvision import transforms\n",
    "\n",
    "def main():\n",
    "    # Load the TorchScript model\n",
    "    model = torch.jit.load('20x/modelpnnx20x.pt')\n",
    "\n",
    "    # Load the images from disk\n",
    "    image_a = Image.open('demo_video/STB-VMM_demo_x20_static_original/frame_000007.png')\n",
    "    image_b = Image.open('demo_video/STB-VMM_demo_x20_static_original/frame_000008.png')\n",
    "\n",
    "    # Resize the images if necessary\n",
    "    desired_size = (384, 384)\n",
    "    image_a = image_a.resize(desired_size)\n",
    "    image_b = image_b.resize(desired_size)\n",
    "\n",
    "    # Convert the images to tensors\n",
    "    transform = transforms.ToTensor()\n",
    "    a_tensor = transform(image_a).unsqueeze(0)\n",
    "    b_tensor = transform(image_b).unsqueeze(0)\n",
    "\n",
    "    # Set up amp tensor\n",
    "    amp_tensor = torch.tensor([20])  # Replace 20 with your desired amplification factor\n",
    "\n",
    "    # Run inference\n",
    "    output_tuple = model(a_tensor, b_tensor, amp_tensor)\n",
    "    output = output_tuple[0]  # Access the desired tensor from the tuple\n",
    "\n",
    "    # Reshape the tensor (remove batch dimension)\n",
    "    output = output.squeeze(0)\n",
    "\n",
    "    # Convert the output tensor to a PIL image\n",
    "    output_image = transforms.ToPILImage()(output.cpu())\n",
    "\n",
    "    # Save the combined image\n",
    "    output_image.save('combined_image.png')\n",
    "\n",
    "\n",
    "if __name__ == '__main__':\n",
    "    main()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "414bb27a",
   "metadata": {},
   "outputs": [
    {
     "ename": "RuntimeError",
     "evalue": "version_ <= kMaxSupportedFileFormatVersionINTERNAL ASSERT FAILED at \"/Users/distiller/project/pytorch/caffe2/serialize/inline_container.cc\":147, please report a bug to PyTorch. Attempted to read a PyTorch file with version 10, but the maximum supported version for reading is 9. Your PyTorch installation may be too old.",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mRuntimeError\u001b[0m                              Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[24], line 40\u001b[0m\n\u001b[1;32m     36\u001b[0m     output_image\u001b[39m.\u001b[39msave(\u001b[39m'\u001b[39m\u001b[39mcombined_image.png\u001b[39m\u001b[39m'\u001b[39m)\n\u001b[1;32m     39\u001b[0m \u001b[39mif\u001b[39;00m \u001b[39m__name__\u001b[39m \u001b[39m==\u001b[39m \u001b[39m'\u001b[39m\u001b[39m__main__\u001b[39m\u001b[39m'\u001b[39m:\n\u001b[0;32m---> 40\u001b[0m     main()\n",
      "Cell \u001b[0;32mIn[24], line 9\u001b[0m, in \u001b[0;36mmain\u001b[0;34m()\u001b[0m\n\u001b[1;32m      7\u001b[0m \u001b[39mdef\u001b[39;00m \u001b[39mmain\u001b[39m():\n\u001b[1;32m      8\u001b[0m     \u001b[39m# Load the TorchScript model\u001b[39;00m\n\u001b[0;32m----> 9\u001b[0m     model \u001b[39m=\u001b[39m torch\u001b[39m.\u001b[39;49mjit\u001b[39m.\u001b[39;49mload(\u001b[39m'\u001b[39;49m\u001b[39m20x/modelpnnx20x.pt\u001b[39;49m\u001b[39m'\u001b[39;49m)\n\u001b[1;32m     11\u001b[0m     \u001b[39m# Load the images from disk\u001b[39;00m\n\u001b[1;32m     12\u001b[0m     image_a \u001b[39m=\u001b[39m Image\u001b[39m.\u001b[39mopen(\u001b[39m'\u001b[39m\u001b[39mdemo_video/STB-VMM_demo_x20_static_original/frame_000007.png\u001b[39m\u001b[39m'\u001b[39m)\n",
      "File \u001b[0;32m/opt/homebrew/Caskroom/miniconda/base/envs/testconv/lib/python3.9/site-packages/torch/jit/_serialization.py:162\u001b[0m, in \u001b[0;36mload\u001b[0;34m(f, map_location, _extra_files)\u001b[0m\n\u001b[1;32m    160\u001b[0m cu \u001b[39m=\u001b[39m torch\u001b[39m.\u001b[39m_C\u001b[39m.\u001b[39mCompilationUnit()\n\u001b[1;32m    161\u001b[0m \u001b[39mif\u001b[39;00m \u001b[39misinstance\u001b[39m(f, \u001b[39mstr\u001b[39m) \u001b[39mor\u001b[39;00m \u001b[39misinstance\u001b[39m(f, pathlib\u001b[39m.\u001b[39mPath):\n\u001b[0;32m--> 162\u001b[0m     cpp_module \u001b[39m=\u001b[39m torch\u001b[39m.\u001b[39;49m_C\u001b[39m.\u001b[39;49mimport_ir_module(cu, \u001b[39mstr\u001b[39;49m(f), map_location, _extra_files)\n\u001b[1;32m    163\u001b[0m \u001b[39melse\u001b[39;00m:\n\u001b[1;32m    164\u001b[0m     cpp_module \u001b[39m=\u001b[39m torch\u001b[39m.\u001b[39m_C\u001b[39m.\u001b[39mimport_ir_module_from_buffer(\n\u001b[1;32m    165\u001b[0m         cu, f\u001b[39m.\u001b[39mread(), map_location, _extra_files\n\u001b[1;32m    166\u001b[0m     )\n",
      "\u001b[0;31mRuntimeError\u001b[0m: version_ <= kMaxSupportedFileFormatVersionINTERNAL ASSERT FAILED at \"/Users/distiller/project/pytorch/caffe2/serialize/inline_container.cc\":147, please report a bug to PyTorch. Attempted to read a PyTorch file with version 10, but the maximum supported version for reading is 9. Your PyTorch installation may be too old."
     ]
    }
   ],
   "source": [
    "#INFERENCING WITH TORCHSCRIPT IN ORDER TO GET ONE OUTPUT IMAGE\n",
    "\n",
    "import torch\n",
    "from PIL import Image\n",
    "from torchvision import transforms\n",
    "\n",
    "def main():\n",
    "    # Load the TorchScript model\n",
    "    model = torch.jit.load('20x/modelpnnx20x.pt')\n",
    "\n",
    "    # Load the images from disk\n",
    "    image_a = Image.open('demo_video/STB-VMM_demo_x20_static_original/frame_000007.png')\n",
    "    image_b = Image.open('demo_video/STB-VMM_demo_x20_static_original/frame_000008.png')\n",
    "\n",
    "    # Resize the images if necessary\n",
    "    desired_size = (384, 384)\n",
    "    image_a = image_a.resize(desired_size)\n",
    "    image_b = image_b.resize(desired_size)\n",
    "\n",
    "    # Convert the images to tensors\n",
    "    transform = transforms.ToTensor()\n",
    "    a_tensor = transform(image_a).unsqueeze(0)\n",
    "    b_tensor = transform(image_b).unsqueeze(0)\n",
    "\n",
    "    # Run inference\n",
    "    output_tuple = model(a_tensor, b_tensor)\n",
    "    output = output_tuple[0]  # Access the desired tensor from the tuple\n",
    "\n",
    "    # Reshape the tensor (remove batch dimension)\n",
    "    output = output.squeeze(0)\n",
    "\n",
    "    # Convert the output tensor to a PIL image\n",
    "    output_image = transforms.ToPILImage()(output.cpu())\n",
    "\n",
    "    # Save the combined image\n",
    "    output_image.save('combined_image.png')\n",
    "\n",
    "\n",
    "if __name__ == '__main__':\n",
    "    main()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "1163a12f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processed 10 pairs out of 150\n",
      "Processed 20 pairs out of 150\n",
      "Processed 30 pairs out of 150\n",
      "Processed 40 pairs out of 150\n",
      "Processed 50 pairs out of 150\n",
      "Processed 60 pairs out of 150\n",
      "Processed 70 pairs out of 150\n",
      "Processed 80 pairs out of 150\n",
      "Processed 90 pairs out of 150\n",
      "Processed 100 pairs out of 150\n",
      "Processed 110 pairs out of 150\n",
      "Processed 120 pairs out of 150\n",
      "Processed 130 pairs out of 150\n",
      "Processed 140 pairs out of 150\n",
      "Processed 150 pairs out of 150\n",
      "Processing completed.\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "from PIL import Image\n",
    "from torchvision import transforms\n",
    "import os\n",
    "\n",
    "def main():\n",
    "    # Load the TorchScript model\n",
    "    model = torch.jit.load('20x/modelpnnx20x.pt')\n",
    "\n",
    "    # Directory path of input frames\n",
    "    frames_dir = 'demo_video/STB-VMM_demo_x20_static_original'\n",
    "\n",
    "    # Output directory path to save generated frames\n",
    "    output_dir = 'output_frames'\n",
    "    if not os.path.exists(output_dir):\n",
    "        os.makedirs(output_dir)\n",
    "\n",
    "    # Resize the images if necessary\n",
    "    desired_size = (384, 384)\n",
    "\n",
    "    # Get the list of frames in the directory\n",
    "    frame_files = sorted(os.listdir(frames_dir))\n",
    "\n",
    "    # Process the frames in pairs and generate output images\n",
    "    num_frames = len(frame_files)\n",
    "    num_pairs = min(num_frames // 2, 150)\n",
    "\n",
    "    for i in range(num_pairs):\n",
    "        # Load the pair of frames\n",
    "        frame_a_path = os.path.join(frames_dir, frame_files[2 * i])\n",
    "        frame_b_path = os.path.join(frames_dir, frame_files[2 * i + 1])\n",
    "\n",
    "        image_a = Image.open(frame_a_path)\n",
    "        image_b = Image.open(frame_b_path)\n",
    "\n",
    "        # Resize the images\n",
    "        image_a = image_a.resize(desired_size)\n",
    "        image_b = image_b.resize(desired_size)\n",
    "\n",
    "        # Convert the images to tensors\n",
    "        transform = transforms.ToTensor()\n",
    "        a_tensor = transform(image_a).unsqueeze(0)\n",
    "        b_tensor = transform(image_b).unsqueeze(0)\n",
    "\n",
    "        # Set up amp tensor and unsqueeze to match the shape expected by the model\n",
    "        amp_tensor = torch.tensor([20]).unsqueeze(1).unsqueeze(1).unsqueeze(1)\n",
    "\n",
    "        # Run inference\n",
    "        output_tuple = model(a_tensor, b_tensor, amp_tensor)\n",
    "        output = output_tuple[0]  # Access the desired tensor from the tuple\n",
    "\n",
    "        # Reshape the tensor (remove batch dimension)\n",
    "        output = output.squeeze(0)\n",
    "\n",
    "        # Convert the output tensor to a PIL image\n",
    "        output_image = transforms.ToPILImage()(output.cpu())\n",
    "\n",
    "        # Save the output image\n",
    "        output_path = os.path.join(output_dir, f'output_{i + 1:03d}.png')\n",
    "        output_image.save(output_path)\n",
    "\n",
    "        if (i + 1) % 10 == 0:\n",
    "            print(f'Processed {i + 1} pairs out of {num_pairs}')\n",
    "\n",
    "    print('Processing completed.')\n",
    "\n",
    "\n",
    "if __name__ == '__main__':\n",
    "    main()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "5200393e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processed 10 pairs out of 150\n",
      "Processed 20 pairs out of 150\n",
      "Processed 30 pairs out of 150\n",
      "Processed 40 pairs out of 150\n",
      "Processed 50 pairs out of 150\n",
      "Processed 60 pairs out of 150\n",
      "Processed 70 pairs out of 150\n",
      "Processed 80 pairs out of 150\n",
      "Processed 90 pairs out of 150\n",
      "Processed 100 pairs out of 150\n",
      "Processed 110 pairs out of 150\n",
      "Processed 120 pairs out of 150\n",
      "Processed 130 pairs out of 150\n",
      "Processed 140 pairs out of 150\n",
      "Processed 150 pairs out of 150\n",
      "Processing completed.\n"
     ]
    }
   ],
   "source": [
    "#INFERENCE WITH ALL FRAMES FROM DIRECTORY\n",
    "\n",
    "import torch\n",
    "from PIL import Image\n",
    "from torchvision import transforms\n",
    "import os\n",
    "\n",
    "def main():\n",
    "    # Load the TorchScript model\n",
    "    model = torch.jit.load('20x/modelpnnx20x.pt')\n",
    "\n",
    "    # Directory path of input frames\n",
    "    frames_dir = 'demo_video/STB-VMM_demo_x20_static_original'\n",
    "\n",
    "    # Output directory path to save generated frames\n",
    "    output_dir = 'output_frames'\n",
    "    if not os.path.exists(output_dir):\n",
    "        os.makedirs(output_dir)\n",
    "\n",
    "    # Resize the images if necessary\n",
    "    desired_size = (384, 384)\n",
    "\n",
    "    # Get the list of frames in the directory\n",
    "    frame_files = sorted(os.listdir(frames_dir))\n",
    "\n",
    "    # Process the frames in pairs and generate output images\n",
    "    num_frames = len(frame_files)\n",
    "    num_pairs = min(num_frames // 2, 150)\n",
    "\n",
    "    for i in range(num_pairs):\n",
    "        # Load the pair of frames\n",
    "        frame_a_path = os.path.join(frames_dir, frame_files[2 * i])\n",
    "        frame_b_path = os.path.join(frames_dir, frame_files[2 * i + 1])\n",
    "\n",
    "        image_a = Image.open(frame_a_path)\n",
    "        image_b = Image.open(frame_b_path)\n",
    "\n",
    "        # Resize the images\n",
    "        image_a = image_a.resize(desired_size)\n",
    "        image_b = image_b.resize(desired_size)\n",
    "\n",
    "        # Convert the images to tensors\n",
    "        transform = transforms.ToTensor()\n",
    "        a_tensor = transform(image_a).unsqueeze(0)\n",
    "        b_tensor = transform(image_b).unsqueeze(0)\n",
    "\n",
    "        # Run inference\n",
    "        output_tuple = model(a_tensor, b_tensor)\n",
    "        output = output_tuple[0]  # Access the desired tensor from the tuple\n",
    "\n",
    "        # Reshape the tensor (remove batch dimension)\n",
    "        output = output.squeeze(0)\n",
    "\n",
    "        # Convert the output tensor to a PIL image\n",
    "        output_image = transforms.ToPILImage()(output.cpu())\n",
    "\n",
    "        # Save the output image\n",
    "        output_path = os.path.join(output_dir, f'output_{i + 1:03d}.png')\n",
    "        output_image.save(output_path)\n",
    "\n",
    "        if (i + 1) % 10 == 0:\n",
    "            print(f'Processed {i + 1} pairs out of {num_pairs}')\n",
    "\n",
    "    print('Processing completed.')\n",
    "\n",
    "\n",
    "if __name__ == '__main__':\n",
    "    main()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "683bd2ad",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "demo\n"
     ]
    },
    {
     "ename": "ValueError",
     "evalue": "pic should be 2/3 dimensional. Got 4 dimensions.",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mValueError\u001b[0m                                Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[10], line 60\u001b[0m\n\u001b[1;32m     57\u001b[0m             \u001b[38;5;28;01mbreak\u001b[39;00m\n\u001b[1;32m     59\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;18m__name__\u001b[39m \u001b[38;5;241m==\u001b[39m \u001b[38;5;124m'\u001b[39m\u001b[38;5;124m__main__\u001b[39m\u001b[38;5;124m'\u001b[39m:\n\u001b[0;32m---> 60\u001b[0m     \u001b[43mmain\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n",
      "Cell \u001b[0;32mIn[10], line 49\u001b[0m, in \u001b[0;36mmain\u001b[0;34m()\u001b[0m\n\u001b[1;32m     46\u001b[0m output \u001b[38;5;241m=\u001b[39m output_tuple[\u001b[38;5;241m0\u001b[39m][\u001b[38;5;241m-\u001b[39m\u001b[38;5;241m3\u001b[39m:, :, :]  \u001b[38;5;66;03m# Take the last 3 channels\u001b[39;00m\n\u001b[1;32m     48\u001b[0m \u001b[38;5;66;03m# Convert the output tensor to a PIL image\u001b[39;00m\n\u001b[0;32m---> 49\u001b[0m output_image \u001b[38;5;241m=\u001b[39m \u001b[43mtransforms\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mToPILImage\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[43m(\u001b[49m\u001b[43moutput\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mcpu\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     51\u001b[0m \u001b[38;5;66;03m# Save the output image\u001b[39;00m\n\u001b[1;32m     52\u001b[0m output_file \u001b[38;5;241m=\u001b[39m \u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124moutput_\u001b[39m\u001b[38;5;132;01m{\u001b[39;00mi\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m.png\u001b[39m\u001b[38;5;124m\"\u001b[39m\n",
      "File \u001b[0;32m/opt/homebrew/Caskroom/miniconda/base/envs/afstudeer/lib/python3.10/site-packages/torchvision/transforms/transforms.py:234\u001b[0m, in \u001b[0;36mToPILImage.__call__\u001b[0;34m(self, pic)\u001b[0m\n\u001b[1;32m    225\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m__call__\u001b[39m(\u001b[38;5;28mself\u001b[39m, pic):\n\u001b[1;32m    226\u001b[0m \u001b[38;5;250m    \u001b[39m\u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m    227\u001b[0m \u001b[38;5;124;03m    Args:\u001b[39;00m\n\u001b[1;32m    228\u001b[0m \u001b[38;5;124;03m        pic (Tensor or numpy.ndarray): Image to be converted to PIL Image.\u001b[39;00m\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m    232\u001b[0m \n\u001b[1;32m    233\u001b[0m \u001b[38;5;124;03m    \"\"\"\u001b[39;00m\n\u001b[0;32m--> 234\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mF\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mto_pil_image\u001b[49m\u001b[43m(\u001b[49m\u001b[43mpic\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mmode\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m/opt/homebrew/Caskroom/miniconda/base/envs/afstudeer/lib/python3.10/site-packages/torchvision/transforms/functional.py:266\u001b[0m, in \u001b[0;36mto_pil_image\u001b[0;34m(pic, mode)\u001b[0m\n\u001b[1;32m    264\u001b[0m \u001b[38;5;28;01melif\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(pic, torch\u001b[38;5;241m.\u001b[39mTensor):\n\u001b[1;32m    265\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m pic\u001b[38;5;241m.\u001b[39mndimension() \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;129;01min\u001b[39;00m {\u001b[38;5;241m2\u001b[39m, \u001b[38;5;241m3\u001b[39m}:\n\u001b[0;32m--> 266\u001b[0m         \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mValueError\u001b[39;00m(\u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mpic should be 2/3 dimensional. Got \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mpic\u001b[38;5;241m.\u001b[39mndimension()\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m dimensions.\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[1;32m    268\u001b[0m     \u001b[38;5;28;01melif\u001b[39;00m pic\u001b[38;5;241m.\u001b[39mndimension() \u001b[38;5;241m==\u001b[39m \u001b[38;5;241m2\u001b[39m:\n\u001b[1;32m    269\u001b[0m         \u001b[38;5;66;03m# if 2D image, add channel dimension (CHW)\u001b[39;00m\n\u001b[1;32m    270\u001b[0m         pic \u001b[38;5;241m=\u001b[39m pic\u001b[38;5;241m.\u001b[39munsqueeze(\u001b[38;5;241m0\u001b[39m)\n",
      "\u001b[0;31mValueError\u001b[0m: pic should be 2/3 dimensional. Got 4 dimensions."
     ]
    }
   ],
   "source": [
    "import os\n",
    "import numpy as np\n",
    "import torch\n",
    "from PIL import Image\n",
    "from torchvision import transforms\n",
    "\n",
    "def main():\n",
    "    # Load the TorchScript model\n",
    "    model = torch.jit.load('20x/modelpnnx20x.pt')\n",
    "\n",
    "    # Check saving directory\n",
    "    save_dir = 'demo'\n",
    "    if not os.path.exists(save_dir):\n",
    "        os.makedirs(save_dir)\n",
    "    print(save_dir)\n",
    "\n",
    "    # Load the images from disk\n",
    "    frame_files = sorted(os.listdir('demo_video/STB-VMM_demo_x20_static_original'))\n",
    "\n",
    "    # Resize the images if necessary\n",
    "    desired_size = (384, 384)\n",
    "\n",
    "    # Process the frames and generate output images\n",
    "    num_frames = len(frame_files)\n",
    "    num_pairs = min(num_frames // 2, 150)\n",
    "\n",
    "    for i in range(num_pairs):\n",
    "        # Load the pair of frames\n",
    "        frame_a_path = os.path.join('demo_video/STB-VMM_demo_x20_static_original', frame_files[2 * i])\n",
    "        frame_b_path = os.path.join('demo_video/STB-VMM_demo_x20_static_original', frame_files[2 * i + 1])\n",
    "\n",
    "        image_a = Image.open(frame_a_path)\n",
    "        image_b = Image.open(frame_b_path)\n",
    "\n",
    "        # Resize the images\n",
    "        image_a = image_a.resize(desired_size)\n",
    "        image_b = image_b.resize(desired_size)\n",
    "\n",
    "        # Convert the images to tensors\n",
    "        transform = transforms.ToTensor()\n",
    "        a_tensor = transform(image_a).unsqueeze(0)\n",
    "        b_tensor = transform(image_b).unsqueeze(0)\n",
    "\n",
    "        # Run inference\n",
    "        output_tuple = model(a_tensor, b_tensor)\n",
    "        output = output_tuple[0][-3:, :, :]  # Take the last 3 channels\n",
    "\n",
    "        # Convert the output tensor to a PIL image\n",
    "        output_image = transforms.ToPILImage()(output.cpu())\n",
    "\n",
    "        # Save the output image\n",
    "        output_file = f\"output_{i}.png\"\n",
    "        output_path = os.path.join(save_dir, output_file)\n",
    "        output_image.save(output_path)\n",
    "\n",
    "        if (i+1) == num_pairs:\n",
    "            break\n",
    "\n",
    "if __name__ == '__main__':\n",
    "    main()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "8eb57d50",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "IOPub data rate exceeded.\n",
      "The notebook server will temporarily stop sending output\n",
      "to the client in order to avoid crashing it.\n",
      "To change this limit, set the config variable\n",
      "`--NotebookApp.iopub_data_rate_limit`.\n",
      "\n",
      "Current values:\n",
      "NotebookApp.iopub_data_rate_limit=1000000.0 (bytes/sec)\n",
      "NotebookApp.rate_limit_window=3.0 (secs)\n",
      "\n"
     ]
    }
   ],
   "source": [
    "import onnx\n",
    "\n",
    "# Load the ONNX model\n",
    "model = onnx.load(\"20x/modelpnnx20x.pnnx.onnx\")\n",
    "\n",
    "# Add the opset version\n",
    "model.opset_import.add()\n",
    "model.opset_import[0].version = 11\n",
    "model.ir_version = 7\n",
    "\n",
    "\n",
    "# Check the model\n",
    "print(model)\n",
    "\n",
    "# Save the modified model\n",
    "onnx.save(model, \"model_with_opset.onnx\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "13f079ac",
   "metadata": {},
   "outputs": [
    {
     "ename": "Fail",
     "evalue": "[ONNXRuntimeError] : 1 : FAIL : Load model from model_with_opset.onnx failed:Invalid tensor data type 0.",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mFail\u001b[0m                                      Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[1], line 43\u001b[0m\n\u001b[1;32m     40\u001b[0m     output_image\u001b[39m.\u001b[39msave(\u001b[39m'\u001b[39m\u001b[39mcombined_image.png\u001b[39m\u001b[39m'\u001b[39m)\n\u001b[1;32m     42\u001b[0m \u001b[39mif\u001b[39;00m \u001b[39m__name__\u001b[39m \u001b[39m==\u001b[39m \u001b[39m'\u001b[39m\u001b[39m__main__\u001b[39m\u001b[39m'\u001b[39m:\n\u001b[0;32m---> 43\u001b[0m     main()\n",
      "Cell \u001b[0;32mIn[1], line 10\u001b[0m, in \u001b[0;36mmain\u001b[0;34m()\u001b[0m\n\u001b[1;32m      6\u001b[0m \u001b[39mdef\u001b[39;00m \u001b[39mmain\u001b[39m():\n\u001b[1;32m      7\u001b[0m     \u001b[39m# Load the ONNX model\u001b[39;00m\n\u001b[1;32m      8\u001b[0m     onnx_model_path \u001b[39m=\u001b[39m \u001b[39m'\u001b[39m\u001b[39mmodel_with_opset.onnx\u001b[39m\u001b[39m'\u001b[39m\n\u001b[0;32m---> 10\u001b[0m     ort_session \u001b[39m=\u001b[39m onnxruntime\u001b[39m.\u001b[39;49mInferenceSession(onnx_model_path)\n\u001b[1;32m     12\u001b[0m     \u001b[39m# Load the images from disk\u001b[39;00m\n\u001b[1;32m     13\u001b[0m     image_a \u001b[39m=\u001b[39m Image\u001b[39m.\u001b[39mopen(\u001b[39m'\u001b[39m\u001b[39mdemo_video/STB-VMM_demo_x20_static_original/frame_000007.png\u001b[39m\u001b[39m'\u001b[39m)\n",
      "File \u001b[0;32m/opt/homebrew/Caskroom/miniconda/base/envs/testconv/lib/python3.9/site-packages/onnxruntime/capi/onnxruntime_inference_collection.py:347\u001b[0m, in \u001b[0;36mInferenceSession.__init__\u001b[0;34m(self, path_or_bytes, sess_options, providers, provider_options, **kwargs)\u001b[0m\n\u001b[1;32m    344\u001b[0m disabled_optimizers \u001b[39m=\u001b[39m kwargs[\u001b[39m\"\u001b[39m\u001b[39mdisabled_optimizers\u001b[39m\u001b[39m\"\u001b[39m] \u001b[39mif\u001b[39;00m \u001b[39m\"\u001b[39m\u001b[39mdisabled_optimizers\u001b[39m\u001b[39m\"\u001b[39m \u001b[39min\u001b[39;00m kwargs \u001b[39melse\u001b[39;00m \u001b[39mNone\u001b[39;00m\n\u001b[1;32m    346\u001b[0m \u001b[39mtry\u001b[39;00m:\n\u001b[0;32m--> 347\u001b[0m     \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49m_create_inference_session(providers, provider_options, disabled_optimizers)\n\u001b[1;32m    348\u001b[0m \u001b[39mexcept\u001b[39;00m \u001b[39mValueError\u001b[39;00m:\n\u001b[1;32m    349\u001b[0m     \u001b[39mif\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_enable_fallback:\n",
      "File \u001b[0;32m/opt/homebrew/Caskroom/miniconda/base/envs/testconv/lib/python3.9/site-packages/onnxruntime/capi/onnxruntime_inference_collection.py:384\u001b[0m, in \u001b[0;36mInferenceSession._create_inference_session\u001b[0;34m(self, providers, provider_options, disabled_optimizers)\u001b[0m\n\u001b[1;32m    382\u001b[0m session_options \u001b[39m=\u001b[39m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_sess_options \u001b[39mif\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_sess_options \u001b[39melse\u001b[39;00m C\u001b[39m.\u001b[39mget_default_session_options()\n\u001b[1;32m    383\u001b[0m \u001b[39mif\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_model_path:\n\u001b[0;32m--> 384\u001b[0m     sess \u001b[39m=\u001b[39m C\u001b[39m.\u001b[39;49mInferenceSession(session_options, \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49m_model_path, \u001b[39mTrue\u001b[39;49;00m, \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49m_read_config_from_model)\n\u001b[1;32m    385\u001b[0m \u001b[39melse\u001b[39;00m:\n\u001b[1;32m    386\u001b[0m     sess \u001b[39m=\u001b[39m C\u001b[39m.\u001b[39mInferenceSession(session_options, \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_model_bytes, \u001b[39mFalse\u001b[39;00m, \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_read_config_from_model)\n",
      "\u001b[0;31mFail\u001b[0m: [ONNXRuntimeError] : 1 : FAIL : Load model from model_with_opset.onnx failed:Invalid tensor data type 0."
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import onnxruntime\n",
    "from torchvision import transforms\n",
    "from PIL import Image\n",
    "\n",
    "def main():\n",
    "    # Load the ONNX model\n",
    "    onnx_model_path = 'model_with_opset.onnx'\n",
    "\n",
    "    ort_session = onnxruntime.InferenceSession(onnx_model_path)\n",
    "\n",
    "    # Load the images from disk\n",
    "    image_a = Image.open('demo_video/STB-VMM_demo_x20_static_original/frame_000007.png')\n",
    "    image_b = Image.open('demo_video/STB-VMM_demo_x20_static_original/frame_000008.png')\n",
    "\n",
    "    # Resize the images if necessary\n",
    "    desired_size = (384, 384)\n",
    "    image_a = image_a.resize(desired_size)\n",
    "    image_b = image_b.resize(desired_size)\n",
    "\n",
    "    # Convert the images to tensors\n",
    "    transform = transforms.ToTensor()\n",
    "    a_tensor = transform(image_a).unsqueeze(0)\n",
    "    b_tensor = transform(image_b).unsqueeze(0)\n",
    "\n",
    "    # Convert the input tensors to NumPy arrays\n",
    "    a_np = a_tensor.numpy()\n",
    "    b_np = b_tensor.numpy()\n",
    "\n",
    "    # Run inference using ONNX Runtime\n",
    "    input_names = ort_session.get_inputs()[0].name\n",
    "    output_names = [output.name for output in ort_session.get_outputs()]\n",
    "    output = ort_session.run(output_names, {input_names: (a_np, b_np)})\n",
    "\n",
    "    # Convert the output tensor to a PIL image\n",
    "    output_tensor = torch.from_numpy(output[0])\n",
    "    output_image = transforms.ToPILImage()(output_tensor.cpu())\n",
    "\n",
    "    # Save the combined image\n",
    "    output_image.save('combined_image.png')\n",
    "\n",
    "if __name__ == '__main__':\n",
    "    main()\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "f4bcfceb",
   "metadata": {},
   "source": [
    "## Everything below can be ignored as this is what was used during prevous testing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "d0cd0fce",
   "metadata": {
    "is_executing": true
   },
   "outputs": [
    {
     "ename": "TypeError",
     "evalue": "forward() takes from 2 to 3 positional arguments but 4 were given",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mTypeError\u001b[0m                                 Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[21], line 23\u001b[0m\n\u001b[1;32m     20\u001b[0m model\u001b[39m.\u001b[39meval()\n\u001b[1;32m     22\u001b[0m \u001b[39m# Export the model to ONNX\u001b[39;00m\n\u001b[0;32m---> 23\u001b[0m torch\u001b[39m.\u001b[39;49monnx\u001b[39m.\u001b[39;49mexport(model, (a_dummy,b_dummy, amp_dummy), \u001b[39m'\u001b[39;49m\u001b[39mmodelforpnnx.onnx\u001b[39;49m\u001b[39m'\u001b[39;49m)\n",
      "File \u001b[0;32m/opt/homebrew/Caskroom/miniconda/base/envs/testconv/lib/python3.9/site-packages/torch/onnx/__init__.py:305\u001b[0m, in \u001b[0;36mexport\u001b[0;34m(model, args, f, export_params, verbose, training, input_names, output_names, operator_export_type, opset_version, do_constant_folding, dynamic_axes, keep_initializers_as_inputs, custom_opsets, export_modules_as_functions)\u001b[0m\n\u001b[1;32m     39\u001b[0m \u001b[39m\u001b[39m\u001b[39mr\u001b[39m\u001b[39m\"\"\"\u001b[39;00m\n\u001b[1;32m     40\u001b[0m \u001b[39mExports a model into ONNX format. If ``model`` is not a\u001b[39;00m\n\u001b[1;32m     41\u001b[0m \u001b[39m:class:`torch.jit.ScriptModule` nor a :class:`torch.jit.ScriptFunction`, this runs\u001b[39;00m\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m    301\u001b[0m \u001b[39m    model to the file ``f`` even if this is raised.\u001b[39;00m\n\u001b[1;32m    302\u001b[0m \u001b[39m\"\"\"\u001b[39;00m\n\u001b[1;32m    304\u001b[0m \u001b[39mfrom\u001b[39;00m \u001b[39mtorch\u001b[39;00m\u001b[39m.\u001b[39;00m\u001b[39monnx\u001b[39;00m \u001b[39mimport\u001b[39;00m utils\n\u001b[0;32m--> 305\u001b[0m \u001b[39mreturn\u001b[39;00m utils\u001b[39m.\u001b[39;49mexport(model, args, f, export_params, verbose, training,\n\u001b[1;32m    306\u001b[0m                     input_names, output_names, operator_export_type, opset_version,\n\u001b[1;32m    307\u001b[0m                     do_constant_folding, dynamic_axes,\n\u001b[1;32m    308\u001b[0m                     keep_initializers_as_inputs, custom_opsets,\n\u001b[1;32m    309\u001b[0m                     export_modules_as_functions)\n",
      "File \u001b[0;32m/opt/homebrew/Caskroom/miniconda/base/envs/testconv/lib/python3.9/site-packages/torch/onnx/utils.py:118\u001b[0m, in \u001b[0;36mexport\u001b[0;34m(model, args, f, export_params, verbose, training, input_names, output_names, operator_export_type, opset_version, do_constant_folding, dynamic_axes, keep_initializers_as_inputs, custom_opsets, export_modules_as_functions)\u001b[0m\n\u001b[1;32m    115\u001b[0m     \u001b[39melse\u001b[39;00m:\n\u001b[1;32m    116\u001b[0m         operator_export_type \u001b[39m=\u001b[39m OperatorExportTypes\u001b[39m.\u001b[39mONNX\n\u001b[0;32m--> 118\u001b[0m _export(model, args, f, export_params, verbose, training, input_names, output_names,\n\u001b[1;32m    119\u001b[0m         operator_export_type\u001b[39m=\u001b[39;49moperator_export_type, opset_version\u001b[39m=\u001b[39;49mopset_version,\n\u001b[1;32m    120\u001b[0m         do_constant_folding\u001b[39m=\u001b[39;49mdo_constant_folding, dynamic_axes\u001b[39m=\u001b[39;49mdynamic_axes,\n\u001b[1;32m    121\u001b[0m         keep_initializers_as_inputs\u001b[39m=\u001b[39;49mkeep_initializers_as_inputs,\n\u001b[1;32m    122\u001b[0m         custom_opsets\u001b[39m=\u001b[39;49mcustom_opsets, export_modules_as_functions\u001b[39m=\u001b[39;49mexport_modules_as_functions)\n",
      "File \u001b[0;32m/opt/homebrew/Caskroom/miniconda/base/envs/testconv/lib/python3.9/site-packages/torch/onnx/utils.py:719\u001b[0m, in \u001b[0;36m_export\u001b[0;34m(model, args, f, export_params, verbose, training, input_names, output_names, operator_export_type, export_type, opset_version, do_constant_folding, dynamic_axes, keep_initializers_as_inputs, fixed_batch_size, custom_opsets, add_node_names, onnx_shape_inference, export_modules_as_functions)\u001b[0m\n\u001b[1;32m    715\u001b[0m     dynamic_axes \u001b[39m=\u001b[39m {}\n\u001b[1;32m    716\u001b[0m _validate_dynamic_axes(dynamic_axes, model, input_names, output_names)\n\u001b[1;32m    718\u001b[0m graph, params_dict, torch_out \u001b[39m=\u001b[39m \\\n\u001b[0;32m--> 719\u001b[0m     _model_to_graph(model, args, verbose, input_names,\n\u001b[1;32m    720\u001b[0m                     output_names, operator_export_type,\n\u001b[1;32m    721\u001b[0m                     val_do_constant_folding,\n\u001b[1;32m    722\u001b[0m                     fixed_batch_size\u001b[39m=\u001b[39;49mfixed_batch_size,\n\u001b[1;32m    723\u001b[0m                     training\u001b[39m=\u001b[39;49mtraining,\n\u001b[1;32m    724\u001b[0m                     dynamic_axes\u001b[39m=\u001b[39;49mdynamic_axes)\n\u001b[1;32m    726\u001b[0m \u001b[39m# TODO: Don't allocate a in-memory string for the protobuf\u001b[39;00m\n\u001b[1;32m    727\u001b[0m defer_weight_export \u001b[39m=\u001b[39m export_type \u001b[39mis\u001b[39;00m \u001b[39mnot\u001b[39;00m ExportTypes\u001b[39m.\u001b[39mPROTOBUF_FILE\n",
      "File \u001b[0;32m/opt/homebrew/Caskroom/miniconda/base/envs/testconv/lib/python3.9/site-packages/torch/onnx/utils.py:499\u001b[0m, in \u001b[0;36m_model_to_graph\u001b[0;34m(model, args, verbose, input_names, output_names, operator_export_type, do_constant_folding, _disable_torch_constant_prop, fixed_batch_size, training, dynamic_axes)\u001b[0m\n\u001b[1;32m    496\u001b[0m \u001b[39mif\u001b[39;00m \u001b[39misinstance\u001b[39m(args, (torch\u001b[39m.\u001b[39mTensor, \u001b[39mint\u001b[39m, \u001b[39mfloat\u001b[39m, \u001b[39mbool\u001b[39m)):\n\u001b[1;32m    497\u001b[0m     args \u001b[39m=\u001b[39m (args, )\n\u001b[0;32m--> 499\u001b[0m graph, params, torch_out, module \u001b[39m=\u001b[39m _create_jit_graph(model, args)\n\u001b[1;32m    501\u001b[0m params_dict \u001b[39m=\u001b[39m _get_named_param_dict(graph, params)\n\u001b[1;32m    503\u001b[0m graph \u001b[39m=\u001b[39m _optimize_graph(graph, operator_export_type,\n\u001b[1;32m    504\u001b[0m                         _disable_torch_constant_prop\u001b[39m=\u001b[39m_disable_torch_constant_prop,\n\u001b[1;32m    505\u001b[0m                         fixed_batch_size\u001b[39m=\u001b[39mfixed_batch_size, params_dict\u001b[39m=\u001b[39mparams_dict,\n\u001b[1;32m    506\u001b[0m                         dynamic_axes\u001b[39m=\u001b[39mdynamic_axes, input_names\u001b[39m=\u001b[39minput_names,\n\u001b[1;32m    507\u001b[0m                         module\u001b[39m=\u001b[39mmodule)\n",
      "File \u001b[0;32m/opt/homebrew/Caskroom/miniconda/base/envs/testconv/lib/python3.9/site-packages/torch/onnx/utils.py:440\u001b[0m, in \u001b[0;36m_create_jit_graph\u001b[0;34m(model, args)\u001b[0m\n\u001b[1;32m    438\u001b[0m     \u001b[39mreturn\u001b[39;00m graph, params, torch_out, \u001b[39mNone\u001b[39;00m\n\u001b[1;32m    439\u001b[0m \u001b[39melse\u001b[39;00m:\n\u001b[0;32m--> 440\u001b[0m     graph, torch_out \u001b[39m=\u001b[39m _trace_and_get_graph_from_model(model, args)\n\u001b[1;32m    441\u001b[0m     torch\u001b[39m.\u001b[39m_C\u001b[39m.\u001b[39m_jit_pass_onnx_lint(graph)\n\u001b[1;32m    442\u001b[0m     state_dict \u001b[39m=\u001b[39m _unique_state_dict(model)\n",
      "File \u001b[0;32m/opt/homebrew/Caskroom/miniconda/base/envs/testconv/lib/python3.9/site-packages/torch/onnx/utils.py:391\u001b[0m, in \u001b[0;36m_trace_and_get_graph_from_model\u001b[0;34m(model, args)\u001b[0m\n\u001b[1;32m    384\u001b[0m \u001b[39mdef\u001b[39;00m \u001b[39m_trace_and_get_graph_from_model\u001b[39m(model, args):\n\u001b[1;32m    385\u001b[0m \n\u001b[1;32m    386\u001b[0m     \u001b[39m# A basic sanity check: make sure the state_dict keys are the same\u001b[39;00m\n\u001b[1;32m    387\u001b[0m     \u001b[39m# before and after running the model.  Fail fast!\u001b[39;00m\n\u001b[1;32m    388\u001b[0m     orig_state_dict_keys \u001b[39m=\u001b[39m _unique_state_dict(model)\u001b[39m.\u001b[39mkeys()\n\u001b[1;32m    390\u001b[0m     trace_graph, torch_out, inputs_states \u001b[39m=\u001b[39m \\\n\u001b[0;32m--> 391\u001b[0m         torch\u001b[39m.\u001b[39;49mjit\u001b[39m.\u001b[39;49m_get_trace_graph(model, args, strict\u001b[39m=\u001b[39;49m\u001b[39mFalse\u001b[39;49;00m, _force_outplace\u001b[39m=\u001b[39;49m\u001b[39mFalse\u001b[39;49;00m, _return_inputs_states\u001b[39m=\u001b[39;49m\u001b[39mTrue\u001b[39;49;00m)\n\u001b[1;32m    392\u001b[0m     warn_on_static_input_change(inputs_states)\n\u001b[1;32m    394\u001b[0m     \u001b[39mif\u001b[39;00m orig_state_dict_keys \u001b[39m!=\u001b[39m _unique_state_dict(model)\u001b[39m.\u001b[39mkeys():\n",
      "File \u001b[0;32m/opt/homebrew/Caskroom/miniconda/base/envs/testconv/lib/python3.9/site-packages/torch/jit/_trace.py:1166\u001b[0m, in \u001b[0;36m_get_trace_graph\u001b[0;34m(f, args, kwargs, strict, _force_outplace, return_inputs, _return_inputs_states)\u001b[0m\n\u001b[1;32m   1164\u001b[0m \u001b[39mif\u001b[39;00m \u001b[39mnot\u001b[39;00m \u001b[39misinstance\u001b[39m(args, \u001b[39mtuple\u001b[39m):\n\u001b[1;32m   1165\u001b[0m     args \u001b[39m=\u001b[39m (args,)\n\u001b[0;32m-> 1166\u001b[0m outs \u001b[39m=\u001b[39m ONNXTracedModule(f, strict, _force_outplace, return_inputs, _return_inputs_states)(\u001b[39m*\u001b[39;49margs, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mkwargs)\n\u001b[1;32m   1167\u001b[0m \u001b[39mreturn\u001b[39;00m outs\n",
      "File \u001b[0;32m/opt/homebrew/Caskroom/miniconda/base/envs/testconv/lib/python3.9/site-packages/torch/nn/modules/module.py:1110\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m   1106\u001b[0m \u001b[39m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[1;32m   1107\u001b[0m \u001b[39m# this function, and just call forward.\u001b[39;00m\n\u001b[1;32m   1108\u001b[0m \u001b[39mif\u001b[39;00m \u001b[39mnot\u001b[39;00m (\u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_backward_hooks \u001b[39mor\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_forward_hooks \u001b[39mor\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_forward_pre_hooks \u001b[39mor\u001b[39;00m _global_backward_hooks\n\u001b[1;32m   1109\u001b[0m         \u001b[39mor\u001b[39;00m _global_forward_hooks \u001b[39mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[0;32m-> 1110\u001b[0m     \u001b[39mreturn\u001b[39;00m forward_call(\u001b[39m*\u001b[39;49m\u001b[39minput\u001b[39;49m, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mkwargs)\n\u001b[1;32m   1111\u001b[0m \u001b[39m# Do not call functions when jit is used\u001b[39;00m\n\u001b[1;32m   1112\u001b[0m full_backward_hooks, non_full_backward_hooks \u001b[39m=\u001b[39m [], []\n",
      "File \u001b[0;32m/opt/homebrew/Caskroom/miniconda/base/envs/testconv/lib/python3.9/site-packages/torch/jit/_trace.py:127\u001b[0m, in \u001b[0;36mONNXTracedModule.forward\u001b[0;34m(self, *args)\u001b[0m\n\u001b[1;32m    124\u001b[0m     \u001b[39melse\u001b[39;00m:\n\u001b[1;32m    125\u001b[0m         \u001b[39mreturn\u001b[39;00m \u001b[39mtuple\u001b[39m(out_vars)\n\u001b[0;32m--> 127\u001b[0m graph, out \u001b[39m=\u001b[39m torch\u001b[39m.\u001b[39;49m_C\u001b[39m.\u001b[39;49m_create_graph_by_tracing(\n\u001b[1;32m    128\u001b[0m     wrapper,\n\u001b[1;32m    129\u001b[0m     in_vars \u001b[39m+\u001b[39;49m module_state,\n\u001b[1;32m    130\u001b[0m     _create_interpreter_name_lookup_fn(),\n\u001b[1;32m    131\u001b[0m     \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49mstrict,\n\u001b[1;32m    132\u001b[0m     \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49m_force_outplace,\n\u001b[1;32m    133\u001b[0m )\n\u001b[1;32m    135\u001b[0m \u001b[39mif\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_return_inputs:\n\u001b[1;32m    136\u001b[0m     \u001b[39mreturn\u001b[39;00m graph, outs[\u001b[39m0\u001b[39m], ret_inputs[\u001b[39m0\u001b[39m]\n",
      "File \u001b[0;32m/opt/homebrew/Caskroom/miniconda/base/envs/testconv/lib/python3.9/site-packages/torch/jit/_trace.py:118\u001b[0m, in \u001b[0;36mONNXTracedModule.forward.<locals>.wrapper\u001b[0;34m(*args)\u001b[0m\n\u001b[1;32m    116\u001b[0m \u001b[39mif\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_return_inputs_states:\n\u001b[1;32m    117\u001b[0m     inputs_states\u001b[39m.\u001b[39mappend(_unflatten(in_args, in_desc))\n\u001b[0;32m--> 118\u001b[0m outs\u001b[39m.\u001b[39mappend(\u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49minner(\u001b[39m*\u001b[39;49mtrace_inputs))\n\u001b[1;32m    119\u001b[0m \u001b[39mif\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_return_inputs_states:\n\u001b[1;32m    120\u001b[0m     inputs_states[\u001b[39m0\u001b[39m] \u001b[39m=\u001b[39m (inputs_states[\u001b[39m0\u001b[39m], trace_inputs)\n",
      "File \u001b[0;32m/opt/homebrew/Caskroom/miniconda/base/envs/testconv/lib/python3.9/site-packages/torch/nn/modules/module.py:1110\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m   1106\u001b[0m \u001b[39m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[1;32m   1107\u001b[0m \u001b[39m# this function, and just call forward.\u001b[39;00m\n\u001b[1;32m   1108\u001b[0m \u001b[39mif\u001b[39;00m \u001b[39mnot\u001b[39;00m (\u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_backward_hooks \u001b[39mor\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_forward_hooks \u001b[39mor\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_forward_pre_hooks \u001b[39mor\u001b[39;00m _global_backward_hooks\n\u001b[1;32m   1109\u001b[0m         \u001b[39mor\u001b[39;00m _global_forward_hooks \u001b[39mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[0;32m-> 1110\u001b[0m     \u001b[39mreturn\u001b[39;00m forward_call(\u001b[39m*\u001b[39;49m\u001b[39minput\u001b[39;49m, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mkwargs)\n\u001b[1;32m   1111\u001b[0m \u001b[39m# Do not call functions when jit is used\u001b[39;00m\n\u001b[1;32m   1112\u001b[0m full_backward_hooks, non_full_backward_hooks \u001b[39m=\u001b[39m [], []\n",
      "File \u001b[0;32m/opt/homebrew/Caskroom/miniconda/base/envs/testconv/lib/python3.9/site-packages/torch/nn/modules/module.py:1098\u001b[0m, in \u001b[0;36mModule._slow_forward\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m   1096\u001b[0m         recording_scopes \u001b[39m=\u001b[39m \u001b[39mFalse\u001b[39;00m\n\u001b[1;32m   1097\u001b[0m \u001b[39mtry\u001b[39;00m:\n\u001b[0;32m-> 1098\u001b[0m     result \u001b[39m=\u001b[39m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49mforward(\u001b[39m*\u001b[39;49m\u001b[39minput\u001b[39;49m, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mkwargs)\n\u001b[1;32m   1099\u001b[0m \u001b[39mfinally\u001b[39;00m:\n\u001b[1;32m   1100\u001b[0m     \u001b[39mif\u001b[39;00m recording_scopes:\n",
      "\u001b[0;31mTypeError\u001b[0m: forward() takes from 2 to 3 positional arguments but 4 were given"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "from models.model2H5A import STBVMM\n",
    "\n",
    "# Define dummy inputs of the appropriate sizes for your model\n",
    "\n",
    "a_dummy = torch.randn(1, 3, 384, 384)\n",
    "b_dummy = torch.randn(1, 3, 384, 384)\n",
    "# c_dummy = torch.randn(1, 3, 384, 384)  # optional during training\n",
    "\n",
    "# For the amplitude parameter amp\n",
    "amp_dummy = torch.randn(1)  # adjusted to be a scalar value\n",
    "\n",
    "# Instantiate your model\n",
    "model = STBVMM()\n",
    "\n",
    "# Load the weights from the checkpoint\n",
    "checkpoint = torch.load('ckpt/ckpt_e49.pth.tar')\n",
    "model.load_state_dict(checkpoint['state_dict'], strict= False)\n",
    "# Make sure the model is in evaluation mode\n",
    "model.eval()\n",
    "\n",
    "# Export the model to ONNX\n",
    "torch.onnx.export(model, (a_dummy,b_dummy, amp_dummy), 'modelforpnnx.onnx')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "617bd840",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/raoulritter/STB-VMM/models/model.py:988: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!\n",
      "  if a.shape != b.shape:\n",
      "/Users/raoulritter/STB-VMM/models/model.py:153: UserWarning: __floordiv__ is deprecated, and its behavior will change in a future version of pytorch. It currently rounds toward 0 (like the 'trunc' function NOT 'floor'). This results in incorrect rounding for negative values. To keep the current behavior, use torch.div(a, b, rounding_mode='trunc'), or for actual floor division, use torch.div(a, b, rounding_mode='floor').\n",
      "  x = x.view(B, H // window_size, window_size, W // window_size, window_size, C)\n",
      "/Users/raoulritter/STB-VMM/models/model.py:371: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!\n",
      "  if self.input_resolution == x_size:\n",
      "/Users/raoulritter/STB-VMM/models/model.py:232: UserWarning: __floordiv__ is deprecated, and its behavior will change in a future version of pytorch. It currently rounds toward 0 (like the 'trunc' function NOT 'floor'). This results in incorrect rounding for negative values. To keep the current behavior, use torch.div(a, b, rounding_mode='trunc'), or for actual floor division, use torch.div(a, b, rounding_mode='floor').\n",
      "  qkv = self.qkv(x).reshape(B_, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)\n",
      "/Users/raoulritter/STB-VMM/models/model.py:169: TracerWarning: Converting a tensor to a Python integer might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!\n",
      "  B = int(windows.shape[0] / (H * W / window_size / window_size))\n",
      "/Users/raoulritter/STB-VMM/models/model.py:170: UserWarning: __floordiv__ is deprecated, and its behavior will change in a future version of pytorch. It currently rounds toward 0 (like the 'trunc' function NOT 'floor'). This results in incorrect rounding for negative values. To keep the current behavior, use torch.div(a, b, rounding_mode='trunc'), or for actual floor division, use torch.div(a, b, rounding_mode='floor').\n",
      "  x = windows.view(B, H // window_size, W // window_size, window_size, window_size, -1)\n",
      "/Users/raoulritter/STB-VMM/models/model.py:245: UserWarning: __floordiv__ is deprecated, and its behavior will change in a future version of pytorch. It currently rounds toward 0 (like the 'trunc' function NOT 'floor'). This results in incorrect rounding for negative values. To keep the current behavior, use torch.div(a, b, rounding_mode='trunc'), or for actual floor division, use torch.div(a, b, rounding_mode='floor').\n",
      "  attn = attn.view(B_ // nW, nW, self.num_heads, N, N) + mask.unsqueeze(1).unsqueeze(0)\n"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[2], line 34\u001b[0m\n\u001b[1;32m     20\u001b[0m model\u001b[39m.\u001b[39meval()\n\u001b[1;32m     21\u001b[0m \u001b[39m# model.pt¸\u001b[39;00m\n\u001b[1;32m     22\u001b[0m \n\u001b[1;32m     23\u001b[0m \u001b[39m# # Export the model to ONNX\u001b[39;00m\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m     32\u001b[0m \u001b[39m#                   export_params=True,  # store the trained parameter weights inside the model file\u001b[39;00m\n\u001b[1;32m     33\u001b[0m \u001b[39m# Export the model to ONNX\u001b[39;00m\n\u001b[0;32m---> 34\u001b[0m torch\u001b[39m.\u001b[39;49monnx\u001b[39m.\u001b[39;49mexport(model, \n\u001b[1;32m     35\u001b[0m                   (a_dummy,b_dummy, amp_dummy), \n\u001b[1;32m     36\u001b[0m                   \u001b[39m'\u001b[39;49m\u001b[39mmodel.onnx\u001b[39;49m\u001b[39m'\u001b[39;49m,\n\u001b[1;32m     37\u001b[0m \n\u001b[1;32m     38\u001b[0m                   opset_version\u001b[39m=\u001b[39;49m\u001b[39m11\u001b[39;49m)\n",
      "File \u001b[0;32m/opt/homebrew/Caskroom/miniconda/base/envs/testconv/lib/python3.9/site-packages/torch/onnx/__init__.py:305\u001b[0m, in \u001b[0;36mexport\u001b[0;34m(model, args, f, export_params, verbose, training, input_names, output_names, operator_export_type, opset_version, do_constant_folding, dynamic_axes, keep_initializers_as_inputs, custom_opsets, export_modules_as_functions)\u001b[0m\n\u001b[1;32m     39\u001b[0m \u001b[39m\u001b[39m\u001b[39mr\u001b[39m\u001b[39m\"\"\"\u001b[39;00m\n\u001b[1;32m     40\u001b[0m \u001b[39mExports a model into ONNX format. If ``model`` is not a\u001b[39;00m\n\u001b[1;32m     41\u001b[0m \u001b[39m:class:`torch.jit.ScriptModule` nor a :class:`torch.jit.ScriptFunction`, this runs\u001b[39;00m\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m    301\u001b[0m \u001b[39m    model to the file ``f`` even if this is raised.\u001b[39;00m\n\u001b[1;32m    302\u001b[0m \u001b[39m\"\"\"\u001b[39;00m\n\u001b[1;32m    304\u001b[0m \u001b[39mfrom\u001b[39;00m \u001b[39mtorch\u001b[39;00m\u001b[39m.\u001b[39;00m\u001b[39monnx\u001b[39;00m \u001b[39mimport\u001b[39;00m utils\n\u001b[0;32m--> 305\u001b[0m \u001b[39mreturn\u001b[39;00m utils\u001b[39m.\u001b[39;49mexport(model, args, f, export_params, verbose, training,\n\u001b[1;32m    306\u001b[0m                     input_names, output_names, operator_export_type, opset_version,\n\u001b[1;32m    307\u001b[0m                     do_constant_folding, dynamic_axes,\n\u001b[1;32m    308\u001b[0m                     keep_initializers_as_inputs, custom_opsets,\n\u001b[1;32m    309\u001b[0m                     export_modules_as_functions)\n",
      "File \u001b[0;32m/opt/homebrew/Caskroom/miniconda/base/envs/testconv/lib/python3.9/site-packages/torch/onnx/utils.py:118\u001b[0m, in \u001b[0;36mexport\u001b[0;34m(model, args, f, export_params, verbose, training, input_names, output_names, operator_export_type, opset_version, do_constant_folding, dynamic_axes, keep_initializers_as_inputs, custom_opsets, export_modules_as_functions)\u001b[0m\n\u001b[1;32m    115\u001b[0m     \u001b[39melse\u001b[39;00m:\n\u001b[1;32m    116\u001b[0m         operator_export_type \u001b[39m=\u001b[39m OperatorExportTypes\u001b[39m.\u001b[39mONNX\n\u001b[0;32m--> 118\u001b[0m _export(model, args, f, export_params, verbose, training, input_names, output_names,\n\u001b[1;32m    119\u001b[0m         operator_export_type\u001b[39m=\u001b[39;49moperator_export_type, opset_version\u001b[39m=\u001b[39;49mopset_version,\n\u001b[1;32m    120\u001b[0m         do_constant_folding\u001b[39m=\u001b[39;49mdo_constant_folding, dynamic_axes\u001b[39m=\u001b[39;49mdynamic_axes,\n\u001b[1;32m    121\u001b[0m         keep_initializers_as_inputs\u001b[39m=\u001b[39;49mkeep_initializers_as_inputs,\n\u001b[1;32m    122\u001b[0m         custom_opsets\u001b[39m=\u001b[39;49mcustom_opsets, export_modules_as_functions\u001b[39m=\u001b[39;49mexport_modules_as_functions)\n",
      "File \u001b[0;32m/opt/homebrew/Caskroom/miniconda/base/envs/testconv/lib/python3.9/site-packages/torch/onnx/utils.py:719\u001b[0m, in \u001b[0;36m_export\u001b[0;34m(model, args, f, export_params, verbose, training, input_names, output_names, operator_export_type, export_type, opset_version, do_constant_folding, dynamic_axes, keep_initializers_as_inputs, fixed_batch_size, custom_opsets, add_node_names, onnx_shape_inference, export_modules_as_functions)\u001b[0m\n\u001b[1;32m    715\u001b[0m     dynamic_axes \u001b[39m=\u001b[39m {}\n\u001b[1;32m    716\u001b[0m _validate_dynamic_axes(dynamic_axes, model, input_names, output_names)\n\u001b[1;32m    718\u001b[0m graph, params_dict, torch_out \u001b[39m=\u001b[39m \\\n\u001b[0;32m--> 719\u001b[0m     _model_to_graph(model, args, verbose, input_names,\n\u001b[1;32m    720\u001b[0m                     output_names, operator_export_type,\n\u001b[1;32m    721\u001b[0m                     val_do_constant_folding,\n\u001b[1;32m    722\u001b[0m                     fixed_batch_size\u001b[39m=\u001b[39;49mfixed_batch_size,\n\u001b[1;32m    723\u001b[0m                     training\u001b[39m=\u001b[39;49mtraining,\n\u001b[1;32m    724\u001b[0m                     dynamic_axes\u001b[39m=\u001b[39;49mdynamic_axes)\n\u001b[1;32m    726\u001b[0m \u001b[39m# TODO: Don't allocate a in-memory string for the protobuf\u001b[39;00m\n\u001b[1;32m    727\u001b[0m defer_weight_export \u001b[39m=\u001b[39m export_type \u001b[39mis\u001b[39;00m \u001b[39mnot\u001b[39;00m ExportTypes\u001b[39m.\u001b[39mPROTOBUF_FILE\n",
      "File \u001b[0;32m/opt/homebrew/Caskroom/miniconda/base/envs/testconv/lib/python3.9/site-packages/torch/onnx/utils.py:503\u001b[0m, in \u001b[0;36m_model_to_graph\u001b[0;34m(model, args, verbose, input_names, output_names, operator_export_type, do_constant_folding, _disable_torch_constant_prop, fixed_batch_size, training, dynamic_axes)\u001b[0m\n\u001b[1;32m    499\u001b[0m graph, params, torch_out, module \u001b[39m=\u001b[39m _create_jit_graph(model, args)\n\u001b[1;32m    501\u001b[0m params_dict \u001b[39m=\u001b[39m _get_named_param_dict(graph, params)\n\u001b[0;32m--> 503\u001b[0m graph \u001b[39m=\u001b[39m _optimize_graph(graph, operator_export_type,\n\u001b[1;32m    504\u001b[0m                         _disable_torch_constant_prop\u001b[39m=\u001b[39;49m_disable_torch_constant_prop,\n\u001b[1;32m    505\u001b[0m                         fixed_batch_size\u001b[39m=\u001b[39;49mfixed_batch_size, params_dict\u001b[39m=\u001b[39;49mparams_dict,\n\u001b[1;32m    506\u001b[0m                         dynamic_axes\u001b[39m=\u001b[39;49mdynamic_axes, input_names\u001b[39m=\u001b[39;49minput_names,\n\u001b[1;32m    507\u001b[0m                         module\u001b[39m=\u001b[39;49mmodule)\n\u001b[1;32m    508\u001b[0m \u001b[39mfrom\u001b[39;00m \u001b[39mtorch\u001b[39;00m\u001b[39m.\u001b[39;00m\u001b[39monnx\u001b[39;00m\u001b[39m.\u001b[39;00m\u001b[39msymbolic_helper\u001b[39;00m \u001b[39mimport\u001b[39;00m _onnx_shape_inference\n\u001b[1;32m    509\u001b[0m \u001b[39mif\u001b[39;00m \u001b[39misinstance\u001b[39m(model, torch\u001b[39m.\u001b[39mjit\u001b[39m.\u001b[39mScriptModule) \u001b[39mor\u001b[39;00m \u001b[39misinstance\u001b[39m(model, torch\u001b[39m.\u001b[39mjit\u001b[39m.\u001b[39mScriptFunction):\n",
      "File \u001b[0;32m/opt/homebrew/Caskroom/miniconda/base/envs/testconv/lib/python3.9/site-packages/torch/onnx/utils.py:232\u001b[0m, in \u001b[0;36m_optimize_graph\u001b[0;34m(graph, operator_export_type, _disable_torch_constant_prop, fixed_batch_size, params_dict, dynamic_axes, input_names, module)\u001b[0m\n\u001b[1;32m    230\u001b[0m     torch\u001b[39m.\u001b[39m_C\u001b[39m.\u001b[39m_jit_pass_onnx_set_dynamic_input_shape(graph, dynamic_axes, input_names)\n\u001b[1;32m    231\u001b[0m torch\u001b[39m.\u001b[39m_C\u001b[39m.\u001b[39m_jit_pass_onnx_lint(graph)\n\u001b[0;32m--> 232\u001b[0m graph \u001b[39m=\u001b[39m torch\u001b[39m.\u001b[39;49m_C\u001b[39m.\u001b[39;49m_jit_pass_onnx(graph, operator_export_type)\n\u001b[1;32m    233\u001b[0m torch\u001b[39m.\u001b[39m_C\u001b[39m.\u001b[39m_jit_pass_onnx_lint(graph)\n\u001b[1;32m    234\u001b[0m torch\u001b[39m.\u001b[39m_C\u001b[39m.\u001b[39m_jit_pass_lint(graph)\n",
      "File \u001b[0;32m/opt/homebrew/Caskroom/miniconda/base/envs/testconv/lib/python3.9/site-packages/torch/onnx/__init__.py:354\u001b[0m, in \u001b[0;36m_run_symbolic_function\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m    352\u001b[0m \u001b[39mdef\u001b[39;00m \u001b[39m_run_symbolic_function\u001b[39m(\u001b[39m*\u001b[39margs, \u001b[39m*\u001b[39m\u001b[39m*\u001b[39mkwargs):\n\u001b[1;32m    353\u001b[0m     \u001b[39mfrom\u001b[39;00m \u001b[39mtorch\u001b[39;00m\u001b[39m.\u001b[39;00m\u001b[39monnx\u001b[39;00m \u001b[39mimport\u001b[39;00m utils\n\u001b[0;32m--> 354\u001b[0m     \u001b[39mreturn\u001b[39;00m utils\u001b[39m.\u001b[39;49m_run_symbolic_function(\u001b[39m*\u001b[39;49margs, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mkwargs)\n",
      "File \u001b[0;32m/opt/homebrew/Caskroom/miniconda/base/envs/testconv/lib/python3.9/site-packages/torch/onnx/utils.py:1066\u001b[0m, in \u001b[0;36m_run_symbolic_function\u001b[0;34m(g, block, n, inputs, env, operator_export_type)\u001b[0m\n\u001b[1;32m   1064\u001b[0m \u001b[39mif\u001b[39;00m op_name \u001b[39m==\u001b[39m \u001b[39m\"\u001b[39m\u001b[39mConstant\u001b[39m\u001b[39m\"\u001b[39m \u001b[39mand\u001b[39;00m \u001b[39mnot\u001b[39;00m n\u001b[39m.\u001b[39mmustBeNone():\n\u001b[1;32m   1065\u001b[0m     \u001b[39mif\u001b[39;00m n\u001b[39m.\u001b[39mkindOf(\u001b[39m\"\u001b[39m\u001b[39mvalue\u001b[39m\u001b[39m\"\u001b[39m) \u001b[39m==\u001b[39m \u001b[39m\"\u001b[39m\u001b[39mt\u001b[39m\u001b[39m\"\u001b[39m:\n\u001b[0;32m-> 1066\u001b[0m         \u001b[39mreturn\u001b[39;00m g\u001b[39m.\u001b[39;49mop(\u001b[39m\"\u001b[39;49m\u001b[39mConstant\u001b[39;49m\u001b[39m\"\u001b[39;49m, value_t\u001b[39m=\u001b[39;49mn[\u001b[39m\"\u001b[39;49m\u001b[39mvalue\u001b[39;49m\u001b[39m\"\u001b[39;49m])\n\u001b[1;32m   1067\u001b[0m     \u001b[39mif\u001b[39;00m n\u001b[39m.\u001b[39mkindOf(\u001b[39m\"\u001b[39m\u001b[39mvalue\u001b[39m\u001b[39m\"\u001b[39m) \u001b[39m==\u001b[39m \u001b[39m\"\u001b[39m\u001b[39ms\u001b[39m\u001b[39m\"\u001b[39m:\n\u001b[1;32m   1068\u001b[0m         \u001b[39mreturn\u001b[39;00m g\u001b[39m.\u001b[39mop(\u001b[39m\"\u001b[39m\u001b[39mConstant\u001b[39m\u001b[39m\"\u001b[39m, value_s\u001b[39m=\u001b[39mn[\u001b[39m\"\u001b[39m\u001b[39mvalue\u001b[39m\u001b[39m\"\u001b[39m])\n",
      "File \u001b[0;32m/opt/homebrew/Caskroom/miniconda/base/envs/testconv/lib/python3.9/site-packages/torch/onnx/utils.py:957\u001b[0m, in \u001b[0;36m_graph_op\u001b[0;34m(g, opname, *raw_args, **kwargs)\u001b[0m\n\u001b[1;32m    955\u001b[0m \u001b[39mif\u001b[39;00m _onnx_shape_inference:\n\u001b[1;32m    956\u001b[0m     \u001b[39mfrom\u001b[39;00m \u001b[39mtorch\u001b[39;00m\u001b[39m.\u001b[39;00m\u001b[39monnx\u001b[39;00m\u001b[39m.\u001b[39;00m\u001b[39msymbolic_helper\u001b[39;00m \u001b[39mimport\u001b[39;00m _export_onnx_opset_version \u001b[39mas\u001b[39;00m opset_version\n\u001b[0;32m--> 957\u001b[0m     torch\u001b[39m.\u001b[39;49m_C\u001b[39m.\u001b[39;49m_jit_pass_onnx_node_shape_type_inference(n, _params_dict, opset_version)\n\u001b[1;32m    959\u001b[0m \u001b[39mif\u001b[39;00m outputs \u001b[39m==\u001b[39m \u001b[39m1\u001b[39m:\n\u001b[1;32m    960\u001b[0m     \u001b[39mreturn\u001b[39;00m n\u001b[39m.\u001b[39moutput()\n",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "import torch\n",
    "from models.model import STBVMM  \n",
    "\n",
    "# Define dummy inputs of the appropriate sizes for your model\n",
    "\n",
    "a_dummy = torch.randn(1, 3, 384,384)\n",
    "b_dummy = torch.randn(1, 3, 384, 384)\n",
    "# c_dummy = torch.randn(1, 3, 384, 384)  # optional during training\n",
    "\n",
    "# For the amplitude parameter amp\n",
    "amp_dummy = torch.randn(5)  # adjusted to be a scalar value\n",
    "\n",
    "# Instantiate your model\n",
    "model = STBVMM()\n",
    "\n",
    "# Load the weights from the checkpoint\n",
    "checkpoint = torch.load('ckpt/ckpt_e49.pth.tar')\n",
    "model.load_state_dict(checkpoint['state_dict'], strict= False)\n",
    "# Make sure the model is in evaluation mode\n",
    "model.eval()\n",
    "# model.pt¸\n",
    "\n",
    "# # Export the model to ONNX\n",
    "# torch.onnx.export(model, \n",
    "#                   (a_dummy, amp_dummy), \n",
    "#                   opset_version=10,\n",
    "#                   'model.onnx')\n",
    "\n",
    "# torch.onnx.export(model,               # model being run\n",
    "#                   example_input,       # model input (or a tuple for multiple inputs)\n",
    "#                   \"model.onnx\",        # where to save the model (can be a file or file-like object)\n",
    "#                   export_params=True,  # store the trained parameter weights inside the model file\n",
    "# Export the model to ONNX\n",
    "torch.onnx.export(model, \n",
    "                  (a_dummy,b_dummy, amp_dummy), \n",
    "                  'model.onnx',\n",
    "                  opset_version=11)\n",
    "\n"
   ]
  }
 ],
 "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.9.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}