transPose3d.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import pickle\n",
    "os.environ[\"KMP_DUPLICATE_LIB_OK\"]=\"TRUE\"\n",
    "import numpy as np\n",
    "import torch\n",
    "from torch import nn\n",
    "from torch import optim\n",
    "from torch.utils.data import DataLoader, Dataset, random_split\n",
    "from tqdm.notebook import tqdm\n",
    "import torch.nn.functional as F\n",
    "import itertools\n",
    "from model import *\n",
    "from dataset_loader import ActionDatasets, create_data_loader\n",
    "from loss_func import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "cuda\n"
     ]
    }
   ],
   "source": [
    "lr = 0.001 # 学习率\n",
    "epoches = 1000 #训练次数\n",
    "batch_size = 256 # 每一个训练批次数量\n",
    "start_epoch = 0\n",
    "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    "print(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "poseS1 = PoseS1().to(device)\n",
    "poseS2 = PoseS2().to(device)\n",
    "poseS3 = PoseS3().to(device)\n",
    "transB1 = TransB1().to(device)\n",
    "transB2 = TransB2().to(device)\n",
    "\n",
    "Ls1 = poseLoss # 叶关节点位置的loss\n",
    "Ls2 = poseLoss # 除胯外的关节位置的loss\n",
    "Ls3 = poseLoss # 除胯外的关节6d旋转的loss\n",
    "Lb1 = crossEntropy # 接触脚的概率\n",
    "Lb2 = ver_n_loss # 连续1,3,9,27帧的位移loss\n",
    "\n",
    "optimizer = optim.Adam(filter(lambda p: p.requires_grad,  itertools.chain(poseS1.parameters(),\n",
    "                                       poseS2.parameters(),\n",
    "                                       poseS3.parameters(),\n",
    "                                       transB1.parameters(),\n",
    "                                       transB2.parameters())), lr=lr)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_checkpoint(filepath):\n",
    "    checkpoint = torch.load(filepath)\n",
    "    global start_epoch\n",
    "    start_epoch = checkpoint['epoch']\n",
    "    models = checkpoint['modes'] # 提取网络结构\n",
    "    models_state_dict = checkpoint['models_state_dict']  # 提取网络结构\n",
    "    for model, state_dict in zip([poseS1,  poseS2, poseS3, transB1, transB2], models_state_dict):\n",
    "        model.load_state_dict(state_dict)  # 加载网络权重参数\n",
    "    optimizer.load_state_dict(checkpoint['optimizer_state_dict'])  # 加载优化器参数\n",
    "    \n",
    "    for parameter in model.parameters():\n",
    "        parameter.requires_grad = True\n",
    "    model.train()\n",
    "    \n",
    "    return models\n",
    "models = load_checkpoint(\"./checkpoint/checkpoint_549_pose1.pth\")\n",
    "for p in poseS1.parameters():\n",
    "    p.requires_grad=False\n",
    "for p in poseS2.parameters():\n",
    "    p.requires_grad=False\n",
    "# for p in poseS3.parameters():\n",
    "#     p.requires_grad=False\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    },
    "tags": []
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "训练集大小36858， 验证集大小9215， 测试集大小46073\n"
     ]
    }
   ],
   "source": [
    "train_loader, test_loader, valid_loader = create_data_loader(batch_size)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "pycharm": {
     "is_executing": true,
     "name": "#%%\n"
    },
    "scrolled": true,
    "tags": []
   },
   "outputs": [
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "80263b5b294e4a92ac7546dcdc549fd8"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "5dece621e6ef4892aec525ebcf2c3295"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "fa86f046f4384254a856a01bee1131b4"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "9f19f38fd9da4fe2aad3cd20c35ec426"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "6df8e080d10f4ed3b2413b43d0d2f7e1"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "abb63a9050d7496581b80b804932571c"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "d6ab1ae77e6d4ee5b9b29256655a7730"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "dc0660aff0a1498ba5fe2bcdf2233700"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "830d988c5eb24e13bd36fe81a71c17ff"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "6f1af5d544374fd8997fc95ec6be0029"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "96eb41d44697488cae2e97ac27507cc9"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "8bdc241fe3684073995638a4e0456167"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "9ea8f5277d0a4c78becc2f96d71ae940"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "05739e8ccaf94bc5ade63d4d50505dea"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "98649de1b7c44fc9a59f9bd884e86652"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "db30308336084b2bbaca23b5722e0367"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "0a70c5f0e46c4616bbf876c63c6cbb68"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "7eeab00e69c14849a924fb3d1e2c824b"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "181a9bd40493454f9e17dd3888e9f56b"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "4a698f7a8f1f4ee6b9e4b089b1a14f20"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "5ffb524eed0541db8ea4da0072fdd5d2"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "24232925925742e785fea47a9e1a352a"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "93489defee48476383950b33dd05c49b"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "c14e03ea7d674668b37a08b7e5b00e14"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "60b8dc8e06f74c6ba3eaef96a4f8ae7d"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "401d5692f70045d79ac4ea7264db5ba5"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "7d36751a757349ebab0a4f738b592438"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "ad1d92c42b174b89830c770ee1864ef0"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "ca6d51247933432b94ffbd2104ed2016"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "3e042f87615a4a38949d508e64a918ea"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "fad9def76ae04c02bfe16e6e29b990f5"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "f704431e85084b379aab3a252382e03d"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "04143ee8ccd647eeaba6f988a7d92542"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "ed806dacbf7d4c0dae0ea831aa78fdf7"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "299ee6b20aa34254b83924b8289cba27"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "0027bb6decaa42f3ad667a2f43dbd726"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "5a0071c97a234e909053b86df76091db"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "68a14757cc2a49a9879e31a0e671e4f2"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "216554affa144b1098c1cc38770da104"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "39080ed1b1fd4926a507d348463a4388"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "6984410057ed4e7dbe46d50a27f0acd3"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "7a52c5d787bf4e5a98e66b0a4bb6598b"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "bcdb80d529154e158f664ad722b4b4ea"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "3b77f4e576e9413f972f1583414c626e"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "b81e171450c042fdbe64f4688f2bbfaa"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "07e4fba0e474469e95ae5ed8f5cfaa15"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "8f10d383853840fab3173f56701304e0"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "ad8e765914a84f99a085c0d23184ae3c"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "0c1dfcab62bc4501b17d830065a6a860"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "dcf21d8bbe834d7b8ad4fa8b7d0ba351"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "8b4a8fa4ba4f4aa4925190d740d577bb"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "66efd202cde547aa8210ffb6035bcdd8"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "362517d387814395a694226ae1bbbbeb"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "2ca78a754ffc4d66b73edae5c79a3c56"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "5884d2692ef74a20ba638ccc9a2c1bec"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "00b4d5a4caa34948b25a083c8e659664"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "db090239d4224cde9ba528d40eecc6ed"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "3529a024e4de4451bd127c9e0e1c4625"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "b3a8b4e00d0847b18023c45dd2190472"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "c1315d9091634b389ee94e132e297e70"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "53fd5aa6fb8949f98008d5da41778d08"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "b5d5920024b748faadba2da7253ca9c5"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "4a837412641b4e9e98a81036a110d690"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "dae451832620415ba57e09c2d4514110"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "fc8dd8804fd94772aaf9f37d738bd345"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "301b87aeb70c453b930f078eaf891507"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "51ca7bc65e7e41acb5c7088d1ea77dbe"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "96ddca53cd844fa698734f48cc1d3f95"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "8b4fc45c4bfc4cd094e7a92c1afcc84b"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "a6d6dab958a2433f913cfe78954f703b"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "9bf46feb064b454483d1c25e936c4caf"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "e05dcf0207384b158f2c1d6574a8e652"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "afc96c98935b41279442989a899dea7a"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "762852fc9cba45bcad6f3317a67f74b6"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "9ad50053e0b649efa3f742b3e1f87642"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "6eff367ae22345628082111d262267d3"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "7f33335eed214d1f9624240c55ebec4f"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "817ed7836c8847e5acbe7b08f4e81534"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "ac5519936c2e495ba36d26d11839d967"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "a85d931a3c8047988e070592d338a689"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "dd0bb28d2b19441a901732e4336ebf74"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "8346b8b09c4f4691aaca4aabfc3e4d37"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "201a5e1433e64c719944886584c813b2"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "474c410a9f194d638a58f36677b735df"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "6fc9cb6a0e984e86b37db88570317f7a"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "8b20f8c46548458a88fb07ddbd911c09"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "bc1c1c6e6e6a43d08ec6b819ef3f3d91"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "5a06f289789a4fc8b42ec7866f80bca4"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "86960b566eb34c959085ae402b616fe8"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "8ed9731ec5844f1b88c533c803728f96"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "59309f2dffb9472984a9966dc092f616"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "f8100621c84e4873af17042b725ddd3e"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "f98f8312b55e496b8eda1e56b7a04e5b"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "047b5a179db84b6085a4476dcbb71e54"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "e131655a446f40a092a9ad685c98ad64"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "c0d17f236aa54c138abb61acad76fac6"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "34e8d52b5c2149cb971c91ff6473aa8e"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "481f6bae62074455b7d3abe0f2ef1a3d"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "74f8a5c6024843ceabb2d80c4fda7ec8"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "cbccc89ec29d452284490c4a15c422f9"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "cc6d11b5c77d4651b517f7669ac6d26e"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "07223a4dc51a4884880c6111d7357a6f"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "a34c4ca85926480f8d6287e2b871b066"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "d23aab56b1f643829530b2bd0bbf5f20"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "6a5720247d264040a05c64589d4c9960"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "f5d39eb89ea24a0b8b5207ca6068fb28"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "30b66b6815a445fbb7dc7298009f1a74"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "de6cc510bda141c3b82b4e10b5bee15c"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "1b01ea98383c46bdac409df3ad14dc78"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "f75a9bc3b20a46bfb1da77f029caa3e0"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "c1a03215c8cd4826acd482fa7a1e3d3b"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "bb3c24179afb42b1961540f43fa02e36"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "9417bd4c782c41ab89bad7d0db06eb91"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "f28b853e7f8a46d68ab97ca9364fcbc2"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "124719599f544781ba2fa5b39c3b0a6b"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "b8c67e8deb534091948424da7817612c"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "1fc83464758c44e9887de3a4b4a6cf07"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "cb9a921ae30e4199982721ff6ccff01e"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "c2c7d2ac379c46b78ecf7c4ae68f53eb"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "2997f6dd3f0a40d89c25bffbdf7d841d"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "78ea38bac00043529f79e83aaeee3eb3"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "4fc3a47817024bf6b32332b1e1e6b4b7"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "8514bd1e2b1c42718415c0b300078e62"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "2476905a59994873bf9b198e7342641b"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "9f4e5405870e4d848489796874e11c56"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "e83d2a5d182d458abd33b8854d7770cc"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "fcb4c8d293f14439b3a80a0fa8690d6c"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "75d805e7ea1746048829bce72f689a6f"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "175adcba49e549f5b2a2b7d2d4ca8fdd"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "40612bdeda794c6aae6b079fc81fa61c"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "2668a983fac849a391fc361357e1e252"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "fffdfc9d830643b085ba23191e24d11b"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "b1d14bbacf3a4c019a7cc9bf984dc7b6"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "2ae03d33e32045c3a3f76e54749c8c1e"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "91b925a695e240639fc22aacb8261238"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "d443b4f91d6046f6822dcb1ef4eb7f72"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "5e0da850a9594cd78b3bb11e9dedf417"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "f0412c04d6c54052ab12f717079e5bc6"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "a3895b915fb247e78db847af7509a1ac"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "4a9e5b736a734f96aa098db2e18aded8"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "84b3905cc733492792508c6ca414add0"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "df8c1c70c85741af81b451faee6adff3"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "98e560080e0b46b79c0bccade827b9e9"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "f12d99726e2b49019426d42a4293df8a"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "3936de775e0649bdbadce5a5c0401991"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "4a89d4234e214d3ab4f888d134060959"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "de560e02c97d47afb04a346b01ad9076"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "6d3aa443e35f435ca0bf54bd340459b1"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "34e8ab7fd98d4682b6f61ce7ac7f5656"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "2c862a91a7b347b7b7a7a86125be38f7"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "78a4b4f5a9de4aca859bbf8f7bc72b01"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "f8b1024c599d4e2184afefb750d099e9"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "c409f28c24264952add971208f06c573"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "20dcf4cdb4df46c7903f5fb83cfcc43c"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "a8a5be2f45ed48d4864136afc278506a"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "e95925e91c2d45aaa460e009a365f064"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "dd35de13462f489a95f5fbf40b75ab2c"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "6c167c949739467a9987a8a5e8de24fd"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "9f19f5687e474fe59c75fef072dd6897"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "89c638ce841948adb7e463b46b54dfdc"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "8f2e51e545344b1eb6e24f34f99d9c78"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "6c79fc76f53641f594e561d85165b8da"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "de7cd9a92dcc42bfabca8ccfa358df19"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "2e391c055d68449388119d8e31d5edba"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "672af6a2466a432d916eeea18b9ff621"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "5793b43ff90646fd8dc72f7c48669bab"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "bba633d5b0ab43279292814dba455d58"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "3c22fd1722724740832ef39eb328aa38"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "ce05c1a5fcf14582bde9ca9962bc27a0"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "262792eb55a5450f8eaaaa7f933d3488"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "28451da437cb428fb9b105fe05e720cb"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "1e7cf83a6042473781eb917d7c6dc8f9"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "b472144effa746f58d5aee6ce3d0fd72"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "627cffa4924a45b1a779c220e91fbbf7"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "0083e15eb2ef435b8986830df767bf6a"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "02d3c18e6942425ca461b45586f31c15"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "9390cb9304e04c83beccd96db0fd6470"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "f2d395b636334a5aaff814da3375b842"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "d84d1b9a1d274204802416cb58670251"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "4b0153711b6b4d1b996f647c57a4d216"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "c1843997f2464ac490d0cfd25277ee42"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "60b1ac71f2c54229aa334f34bd0da19c"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "d821cce59bdf45eaaec33dba067517ef"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "9315bf684ad34e289fdcefcd940c2f88"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "  0%|          | 0/143 [00:00<?, ?it/s, training]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "f22b245b33b94a4ca78c4d1f6337d30e"
      }
     },
     "metadata": {}
    }
   ],
   "source": [
    "for epoch in range(start_epoch, epoches):\n",
    "    i = 0\n",
    "    loss_sum = 0\n",
    "    loss_dict = {\"loss_ls1\":0, \"loss_ls2\":0, \"loss_ls3\":0, \"loss_lb1\":0, \"loss_lb2\":0, \"foot_acc\":0}\n",
    "    bar = tqdm(enumerate(iter(train_loader)),postfix=\"training\", total =len(train_loader))\n",
    "\n",
    "    for idx, (seq_len, x0, p_leaf_gt, p_all_gt, pose_6d_gt, support_leg_gt, root_velocity_gt, root_ori, mask) in bar:\n",
    "        poseS1.train()\n",
    "        poseS2.train()\n",
    "        poseS3.train()\n",
    "        transB1.train()\n",
    "        transB2.train()\n",
    "\n",
    "        x0 = x0.to(device)\n",
    "        p_leaf_gt = p_leaf_gt.to(device)\n",
    "        p_all_gt = p_all_gt.to(device)\n",
    "        pose_6d_gt = pose_6d_gt.to(device)\n",
    "        support_leg_gt = support_leg_gt.to(device)\n",
    "        root_velocity_gt = root_velocity_gt.to(device)\n",
    "        root_ori = root_ori.to(device)\n",
    "        mask = mask.to(device)\n",
    "        \n",
    "\n",
    "        # --------pose 1\n",
    "        p_leaf = poseS1(x0, seq_len)\n",
    "        loss_ls1 = Ls1(p_leaf, p_leaf_gt.to(device) + torch.normal(mean=0, std=0.04, size=p_leaf_gt.shape).to(device), mask)\n",
    "        loss_dict[\"loss_ls1\"] += loss_ls1.item()\n",
    "        # --------pose 2\n",
    "        x1 = torch.cat([p_leaf, x0], dim=-1)\n",
    "        p_all = poseS2(x1, seq_len)\n",
    "        loss_ls2 = Ls2(p_all, p_all_gt.to(device) + torch.normal(mean=0, std=0.025, size=p_all_gt.shape).to(device), mask)\n",
    "        loss_dict[\"loss_ls2\"] += loss_ls2.item()\n",
    "        # --------pose 3\n",
    "        x2 = torch.cat([p_all, x0], dim=-1)\n",
    "        r6d_all = poseS3(x2, seq_len)\n",
    "        loss_ls3 = Ls3(r6d_all, pose_6d_gt, mask)\n",
    "        loss_dict[\"loss_ls3\"] += loss_ls3.item()\n",
    "        # --------transB1\n",
    "        support_leg_prob = transB1(x1, seq_len)\n",
    "\n",
    "        loss_lb1 = Lb1(support_leg_prob, support_leg_gt, mask)\n",
    "        foot_acc = foot_accuracy(support_leg_prob, support_leg_gt, mask)\n",
    "        loss_dict[\"loss_lb1\"] += loss_lb1.item()\n",
    "        loss_dict[\"foot_acc\"] += foot_acc.item() * 100\n",
    "        # --------transB2\n",
    "        ve_hat = transB2(x2 + torch.normal(mean=0, std=0.025, size=x2.shape).to(device), seq_len)\n",
    "        loss_lb2 = Lb2(ve_hat , root_velocity_gt , mask)\n",
    "        loss_dict[\"loss_lb2\"] += loss_lb2.item()\n",
    "        # --------end\n",
    "        \n",
    "        # cal angle error\n",
    "        #avg_error = compute_angle_dif(r6d_all.detach().cpu(), pose_6d_gt.detach().cpu(), mask.detach().cpu())\n",
    "        avg_error = 0\n",
    "        loss_total = loss_lb2\n",
    "        \n",
    "        optimizer.zero_grad()\n",
    "        loss_sum += loss_total.item()\n",
    "        loss_total.backward()\n",
    "        optimizer.step()\n",
    "        i+=1\n",
    "        bar.set_description(f\"[Epoch {epoch+1}/ {epoches}]\")\n",
    "        bar.set_postfix(total_loss= loss_sum / i, avg_angle_error=avg_error,**{key: value / i for key, value in loss_dict.items()})\n",
    "\n",
    "        if idx%30 == 0:\n",
    "            checkpoint = {'modes': [PoseS1(), PoseS2(),PoseS3(),TransB1(),TransB2()],  \n",
    "    'models_state_dict': [poseS1.state_dict(),poseS2.state_dict(), poseS3.state_dict(),transB1.state_dict(),transB2.state_dict()],\n",
    "    'optimizer_state_dict': optimizer.state_dict(),'epoch': epoch}\n",
    "            torch.save(checkpoint, f'checkpoint/checkpoint_latest.pth')\n",
    "            \n",
    "            # get sample\n",
    "            sample_idx = 0\n",
    "            v_gt = root_velocity_gt[sample_idx][mask[sample_idx].int().bool()].reshape(-1, 3) / 60\n",
    "            v_pre = ve_hat[sample_idx][mask[sample_idx].int().bool()].reshape(-1, 3)\n",
    "            foot_pre = support_leg_prob[sample_idx][mask[sample_idx].int().bool()].reshape(-1, 2)\n",
    "            foot_gt = support_leg_gt[sample_idx][mask[sample_idx].int().bool()].reshape(-1, 2)\n",
    "\n",
    "        # print(v_gt)\n",
    "            sample_dict = {\"pose\": r6d_all[sample_idx][mask[sample_idx].int().bool()].detach().cpu().numpy(), \n",
    "                                            'pose_gt':pose_6d_gt[sample_idx][mask[sample_idx].int().bool()].detach().cpu().numpy(), \n",
    "                                            'leg': support_leg_prob[sample_idx][mask[sample_idx].int().bool()].detach().cpu().numpy(),\n",
    "                                            'leg_gt': support_leg_gt[sample_idx][mask[sample_idx].int().bool()].detach().cpu().numpy(),\n",
    "                                        'foot_pre': foot_pre.detach().cpu().numpy(),\n",
    "                                        'foot_gt': foot_gt.detach().cpu().numpy(),\n",
    "                                        'v': v_pre.detach().cpu().numpy(),\n",
    "                                            'v_gt': v_gt.detach().cpu().numpy(),\n",
    "                                            'root_ori':root_ori[sample_idx][mask[sample_idx].int().bool()].detach().cpu().numpy()\n",
    "                                            }\n",
    "            pickle.dump(sample_dict, open(f'sample/sample_data_{epoch}_{idx}.pkl', 'wb'))\n",
    "\n",
    "\n",
    "    checkpoint = {'modes': [PoseS1(), PoseS2(),PoseS3(),TransB1(),TransB2()],  \n",
    "    'models_state_dict': [poseS1.state_dict(),poseS2.state_dict(), poseS3.state_dict(),transB1.state_dict(),transB2.state_dict()],\n",
    "    'optimizer_state_dict': optimizer.state_dict(),'epoch': epoch}\n",
    "    torch.save(checkpoint, f'checkpoint/checkpoint_{epoch+1}_pose1.pth')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.8.10 64-bit",
   "language": "python",
   "name": "python3810jvsc74a57bd0e1617654b9a844d1d5c3f39905eff0ee5c78cb466aa242fb2d2abfff0dae271b"
  },
  "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.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}