train-mgif.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install dgl-cu90 "
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "!nvidia-smi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "#train.py\n",
    "from tqdm import trange\n",
    "import torch\n",
    "\n",
    "from torch.utils.data import DataLoader\n",
    "\n",
    "from logger import Logger\n",
    "from modules.model_m import GeneratorFullModel, DiscriminatorFullModel\n",
    "\n",
    "from torch.optim.lr_scheduler import MultiStepLR\n",
    "\n",
    "from sync_batchnorm import DataParallelWithCallback\n",
    "\n",
    "from frames_dataset import DatasetRepeater"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "#run.py \n",
    "import matplotlib\n",
    "\n",
    "matplotlib.use('Agg')\n",
    "\n",
    "import os, sys\n",
    "import yaml\n",
    "from argparse import ArgumentParser\n",
    "from time import gmtime, strftime\n",
    "from shutil import copy\n",
    "\n",
    "from frames_dataset import FramesDataset\n",
    "\n",
    "from modules.generator import OcclusionAwareGenerator\n",
    "from modules.discriminator import MultiScaleDiscriminator\n",
    "from modules.keypoint_detector import KPDetector\n",
    "\n",
    "import torch\n",
    "\n",
    "#from train import train\n",
    "from reconstruction import reconstruction\n",
    "from animate import animate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from modules.stn_a import STN"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Build Model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "由 run.py 改來"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "class opt:\n",
    "    config = \"config/mgif-256-m2.yaml\"\n",
    "    mode = \"train\"\n",
    "    log_dir = 'log'\n",
    "    checkpoint = None\n",
    "    device_ids = \"0\" #\"0,1,2,3\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "fn = lambda x: list(map(int, x.split(',')))\n",
    "opt.device_ids = fn(opt.device_ids)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/opt/conda/lib/python3.7/site-packages/ipykernel_launcher.py:2: YAMLLoadWarning: calling yaml.load() without Loader=... is deprecated, as the default Loader is unsafe. Please read https://msg.pyyaml.org/load for full details.\n",
      "  \n"
     ]
    }
   ],
   "source": [
    "with open(opt.config) as f:\n",
    "    config = yaml.load(f)\n",
    "\n",
    "if opt.checkpoint is not None:\n",
    "    log_dir = os.path.join(*os.path.split(opt.checkpoint)[:-1])\n",
    "else:\n",
    "    log_dir = os.path.join(opt.log_dir, os.path.basename(opt.config).split('.')[0])\n",
    "    log_dir += ' ' + strftime(\"%d_%m_%y_%H.%M.%S\", gmtime())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "generator = OcclusionAwareGenerator(**config['model_params']['generator_params'],\n",
    "                                    **config['model_params']['common_params'])\n",
    "discriminator = MultiScaleDiscriminator(**config['model_params']['discriminator_params'],\n",
    "                                        **config['model_params']['common_params'])\n",
    "kp_detector = KPDetector(**config['model_params']['kp_detector_params'],\n",
    "                         **config['model_params']['common_params'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 定義初始圖\n",
    "\"\"\"\n",
    "#-1:背景  ![](https://i.imgur.com/CnW6Uy8.gif)\n",
    "0:後半  ![](https://i.imgur.com/XGRK0ul.gif)\n",
    "1:外前腳 ![](https://i.imgur.com/7ngPynF.gif)\n",
    "2:耳朵  ![](https://i.imgur.com/fZjnGGQ.gif)\n",
    "3:頭上  ![](https://i.imgur.com/uGS8Eu0.gif)\n",
    "4:內後腳  ![](https://i.imgur.com/FGIHvpb.gif)\n",
    "5:鼻子  ![](https://i.imgur.com/NoDo30g.gif)\n",
    "6:內側腳(?)  ![](https://i.imgur.com/mA0hRHP.gif)\n",
    "7:外後腳  ![](https://i.imgur.com/AS3d08V.gif)\n",
    "8:軀幹  ![](https://i.imgur.com/Yw3K9qb.gif)\n",
    "9:尾巴 ![](https://i.imgur.com/DxFcCtj.gif)\n",
    "\"\"\"\n",
    "adjmatrix_directed = torch.zeros((10,10))+0.2 #[10, 10]\n",
    "select_edge =  [(3,2), (5,2)] #頭部\n",
    "select_edge += [(1,8), (6,8)] #前半\n",
    "select_edge += [(4,0), (7,0), (9,0)] #後半\n",
    "select_edge += [(8,2), (8,0)] #總體\n",
    "select_edge += [ (e[1],e[0]) for e in select_edge] #undirect\n",
    "select_edge += [ (i,i) for i in range(10)] #to self\n",
    "for e in select_edge:\n",
    "    adjmatrix_directed[e[0],e[1]] = 0.8"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(0)"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.sum(adjmatrix_directed>0.8)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'\\nbs = 5\\nx = torch.rand((bs,10,2))\\nkp_refiner.assign_mask_weight(adjmatrix_directed)\\ny, a, _ = kp_refiner(x, mask_type=\"soft\")\\nprint( y.shape, a.shape )\\nprint( kp_refiner.mask_weight )\\n'"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import torch\n",
    "from modules.graphattn_a1 import MultiHeadAttention\n",
    "kp_refiner = MultiHeadAttention(2,2,1)\n",
    "kp_refiner.assign_mask_weight(adjmatrix_directed)\n",
    "\n",
    "#try\n",
    "\"\"\"\n",
    "bs = 5\n",
    "x = torch.rand((bs,10,2))\n",
    "kp_refiner.assign_mask_weight(adjmatrix_directed)\n",
    "y, a, _ = kp_refiner(x, mask_type=\"soft\")\n",
    "print( y.shape, a.shape )\n",
    "print( kp_refiner.mask_weight )\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "#modified # added 20200609\n",
    "from modules.graphattn_a import GraphMHAttn,draw_graph \n",
    "kp_refiner_ = GraphMHAttn(dim_kp=2, num_kp=10, n_head=1) \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "#modified # added 20200616\n",
    "from modules.stn_a import STN\n",
    "heatmap_stn = STN((1,64,64))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "#kp_refiner=None\n",
    "heatmap_stn=None\n",
    "oval_heatmap=False if (heatmap_stn is None) else True"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Load Model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "由 demo.py 改來"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "mkdir: cannot create directory ‘demo/202006271200/’: File exists\n"
     ]
    }
   ],
   "source": [
    "!mkdir \"demo/202006271200/\"\n",
    "img_save_folder = \"demo/202006271200/\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "class opt2:\n",
    "    config = opt.config\n",
    "    checkpoint = '../../public/first-order-model/checkpoint/mgif-cpk.pth.tar'\n",
    "    #checkpoint =  \"demo/202006221000/mgif_cpk_newattngraph_300.tar\"\n",
    "    cpu = True\n",
    "    relative = True\n",
    "    adapt_scale = True\n",
    "    find_best_frame = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "# load model\n",
    "checkpoint_path = opt2.checkpoint\n",
    "\n",
    "#checkpoint.seed(0)\n",
    "if opt2.cpu:\n",
    "    checkpoint = torch.load(checkpoint_path, map_location=torch.device('cpu'))\n",
    "else:\n",
    "    checkpoint = torch.load(checkpoint_path)\n",
    "\n",
    "generator.load_state_dict(checkpoint['generator'], strict=False)\n",
    "kp_detector.load_state_dict(checkpoint['kp_detector'])\n",
    "discriminator.load_state_dict(checkpoint['discriminator'])\n",
    "if 'kp_refiner' in checkpoint: \n",
    "    kp_refiner.load_state_dict(checkpoint['kp_refiner'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Fix Orig Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "def freeze_model(model):\n",
    "    model.eval()\n",
    "    for params in model.parameters():\n",
    "        params.requires_grad = False\n",
    "        \n",
    "def unfreeze_model(model):\n",
    "    model.train()\n",
    "    for params in model.parameters():\n",
    "        params.requires_grad = True"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "freeze_generator = False\n",
    "freeze_discriminator = False\n",
    "freeze_kp_detector = False"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "#只訓練kp_refiner\n",
    "freeze_generator = True\n",
    "freeze_discriminator = True\n",
    "freeze_kp_detector = True\n",
    "if freeze_generator: freeze_model(generator)\n",
    "if freeze_discriminator: freeze_model(discriminator)\n",
    "if freeze_kp_detector: freeze_model(kp_detector)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "#只訓練dense_motion_network\n",
    "freeze_generator = True\n",
    "freeze_discriminator = False\n",
    "freeze_kp_detector = True\n",
    "freeze_model(kp_detector)\n",
    "freeze_model(generator)\n",
    "unfreeze_model(generator.dense_motion_network)"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "if freeze_generator: freeze_model(generator)\n",
    "if freeze_discriminator: freeze_model(discriminator)\n",
    "if freeze_kp_detector: freeze_model(kp_detector)"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "unfreeze_model(generator)\n",
    "unfreeze_model(discriminator)\n",
    "unfreeze_model(kp_detector)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Train"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "由 run.py 改來"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "generator.train()\n",
    "discriminator.train()\n",
    "kp_detector.train()\n",
    "if not heatmap_stn is None: heatmap_stn.train()\n",
    "if not kp_refiner is None: kp_refiner.train()\n",
    "if torch.cuda.is_available():\n",
    "    generator.to(opt.device_ids[0])\n",
    "    discriminator.to(opt.device_ids[0])\n",
    "    kp_detector.to(opt.device_ids[0])\n",
    "    if not heatmap_stn is None:heatmap_stn.to(opt.device_ids[0])\n",
    "    if not kp_refiner  is None:kp_refiner.to(opt.device_ids[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Use predefined train-test split.\n"
     ]
    }
   ],
   "source": [
    "!rm -rf ../../public/first-order-model/moving-gif/train/.ipynb_checkpoints\n",
    "!rm -rf ../../public/first-order-model/moving-gif/test/.ipynb_checkpoints\n",
    "dataset = FramesDataset(is_train=(opt.mode == 'train'), **config['dataset_params'])"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "#dataset dim check\n",
    "for x in dataset:\n",
    "    if x['driving'].shape[0]!=3 or x['source'].shape[0]!=3:\n",
    "        print(x['name'], x['driving'].shape, x['source'].shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "if not os.path.exists(log_dir):\n",
    "    os.makedirs(log_dir)\n",
    "if not os.path.exists(os.path.join(log_dir, os.path.basename(opt.config))):\n",
    "    copy(opt.config, log_dir)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "由 train.py 改來"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "checkpoint = opt.checkpoint\n",
    "device_ids = opt.device_ids"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_params = config['train_params']\n",
    "\n",
    "optimizer_generator = torch.optim.Adam(generator.parameters(), lr=train_params['lr_generator'], betas=(0.5, 0.999))\n",
    "optimizer_discriminator = torch.optim.Adam(discriminator.parameters(), lr=train_params['lr_discriminator']*2, betas=(0.5, 0.999))\n",
    "optimizer_kp_detector = torch.optim.Adam(kp_detector.parameters(), lr=train_params['lr_kp_detector'], betas=(0.5, 0.999))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "start_epoch = 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "scheduler_generator = MultiStepLR(optimizer_generator, train_params['epoch_milestones'], gamma=0.1,\n",
    "                                  last_epoch=start_epoch - 1)\n",
    "scheduler_discriminator = MultiStepLR(optimizer_discriminator, train_params['epoch_milestones'], gamma=0.1,\n",
    "                                      last_epoch=start_epoch - 1)\n",
    "scheduler_kp_detector = MultiStepLR(optimizer_kp_detector, train_params['epoch_milestones'], gamma=0.1,\n",
    "                                    last_epoch=-1 + start_epoch * (train_params['lr_kp_detector'] != 0))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "# added 20200616\n",
    "if not kp_refiner is None:\n",
    "    optimizer_kp_refiner = torch.optim.Adam(kp_refiner.parameters(), lr=train_params['lr_kp_refiner'], betas=(0.5, 0.999))\n",
    "    scheduler_kp_refiner = MultiStepLR(optimizer_kp_detector, train_params['epoch_milestones'], gamma=0.1,\n",
    "                                        last_epoch=-1 + start_epoch * (train_params['lr_kp_detector'] != 0))\n",
    "\n",
    "if not heatmap_stn is None:\n",
    "    optimizer_heatmap_stn = torch.optim.Adam(heatmap_stn.parameters(), lr=train_params['lr_heatmap_stn'], betas=(0.5, 0.999))\n",
    "    scheduler_heatmap_stn = MultiStepLR(optimizer_kp_detector, train_params['epoch_milestones'], gamma=0.1,\n",
    "                                        last_epoch=-1 + start_epoch * (train_params['lr_kp_detector'] != 0))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [],
   "source": [
    "if 'num_repeats' in train_params or train_params['num_repeats'] != 1:\n",
    "    dataset = DatasetRepeater(dataset, train_params['num_repeats'])\n",
    "dataloader = DataLoader(dataset, batch_size=train_params['batch_size'], shuffle=True, num_workers=6, drop_last=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "generator_full = GeneratorFullModel(kp_detector, generator, discriminator, train_params)\n",
    "discriminator_full = DiscriminatorFullModel(kp_detector, generator, discriminator, train_params)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [],
   "source": [
    "if torch.cuda.is_available():\n",
    "    generator_full = DataParallelWithCallback(generator_full, device_ids=device_ids)\n",
    "    discriminator_full = DataParallelWithCallback(discriminator_full, device_ids=device_ids)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "MultiHeadAttention(\n",
       "  in_features=2, head_num=1, bias=True, activation=<function relu at 0x7f54158329e0>\n",
       "  (linear_q): Linear(in_features=2, out_features=2, bias=True)\n",
       "  (linear_k): Linear(in_features=2, out_features=2, bias=True)\n",
       "  (linear_v): Linear(in_features=2, out_features=2, bias=True)\n",
       "  (linear_o): Linear(in_features=2, out_features=2, bias=True)\n",
       ")"
      ]
     },
     "execution_count": 33,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "kp_refiner"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|          | 0/100 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "has kp_refiner:True, \n",
      "has heatmap_stn:False, oval_heatmap:False\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/opt/conda/lib/python3.7/site-packages/torch/nn/functional.py:2705: UserWarning: Default grid_sample and affine_grid behavior has changed to align_corners=False since 1.3.0. Please specify align_corners=True if the old behavior is desired. See the documentation of grid_sample for details.\n",
      "  warnings.warn(\"Default grid_sample and affine_grid behavior has changed \"\n",
      "/opt/conda/lib/python3.7/site-packages/torch/nn/functional.py:2506: UserWarning: Default upsampling behavior when mode=bilinear is changed to align_corners=False since 0.4.0. Please specify align_corners=True if the old behavior is desired. See the documentation of nn.Upsample for details.\n",
      "  \"See the documentation of nn.Upsample for details.\".format(mode))\n",
      "/opt/conda/lib/python3.7/site-packages/torch/nn/functional.py:1351: UserWarning: nn.functional.sigmoid is deprecated. Use torch.sigmoid instead.\n",
      "  warnings.warn(\"nn.functional.sigmoid is deprecated. Use torch.sigmoid instead.\")\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "    #edge_weights(>=0.0): 1.00 %\n",
      "    #edge_weights(>=0.1): 1.00 %\n",
      "    #edge_weights(>=0.2): 1.00 %\n",
      "    #edge_weights(>=0.3): 0.28 %\n",
      "    #edge_weights(>=0.4): 0.28 %\n",
      "    #edge_weights(>=0.5): 0.28 %\n",
      "    #edge_weights(>=0.6): 0.28 %\n",
      "    #edge_weights(>=0.7): 0.28 %\n",
      "    #edge_weights(>=0.8): 0.28 %\n",
      "    #edge_weights(>=0.9): 0.00 %\n",
      "    #edge_weights(>=1.0): 0.00 %\n",
      "[0-0]loss: 136.2640, adj #0: 72, adj #1:28\n",
      "[0-30]loss: 143.2298, adj #0: 72, adj #1:28\n",
      "    #edge_weights(>=0.0): 1.00 %\n",
      "    #edge_weights(>=0.1): 1.00 %\n",
      "    #edge_weights(>=0.2): 0.65 %\n",
      "    #edge_weights(>=0.3): 0.28 %\n",
      "    #edge_weights(>=0.4): 0.28 %\n",
      "    #edge_weights(>=0.5): 0.28 %\n",
      "    #edge_weights(>=0.6): 0.28 %\n",
      "    #edge_weights(>=0.7): 0.28 %\n",
      "    #edge_weights(>=0.8): 0.13 %\n",
      "    #edge_weights(>=0.9): 0.00 %\n",
      "    #edge_weights(>=1.0): 0.00 %\n",
      "[0-60]loss: 92.6280, adj #0: 72, adj #1:28\n",
      "[0-90]loss: 120.4336, adj #0: 72, adj #1:28\n",
      "    #edge_weights(>=0.0): 1.00 %\n",
      "    #edge_weights(>=0.1): 1.00 %\n",
      "    #edge_weights(>=0.2): 0.68 %\n",
      "    #edge_weights(>=0.3): 0.28 %\n",
      "    #edge_weights(>=0.4): 0.28 %\n",
      "    #edge_weights(>=0.5): 0.28 %\n",
      "    #edge_weights(>=0.6): 0.28 %\n",
      "    #edge_weights(>=0.7): 0.28 %\n",
      "    #edge_weights(>=0.8): 0.12 %\n",
      "    #edge_weights(>=0.9): 0.00 %\n",
      "    #edge_weights(>=1.0): 0.00 %\n",
      "[0-120]loss: 81.3103, adj #0: 72, adj #1:28\n",
      "[0-150]loss: 180.0942, adj #0: 72, adj #1:28\n",
      "    #edge_weights(>=0.0): 1.00 %\n",
      "    #edge_weights(>=0.1): 1.00 %\n",
      "    #edge_weights(>=0.2): 0.68 %\n",
      "    #edge_weights(>=0.3): 0.28 %\n",
      "    #edge_weights(>=0.4): 0.28 %\n",
      "    #edge_weights(>=0.5): 0.28 %\n",
      "    #edge_weights(>=0.6): 0.28 %\n",
      "    #edge_weights(>=0.7): 0.28 %\n",
      "    #edge_weights(>=0.8): 0.12 %\n",
      "    #edge_weights(>=0.9): 0.00 %\n",
      "    #edge_weights(>=1.0): 0.00 %\n",
      "[0-180]loss: 173.4882, adj #0: 72, adj #1:28\n",
      "[0-210]loss: 74.4455, adj #0: 72, adj #1:28\n",
      "    #edge_weights(>=0.0): 1.00 %\n",
      "    #edge_weights(>=0.1): 1.00 %\n",
      "    #edge_weights(>=0.2): 0.70 %\n",
      "    #edge_weights(>=0.3): 0.28 %\n",
      "    #edge_weights(>=0.4): 0.28 %\n",
      "    #edge_weights(>=0.5): 0.28 %\n",
      "    #edge_weights(>=0.6): 0.28 %\n",
      "    #edge_weights(>=0.7): 0.28 %\n",
      "    #edge_weights(>=0.8): 0.12 %\n",
      "    #edge_weights(>=0.9): 0.00 %\n",
      "    #edge_weights(>=1.0): 0.00 %\n",
      "[0-240]loss: 140.3449, adj #0: 72, adj #1:28\n",
      "[0-270]loss: 130.7182, adj #0: 72, adj #1:28\n",
      "    #edge_weights(>=0.0): 1.00 %\n",
      "    #edge_weights(>=0.1): 1.00 %\n",
      "    #edge_weights(>=0.2): 0.68 %\n",
      "    #edge_weights(>=0.3): 0.28 %\n",
      "    #edge_weights(>=0.4): 0.28 %\n",
      "    #edge_weights(>=0.5): 0.28 %\n",
      "    #edge_weights(>=0.6): 0.28 %\n",
      "    #edge_weights(>=0.7): 0.28 %\n",
      "    #edge_weights(>=0.8): 0.11 %\n",
      "    #edge_weights(>=0.9): 0.00 %\n",
      "    #edge_weights(>=1.0): 0.00 %\n",
      "[0-300]loss: 178.9117, adj #0: 72, adj #1:28\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/opt/conda/lib/python3.7/site-packages/torch/optim/lr_scheduler.py:122: UserWarning: Detected call of `lr_scheduler.step()` before `optimizer.step()`. In PyTorch 1.1.0 and later, you should call them in the opposite order: `optimizer.step()` before `lr_scheduler.step()`.  Failure to do this will result in PyTorch skipping the first value of the learning rate schedule. See more details at https://pytorch.org/docs/stable/optim.html#how-to-adjust-learning-rate\n",
      "  \"https://pytorch.org/docs/stable/optim.html#how-to-adjust-learning-rate\", UserWarning)\n",
      "  1%|          | 1/100 [00:50<1:22:35, 50.06s/it]\n"
     ]
    }
   ],
   "source": [
    "print(\"has kp_refiner:{}, \\nhas heatmap_stn:{}, oval_heatmap:{}\" \\\n",
    "      .format(not kp_refiner is None, not heatmap_stn is None, oval_heatmap))\n",
    "with Logger(log_dir=log_dir, visualizer_params=config['visualizer_params'], checkpoint_freq=train_params['checkpoint_freq']) as logger:\n",
    "    for epoch in trange(start_epoch, train_params['num_epochs']):#每一回合\n",
    "        if epoch>0: break\n",
    "        for j, x in enumerate(dataloader):\n",
    "            if j > 300: break\n",
    "            #print(0)\n",
    "            losses_generator, generated, (adj_source, adj_weights_source, adj_driving, adj_weights_driving) \\\n",
    "                = generator_full(x, kp_refiner=kp_refiner, \n",
    "                                 )\n",
    "            #loss_adj = torch.mean(masks_source * preserve_non_topk(masks_source, k=9) + \\\n",
    "            #                      masks_driving * preserve_non_topk(masks_driving, k=9))\n",
    "            #print(1)\n",
    "            loss_values = [val.mean() for val in losses_generator.values()]\n",
    "            loss = sum(loss_values)#+100*loss_adj\n",
    "            \n",
    "            #print(generator.dense_motion_network.kp_variance)\n",
    "\n",
    "            loss.backward()\n",
    "            if not freeze_generator:\n",
    "                optimizer_generator.step()\n",
    "                optimizer_generator.zero_grad()\n",
    "            if not heatmap_stn is None:\n",
    "                optimizer_heatmap_stn.step()\n",
    "                optimizer_heatmap_stn.zero_grad()  \n",
    "            if not kp_refiner is None:\n",
    "                optimizer_kp_refiner.step()\n",
    "                optimizer_kp_refiner.zero_grad()\n",
    "            if not freeze_kp_detector:\n",
    "                optimizer_kp_detector.step()\n",
    "                optimizer_kp_detector.zero_grad()\n",
    "            \n",
    "            #print\n",
    "            if j%60 == 0:\n",
    "                for i in range(11):  \n",
    "                    r_ = torch.sum(adj_weights_source>=(i/10))/np.product(adj_weights_source.shape)\n",
    "                    print(\"    #edge_weights(>={}): {:.2f} %\".format(i/10, r_)  )\n",
    "            if j%30 == 0:\n",
    "                if not kp_refiner is None:\n",
    "                    adj = np.array(adj_source.cpu())\n",
    "                    #G = kp_refiner.get_graph_by_adjmatrix(adj_source, draw=True)\n",
    "                    #draw_graph(G,save_pth=f'nx_fig/epo{epoch}-{j}-nx.png') #圖片存檔\n",
    "                    #np.save(f'nx_fig/epo{epoch}-{j}-adj.npy', adj  )\n",
    "                    print(\"[{}-{}]loss: {:.4f}, adj #0: {}, adj #1:{}\" \\\n",
    "                          .format(epoch,j,loss.cpu().item(), np.sum(adj==0), np.sum(adj==1)))\n",
    "                else:\n",
    "                    print(\"[{}-{}]loss: {:.4f}\" \\\n",
    "                      .format(epoch,j,loss.cpu().item()))\n",
    "    \n",
    "            if train_params['loss_weights']['generator_gan'] != 0:\n",
    "                optimizer_discriminator.zero_grad()\n",
    "                losses_discriminator = discriminator_full(x, generated)\n",
    "                loss_values = [val.mean() for val in losses_discriminator.values()]\n",
    "                loss = sum(loss_values)\n",
    "\n",
    "                loss.backward()\n",
    "                if not freeze_discriminator:\n",
    "                    optimizer_discriminator.step()\n",
    "                    optimizer_discriminator.zero_grad()\n",
    "\n",
    "            else:\n",
    "                losses_discriminator = {}\n",
    " \n",
    "            \n",
    "            losses_generator.update(losses_discriminator)\n",
    "            losses = {key: value.mean().detach().data.cpu().numpy() for key, value in losses_generator.items()}\n",
    "            logger.log_iter(losses=losses)\n",
    "\n",
    "        \n",
    "        scheduler_generator.step()\n",
    "        scheduler_discriminator.step()\n",
    "        scheduler_kp_detector.step()\n",
    "        if not kp_refiner is None: \n",
    "            scheduler_kp_refiner.step()\n",
    "        if not heatmap_stn is None: \n",
    "            scheduler_heatmap_stn.step()\n",
    "\n",
    "        logger.log_epoch(epoch, {'generator': generator,\n",
    "                                 'discriminator': discriminator,\n",
    "                                 'kp_detector': kp_detector,\n",
    "                                 'optimizer_generator': optimizer_generator,\n",
    "                                 'optimizer_discriminator': optimizer_discriminator,\n",
    "                                 'optimizer_kp_detector': optimizer_kp_detector}, inp=x, out=generated)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [],
   "source": [
    "#save model\n",
    "op_checkpoint = {\n",
    "    'generator' : generator.state_dict(),\n",
    "    'kp_detector': kp_detector.state_dict(),\n",
    "    'discriminator': discriminator.state_dict(),\n",
    "    'kp_refiner':kp_refiner.state_dict(),\n",
    "    \n",
    "}\n",
    "torch.save(op_checkpoint, \n",
    "           img_save_folder+f\"mgif_cpk_newattngraph_300(only_dm).tar\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}