train_gae_starmap_multisamples_scrnaseq_validation.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "import os\n",
    "\n",
    "\n",
    "import scanpy\n",
    "import numpy as np\n",
    "import scipy.sparse as sp\n",
    "\n",
    "import torch\n",
    "from torch import optim\n",
    "\n",
    "import gae.gae.optimizer as optimizer\n",
    "import gae.gae.model\n",
    "import gae.gae.preprocessing as preprocessing\n",
    "\n",
    "import pickle\n",
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns\n",
    "import umap\n",
    "import pandas as pd\n",
    "from sklearn.preprocessing import scale\n",
    "from sklearn.preprocessing import MinMaxScaler\n",
    "from sklearn.metrics import pairwise_distances"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Settings\n",
    "# os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"0\" \n",
    "use_cuda=True\n",
    "fastmode=False #Validate during training pass\n",
    "seed=6\n",
    "epochs=20000\n",
    "saveFreq=20\n",
    "batchsize=2000\n",
    "lr=0.0001 #initial learning rate\n",
    "lr_adv=0.001\n",
    "weight_decay=0 #Weight for L2 loss on embedding matrix.\n",
    "\n",
    "hidden1=30000 #Number of units in hidden layer 1\n",
    "hidden2=30000 #Number of units in hidden layer 2\n",
    "# hidden3=2048\n",
    "# hidden4=2048\n",
    "# hidden5=128\n",
    "fc_dim1=30000\n",
    "# fc_dim2=128\n",
    "# fc_dim3=128\n",
    "# fc_dim4=128\n",
    "# gcn_dim1=2600\n",
    "# clf_hidden=256\n",
    "adv_hidden=128\n",
    "\n",
    "dropout=0.01\n",
    "testNodes=0.1 #fraction of total nodes for testing\n",
    "valNodes=0.05 #fraction of total nodes for validation\n",
    "XreconWeight=20\n",
    "# clfweight=20\n",
    "advWeight=2\n",
    "# randFeatureSubset=None\n",
    "model_str='fc1_dca_sharded'\n",
    "clf=None\n",
    "adv=None  #'clf_fc1_eq'  #'clf_fc1_control_eq' #'clf_fc1_control'  #'clf_fc1'\n",
    "protein=None #'nearest' #None #'scaled_binary'\n",
    "adj_decodeName=None #gala or None\n",
    "ridgeL=0.01\n",
    "shareGenePi=True\n",
    "\n",
    "targetBatch=None\n",
    "training_sample_X='logminmax'\n",
    "switchFreq=10\n",
    "standardizeX=False\n",
    "tissue='Bone_Marrow'\n",
    "name='allk20XA_02_dca_over_FCXonly_scrnaseq_'+tissue \n",
    "logsavepath='/data/xinyi/log/train_gae_scrnaseq/'+name\n",
    "modelsavepath='/data/xinyi/models/train_gae_scrnaseq/'+name\n",
    "plotsavepath='/data/xinyi/plots/train_gae_scrnaseq/'+name\n",
    "datadir='/data/xinyi/STACI/scrnaseq/'\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "data=scanpy.read_h5ad(os.path.join(datadir,'Immune_ALL_human.h5ad'))\n",
    "data=data[:,np.array(np.sum(data.layers['counts'],axis=0)>3).flatten()]\n",
    "if tissue!= 'all':\n",
    "    data=data[data.obs['tissue']==tissue]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "if training_sample_X=='logminmax':\n",
    "    data_train=np.log2(data.layers['counts']+1/2)\n",
    "    scaler = MinMaxScaler()\n",
    "    data_train=np.transpose(scaler.fit_transform(np.transpose(data_train)))\n",
    "    data_train=torch.tensor(data_train)\n",
    "    \n",
    "if 'dca' in model_str:\n",
    "    data_raw=torch.tensor(data.layers['counts']+1/2)\n",
    "    \n",
    "\n",
    "# Set cuda and seed\n",
    "np.random.seed(seed)\n",
    "if use_cuda and (not torch.cuda.is_available()):\n",
    "    print('cuda not available')\n",
    "    use_cuda=False\n",
    "torch.manual_seed(seed)\n",
    "if use_cuda:\n",
    "    torch.cuda.manual_seed(seed)\n",
    "    torch.backends.cudnn.enabled = True\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "if not os.path.exists(logsavepath):\n",
    "    os.mkdir(logsavepath)\n",
    "if not os.path.exists(modelsavepath):\n",
    "    os.mkdir(modelsavepath)\n",
    "if not os.path.exists(plotsavepath):\n",
    "    os.mkdir(plotsavepath)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "# loop over all train/validation sets\n",
    "np.random.seed(seed)\n",
    "torch.manual_seed(seed)\n",
    "if use_cuda:\n",
    "    torch.cuda.manual_seed(seed)\n",
    "    torch.backends.cudnn.enabled = True\n",
    "    \n",
    "num_features=data_train.shape[1]\n",
    "mse=torch.nn.MSELoss()\n",
    "# mse=torch.nn.MSELoss(reduction=None)\n",
    "# Create model\n",
    "if model_str=='gcn_vae_xa':\n",
    "    model = gae.gae.model.GCNModelVAE_XA(num_features, hidden1, hidden2,fc_dim1,fc_dim2,fc_dim3,fc_dim4, dropout)\n",
    "    loss_kl=optimizer.optimizer_kl\n",
    "    loss_x=optimizer.optimizer_MSE\n",
    "    loss_a=optimizer.optimizer_CE\n",
    "elif model_str == 'gcn_vae_gcnX_inprA':\n",
    "    model = gae.gae.model.GCNModelVAE_gcnX_inprA(num_features, hidden1, hidden2,gcn_dim1, dropout)\n",
    "    loss_kl=optimizer.optimizer_kl\n",
    "    loss_x=optimizer.optimizer_MSE\n",
    "    loss_a=optimizer.optimizer_CE\n",
    "    \n",
    "elif model_str == 'gcn_vae_gcnX_inprA_w':\n",
    "    model = gae.gae.model.GCNModelVAE_gcnX_inprA_w(num_features, hidden1, hidden2,gcn_dim1, dropout)\n",
    "    loss_kl=optimizer.optimizer_kl\n",
    "    loss_x=optimizer.optimizer_MSE\n",
    "    loss_a=optimizer.optimizer_CE    \n",
    "    \n",
    "elif model_str=='gcn_vae_xa_e3':\n",
    "    model = gae.gae.model.GCNModelVAE_XA_e3(num_features, hidden1, hidden2,hidden3,fc_dim1,fc_dim2,fc_dim3,fc_dim4, dropout)\n",
    "    loss_kl=optimizer.optimizer_kl\n",
    "    loss_x=optimizer.optimizer_MSE\n",
    "    loss_a=optimizer.optimizer_CE\n",
    "    \n",
    "elif model_str == 'gcn_vae_xa_e1':\n",
    "    model = gae.gae.model.GCNModelVAE_XA_e1(num_features, hidden1, dropout)\n",
    "    loss_kl=optimizer.optimizer_kl\n",
    "    loss_x=optimizer.optimizer_MSE\n",
    "    loss_a=optimizer.optimizer_CE\n",
    "\n",
    "elif model_str == 'gcn_vae_xa_e2_d1':\n",
    "    model = gae.gae.model.GCNModelVAE_XA_e2_d1(num_features, hidden1,hidden2, dropout)\n",
    "    loss_kl=optimizer.optimizer_kl\n",
    "    loss_x=optimizer.optimizer_MSE\n",
    "    loss_a=optimizer.optimizer_CE\n",
    "    \n",
    "elif model_str=='gcn_vae_xa_e2_d1_dca':\n",
    "    model = gae.gae.model.GCNModelVAE_XA_e2_d1_DCA(num_features, hidden1,hidden2,fc_dim1, dropout)\n",
    "    loss_kl=optimizer.optimizer_kl\n",
    "    loss_x=optimizer.optimizer_zinb\n",
    "    loss_a=optimizer.optimizer_CE\n",
    "\n",
    "elif model_str=='gcn_vae_xa_e2_d1_dca_fca':\n",
    "    model = gae.gae.model.GCNModelVAE_XA_e2_d1_DCA_fca(num_features, hidden1,hidden2,fc_dim1, dropout)\n",
    "    loss_kl=optimizer.optimizer_kl\n",
    "    loss_x=optimizer.optimizer_zinb\n",
    "    loss_a=optimizer.optimizer_CE\n",
    "    \n",
    "elif model_str=='gcn_vae_xa_e2_d1_dcaFork':\n",
    "    model = gae.gae.model.GCNModelVAE_XA_e2_d1_DCAfork(num_features, hidden1,hidden2,fc_dim1, dropout)\n",
    "    loss_kl=optimizer.optimizer_kl\n",
    "    loss_x=optimizer.optimizer_zinb\n",
    "    loss_a=optimizer.optimizer_CE\n",
    "\n",
    "elif model_str=='gcn_vae_xa_e2_d1_dcaElemPi':\n",
    "    model = gae.gae.model.GCNModelVAE_XA_e2_d1_DCAelemPi(num_features, hidden1,hidden2,fc_dim1, dropout,shareGenePi)\n",
    "    loss_kl=optimizer.optimizer_kl\n",
    "    loss_x=optimizer.optimizer_zinb\n",
    "    loss_a=optimizer.optimizer_CE\n",
    "\n",
    "elif model_str=='gcn_vae_xa_e2_d1_dcaConstantDisp':\n",
    "    model = gae.gae.model.GCNModelVAE_XA_e2_d1_DCA_constantDisp(num_features, hidden1,hidden2,fc_dim1, dropout,shareGenePi)\n",
    "    loss_kl=optimizer.optimizer_kl\n",
    "    loss_x=optimizer.optimizer_zinb\n",
    "    loss_a=optimizer.optimizer_CE    \n",
    "    \n",
    "elif model_str == 'gcn_vae_xa_e4_d1':\n",
    "    model = gae.gae.model.GCNModelVAE_XA_e4_d1(num_features, hidden1,hidden2,hidden3,hidden4, dropout)\n",
    "    loss_kl=optimizer.optimizer_kl\n",
    "    loss_x=optimizer.optimizer_MSE\n",
    "    loss_a=optimizer.optimizer_CE    \n",
    "    \n",
    "elif model_str=='fc':\n",
    "    model = gae.gae.model.FCVAE(num_features, hidden1, hidden2,hidden3,hidden4,hidden5,fc_dim1,fc_dim2,fc_dim3,fc_dim4, dropout)\n",
    "    loss_kl=optimizer.optimizer_kl\n",
    "    loss_x=optimizer.optimizer_MSE\n",
    "#     loss_x=mse\n",
    "    loss_a=optimizer.optimizer_CE\n",
    "elif model_str=='fcae':\n",
    "    model = gae.gae.model.FCAE(num_features, hidden1, hidden2,hidden3,hidden4,hidden5,fc_dim1,fc_dim2,fc_dim3,fc_dim4, dropout)\n",
    "    loss_x=optimizer.optimizer_MSE\n",
    "#     loss_x=mse\n",
    "\n",
    "elif model_str=='fcae1':\n",
    "    model = gae.gae.model.FCAE1(num_features, dropout,hidden1)\n",
    "    loss_x=optimizer.optimizer_MSE\n",
    "#     loss_x=mse\n",
    "elif model_str=='fcae2':\n",
    "    model = gae.gae.model.FCAE2(num_features, dropout,hidden1,hidden2)\n",
    "    loss_x=optimizer.optimizer_MSE\n",
    "\n",
    "elif model_str=='fc1':\n",
    "    model = gae.gae.model.FCVAE1(num_features, hidden1,dropout)\n",
    "    loss_x=optimizer.optimizer_MSE\n",
    "    loss_kl=optimizer.optimizer_kl\n",
    "    loss_a=optimizer.optimizer_CE\n",
    "\n",
    "elif model_str=='fc1_fca':\n",
    "    model = gae.gae.model.FCVAE1_fca(num_features, hidden1,dropout)\n",
    "    loss_x=optimizer.optimizer_MSE\n",
    "    loss_kl=optimizer.optimizer_kl\n",
    "    loss_a=optimizer.optimizer_CE    \n",
    "    \n",
    "elif model_str=='fc1_dca':\n",
    "    model = gae.gae.model.FCVAE1_DCA(num_features, hidden1,fc_dim1, dropout)\n",
    "    loss_x=optimizer.optimizer_zinb\n",
    "    loss_kl=optimizer.optimizer_kl\n",
    "    loss_a=optimizer.optimizer_CE\n",
    "    \n",
    "elif model_str=='fc1_dca_sharded':\n",
    "    model = gae.gae.model.FCVAE1_DCA_sharded(num_features, hidden1,fc_dim1, dropout)\n",
    "    loss_x=optimizer.optimizer_zinb\n",
    "    loss_kl=optimizer.optimizer_kl\n",
    "    loss_a=optimizer.optimizer_CE\n",
    "    \n",
    "if clf=='clf_fc1':\n",
    "    modelClf=gae.gae.model.Clf_fc1(hidden2, dropout,clf_hidden,ct_unique.size)\n",
    "    loss_clf=torch.nn.CrossEntropyLoss()\n",
    "    \n",
    "if adv=='clf_fc1' or adv=='clf_fc1_eq' or adv=='clf_fc1_control' or adv=='clf_fc1_control_eq':\n",
    "    modelAdv=gae.gae.model.Clf_fc1(hidden2, dropout,adv_hidden,sampleLabellist_ae['control13'].size()[1])\n",
    "    loss_adv=optimizer.optimizer_CEclf\n",
    "    \n",
    "if adv=='clf_linear1' or adv=='clf_linear1_control':\n",
    "    modelAdv=gae.gae.model.Clf_linear1(hidden2, dropout,sampleLabellist_ae['control13'].size()[1])\n",
    "    loss_adv=optimizer.optimizer_CEclf\n",
    "    \n",
    "if 'NB' in name:\n",
    "    print('using NB loss for X')\n",
    "    loss_x=optimizer.optimizer_nb\n",
    "    \n",
    "# if use_cuda:\n",
    "#     model.cuda()\n",
    "    \n",
    "\n",
    "optimizerVAEXA = optim.Adam(model.parameters(), lr=lr, weight_decay=weight_decay)\n",
    "if clf:\n",
    "    optimizerClf=optim.Adam(modelClf.parameters(), lr=lr, weight_decay=weight_decay)\n",
    "if adv:\n",
    "    optimizerAdv=optim.Adam(modelAdv.parameters(), lr=lr_adv, weight_decay=weight_decay)    \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "pctVal=0.05\n",
    "pctTest=0.1\n",
    "np.random.seed(3)\n",
    "allIdx=np.arange(data.shape[0])\n",
    "np.random.shuffle(allIdx)\n",
    "valIdx=allIdx[:int(pctVal*data.shape[0])]\n",
    "testIdx=allIdx[int(pctVal*data.shape[0]):(int(pctVal*data.shape[0])+int(pctTest*data.shape[0]))]\n",
    "trainIdx=allIdx[(int(pctVal*data.shape[0])+int(pctTest*data.shape[0])):]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " Epoch: 0000 loss_train: 3.9689 loss_kl_train: 238.9152 loss_x_train: 3.9450 loss_val: 6.3519 time: 7.0443s\n",
      " Epoch: 0001 loss_train: 2.0433 loss_kl_train: 797.7281 loss_x_train: 1.9635 loss_val: 4.2195 time: 5.7016s\n",
      " Epoch: 0002 loss_train: 3.3342 loss_kl_train: 7616.9347 loss_x_train: 2.5725 loss_val: 4.3637 time: 5.7282s\n",
      " Epoch: 0003 loss_train: 2.8537 loss_kl_train: 7570.0473 loss_x_train: 2.0967 loss_val: 3.2791 time: 5.7369s\n",
      " Epoch: 0004 loss_train: 3.3372 loss_kl_train: 13723.3738 loss_x_train: 1.9649 loss_val: 2.7111 time: 5.7179s\n",
      " Epoch: 0005 loss_train: 6.4816 loss_kl_train: 46460.9623 loss_x_train: 1.8356 loss_val: 2.4949 time: 5.7989s\n",
      " Epoch: 0006 loss_train: 6.2741 loss_kl_train: 41419.7039 loss_x_train: 2.1321 loss_val: 2.7203 time: 5.7441s\n",
      " Epoch: 0007 loss_train: 5.5228 loss_kl_train: 35551.0493 loss_x_train: 1.9677 loss_val: 2.5361 time: 5.7528s\n",
      " Epoch: 0008 loss_train: 6.5626 loss_kl_train: 47818.9674 loss_x_train: 1.7807 loss_val: 2.3176 time: 5.6544s\n",
      " Epoch: 0009 loss_train: 16.8622 loss_kl_train: 149625.8974 loss_x_train: 1.8997 loss_val: 2.2463 time: 5.7158s\n",
      " Epoch: 0010 loss_train: 24.1846 loss_kl_train: 220065.6125 loss_x_train: 2.1781 loss_val: 2.9254 time: 5.7242s\n",
      " Epoch: 0011 loss_train: 7.4861 loss_kl_train: 56296.4058 loss_x_train: 1.8564 loss_val: 2.7588 time: 5.7112s\n",
      " Epoch: 0012 loss_train: 69.8296 loss_kl_train: 680274.6365 loss_x_train: 1.8022 loss_val: 2.6659 time: 5.6746s\n",
      " Epoch: 0013 loss_train: 73.7921 loss_kl_train: 718765.6931 loss_x_train: 1.9156 loss_val: 3.1049 time: 5.7424s\n",
      " Epoch: 0014 loss_train: 93.7912 loss_kl_train: 916905.4028 loss_x_train: 2.1006 loss_val: 3.8560 time: 5.7849s\n",
      " Epoch: 0015 loss_train: 20.5388 loss_kl_train: 187292.1446 loss_x_train: 1.8096 loss_val: 3.5711 time: 5.6865s\n",
      " Epoch: 0016 loss_train: 1.8308 loss_kl_train: 4516.1333 loss_x_train: 1.3791 loss_val: 2.9638 time: 5.7268s\n",
      " Epoch: 0017 loss_train: 1.3949 loss_kl_train: 352.4911 loss_x_train: 1.3597 loss_val: 2.5262 time: 5.7923s\n",
      " Epoch: 0018 loss_train: 1.3708 loss_kl_train: 202.5596 loss_x_train: 1.3505 loss_val: 2.1667 time: 5.7513s\n",
      " Epoch: 0019 loss_train: 1.3596 loss_kl_train: 167.8010 loss_x_train: 1.3428 loss_val: 1.9220 time: 5.8035s\n",
      " Epoch: 0020 loss_train: 1.3538 loss_kl_train: 163.9482 loss_x_train: 1.3374 loss_val: 1.7502 time: 5.8166s\n",
      " Epoch: 0021 loss_train: 1.3482 loss_kl_train: 157.9078 loss_x_train: 1.3324 loss_val: 1.6220 time: 5.7644s\n",
      " Epoch: 0022 loss_train: 1.3436 loss_kl_train: 154.0059 loss_x_train: 1.3282 loss_val: 1.5330 time: 5.6959s\n",
      " Epoch: 0023 loss_train: 1.3402 loss_kl_train: 151.0217 loss_x_train: 1.3251 loss_val: 1.4676 time: 5.6722s\n",
      " Epoch: 0024 loss_train: 1.3369 loss_kl_train: 148.2386 loss_x_train: 1.3221 loss_val: 1.4233 time: 5.7887s\n",
      " Epoch: 0025 loss_train: 1.3351 loss_kl_train: 150.4741 loss_x_train: 1.3201 loss_val: 1.3932 time: 5.7885s\n",
      " Epoch: 0026 loss_train: 1.3317 loss_kl_train: 142.7209 loss_x_train: 1.3174 loss_val: 1.3692 time: 5.7101s\n",
      " Epoch: 0027 loss_train: 1.3296 loss_kl_train: 139.0840 loss_x_train: 1.3157 loss_val: 1.3566 time: 5.7651s\n",
      " Epoch: 0028 loss_train: 1.3274 loss_kl_train: 136.6392 loss_x_train: 1.3138 loss_val: 1.3486 time: 5.7828s\n",
      " Epoch: 0029 loss_train: 1.3257 loss_kl_train: 134.6355 loss_x_train: 1.3122 loss_val: 1.3418 time: 5.7669s\n",
      " Epoch: 0030 loss_train: 1.3243 loss_kl_train: 132.3868 loss_x_train: 1.3110 loss_val: 1.3370 time: 5.7824s\n",
      " Epoch: 0031 loss_train: 1.3227 loss_kl_train: 130.2826 loss_x_train: 1.3096 loss_val: 1.3346 time: 5.8076s\n",
      " Epoch: 0032 loss_train: 1.3214 loss_kl_train: 128.5451 loss_x_train: 1.3086 loss_val: 1.3328 time: 5.7093s\n",
      " Epoch: 0033 loss_train: 1.3201 loss_kl_train: 126.6087 loss_x_train: 1.3074 loss_val: 1.3306 time: 5.7089s\n",
      " Epoch: 0034 loss_train: 1.3190 loss_kl_train: 125.8299 loss_x_train: 1.3064 loss_val: 1.3297 time: 5.7121s\n",
      " Epoch: 0035 loss_train: 1.3179 loss_kl_train: 123.7348 loss_x_train: 1.3056 loss_val: 1.3286 time: 5.7707s\n",
      " Epoch: 0036 loss_train: 1.3169 loss_kl_train: 121.7334 loss_x_train: 1.3047 loss_val: 1.3274 time: 5.7196s\n",
      " Epoch: 0037 loss_train: 1.3160 loss_kl_train: 120.4061 loss_x_train: 1.3039 loss_val: 1.3264 time: 5.7454s\n",
      " Epoch: 0038 loss_train: 1.3152 loss_kl_train: 118.8934 loss_x_train: 1.3033 loss_val: 1.3256 time: 5.7783s\n",
      " Epoch: 0039 loss_train: 1.3144 loss_kl_train: 117.3595 loss_x_train: 1.3027 loss_val: 1.3248 time: 5.7550s\n",
      " Epoch: 0040 loss_train: 1.3135 loss_kl_train: 116.2877 loss_x_train: 1.3019 loss_val: 1.3240 time: 5.6666s\n",
      " Epoch: 0041 loss_train: 1.3128 loss_kl_train: 115.0381 loss_x_train: 1.3013 loss_val: 1.3237 time: 5.6848s\n",
      " Epoch: 0042 loss_train: 1.3121 loss_kl_train: 113.2214 loss_x_train: 1.3008 loss_val: 1.3228 time: 5.7887s\n",
      " Epoch: 0043 loss_train: 1.3115 loss_kl_train: 112.4489 loss_x_train: 1.3002 loss_val: 1.3224 time: 5.7586s\n",
      " Epoch: 0044 loss_train: 1.3109 loss_kl_train: 111.3089 loss_x_train: 1.2997 loss_val: 1.3217 time: 5.7595s\n",
      " Epoch: 0045 loss_train: 1.3102 loss_kl_train: 110.2377 loss_x_train: 1.2992 loss_val: 1.3208 time: 5.7950s\n",
      " Epoch: 0046 loss_train: 1.3096 loss_kl_train: 108.8571 loss_x_train: 1.2987 loss_val: 1.3203 time: 5.8051s\n",
      " Epoch: 0047 loss_train: 1.3092 loss_kl_train: 107.8419 loss_x_train: 1.2985 loss_val: 1.3205 time: 5.6987s\n",
      " Epoch: 0048 loss_train: 1.3087 loss_kl_train: 107.0842 loss_x_train: 1.2980 loss_val: 1.3204 time: 5.6685s\n",
      " Epoch: 0049 loss_train: 1.3082 loss_kl_train: 106.1467 loss_x_train: 1.2976 loss_val: 1.3200 time: 5.7605s\n",
      " Epoch: 0050 loss_train: 1.3078 loss_kl_train: 104.9371 loss_x_train: 1.2973 loss_val: 1.3188 time: 5.7883s\n",
      " Epoch: 0051 loss_train: 1.3073 loss_kl_train: 104.0569 loss_x_train: 1.2969 loss_val: 1.3181 time: 5.8267s\n",
      " Epoch: 0052 loss_train: 1.3069 loss_kl_train: 102.9183 loss_x_train: 1.2966 loss_val: 1.3177 time: 5.8151s\n",
      " Epoch: 0053 loss_train: 1.3065 loss_kl_train: 102.2044 loss_x_train: 1.2963 loss_val: 1.3175 time: 5.7891s\n",
      " Epoch: 0054 loss_train: 1.3061 loss_kl_train: 101.2913 loss_x_train: 1.2960 loss_val: 1.3175 time: 5.7802s\n",
      " Epoch: 0055 loss_train: 1.3058 loss_kl_train: 100.4057 loss_x_train: 1.2958 loss_val: 1.3175 time: 5.7950s\n",
      " Epoch: 0056 loss_train: 1.3054 loss_kl_train: 99.5951 loss_x_train: 1.2954 loss_val: 1.3172 time: 5.8112s\n",
      " Epoch: 0057 loss_train: 1.3050 loss_kl_train: 98.8472 loss_x_train: 1.2951 loss_val: 1.3168 time: 5.7416s\n",
      " Epoch: 0058 loss_train: 1.3047 loss_kl_train: 98.0908 loss_x_train: 1.2949 loss_val: 1.3161 time: 5.7487s\n",
      " Epoch: 0059 loss_train: 1.3043 loss_kl_train: 97.1884 loss_x_train: 1.2946 loss_val: 1.3157 time: 5.8264s\n",
      " Epoch: 0060 loss_train: 1.3041 loss_kl_train: 96.7176 loss_x_train: 1.2944 loss_val: 1.3155 time: 5.8523s\n",
      " Epoch: 0061 loss_train: 1.3038 loss_kl_train: 95.7725 loss_x_train: 1.2942 loss_val: 1.3154 time: 5.8312s\n",
      " Epoch: 0062 loss_train: 1.3034 loss_kl_train: 94.9978 loss_x_train: 1.2939 loss_val: 1.3151 time: 5.8448s\n",
      " Epoch: 0063 loss_train: 1.3032 loss_kl_train: 94.2533 loss_x_train: 1.2938 loss_val: 1.3150 time: 5.7551s\n",
      " Epoch: 0064 loss_train: 1.3029 loss_kl_train: 93.6508 loss_x_train: 1.2936 loss_val: 1.3150 time: 5.7829s\n",
      " Epoch: 0065 loss_train: 1.3026 loss_kl_train: 92.9798 loss_x_train: 1.2933 loss_val: 1.3151 time: 5.8200s\n",
      " Epoch: 0066 loss_train: 1.3025 loss_kl_train: 92.5032 loss_x_train: 1.2932 loss_val: 1.3150 time: 5.7945s\n",
      " Epoch: 0067 loss_train: 1.3021 loss_kl_train: 92.0502 loss_x_train: 1.2929 loss_val: 1.3146 time: 5.7825s\n",
      " Epoch: 0068 loss_train: 1.3019 loss_kl_train: 91.2879 loss_x_train: 1.2928 loss_val: 1.3144 time: 5.7062s\n",
      " Epoch: 0069 loss_train: 1.3017 loss_kl_train: 90.4419 loss_x_train: 1.2926 loss_val: 1.3137 time: 5.7797s\n",
      " Epoch: 0070 loss_train: 1.3015 loss_kl_train: 89.9876 loss_x_train: 1.2925 loss_val: 1.3134 time: 5.6903s\n",
      " Epoch: 0071 loss_train: 1.3012 loss_kl_train: 89.3558 loss_x_train: 1.2923 loss_val: 1.3133 time: 5.8032s\n",
      " Epoch: 0072 loss_train: 1.3010 loss_kl_train: 88.6536 loss_x_train: 1.2921 loss_val: 1.3131 time: 5.8019s\n",
      " Epoch: 0073 loss_train: 1.3008 loss_kl_train: 88.3035 loss_x_train: 1.2919 loss_val: 1.3130 time: 5.8288s\n",
      " Epoch: 0074 loss_train: 1.3005 loss_kl_train: 87.7368 loss_x_train: 1.2917 loss_val: 1.3130 time: 5.8629s\n",
      " Epoch: 0075 loss_train: 1.3003 loss_kl_train: 87.1135 loss_x_train: 1.2916 loss_val: 1.3133 time: 5.7196s\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " Epoch: 0076 loss_train: 1.3001 loss_kl_train: 86.5619 loss_x_train: 1.2914 loss_val: 1.3134 time: 5.8011s\n",
      " Epoch: 0077 loss_train: 1.2999 loss_kl_train: 86.0599 loss_x_train: 1.2913 loss_val: 1.3128 time: 5.8017s\n",
      " Epoch: 0078 loss_train: 1.2997 loss_kl_train: 85.5457 loss_x_train: 1.2911 loss_val: 1.3124 time: 5.7860s\n",
      " Epoch: 0079 loss_train: 1.2996 loss_kl_train: 85.1380 loss_x_train: 1.2910 loss_val: 1.3123 time: 5.7353s\n",
      " Epoch: 0080 loss_train: 1.2993 loss_kl_train: 84.4747 loss_x_train: 1.2909 loss_val: 1.3122 time: 5.7936s\n",
      " Epoch: 0081 loss_train: 1.2992 loss_kl_train: 83.9703 loss_x_train: 1.2908 loss_val: 1.3118 time: 5.7792s\n",
      " Epoch: 0082 loss_train: 1.2990 loss_kl_train: 83.4681 loss_x_train: 1.2906 loss_val: 1.3118 time: 5.6784s\n",
      " Epoch: 0083 loss_train: 1.2988 loss_kl_train: 83.0953 loss_x_train: 1.2905 loss_val: 1.3118 time: 5.6802s\n",
      " Epoch: 0084 loss_train: 1.2986 loss_kl_train: 82.5636 loss_x_train: 1.2904 loss_val: 1.3120 time: 5.7812s\n",
      " Epoch: 0085 loss_train: 1.2984 loss_kl_train: 82.1230 loss_x_train: 1.2902 loss_val: 1.3116 time: 5.7764s\n",
      " Epoch: 0086 loss_train: 1.2983 loss_kl_train: 81.8140 loss_x_train: 1.2901 loss_val: 1.3114 time: 5.8044s\n",
      " Epoch: 0087 loss_train: 1.2981 loss_kl_train: 81.2738 loss_x_train: 1.2900 loss_val: 1.3113 time: 5.8066s\n",
      " Epoch: 0088 loss_train: 1.2979 loss_kl_train: 80.8070 loss_x_train: 1.2898 loss_val: 1.3114 time: 5.7891s\n",
      " Epoch: 0089 loss_train: 1.2978 loss_kl_train: 80.4318 loss_x_train: 1.2897 loss_val: 1.3110 time: 5.7668s\n",
      " Epoch: 0090 loss_train: 1.2976 loss_kl_train: 79.9947 loss_x_train: 1.2896 loss_val: 1.3106 time: 5.8205s\n",
      " Epoch: 0091 loss_train: 1.2975 loss_kl_train: 79.5646 loss_x_train: 1.2896 loss_val: 1.3107 time: 5.8147s\n",
      " Epoch: 0092 loss_train: 1.2974 loss_kl_train: 79.2061 loss_x_train: 1.2894 loss_val: 1.3107 time: 5.8668s\n",
      " Epoch: 0093 loss_train: 1.2972 loss_kl_train: 78.7208 loss_x_train: 1.2893 loss_val: 1.3107 time: 5.8268s\n",
      " Epoch: 0094 loss_train: 1.2970 loss_kl_train: 78.1996 loss_x_train: 1.2892 loss_val: 1.3106 time: 5.8031s\n",
      " Epoch: 0095 loss_train: 1.2969 loss_kl_train: 77.9498 loss_x_train: 1.2891 loss_val: 1.3103 time: 5.8194s\n",
      " Epoch: 0096 loss_train: 1.2967 loss_kl_train: 77.5420 loss_x_train: 1.2890 loss_val: 1.3108 time: 5.8295s\n",
      " Epoch: 0097 loss_train: 1.2966 loss_kl_train: 77.2490 loss_x_train: 1.2889 loss_val: 1.3104 time: 5.8240s\n",
      " Epoch: 0098 loss_train: 1.2965 loss_kl_train: 76.8947 loss_x_train: 1.2888 loss_val: 1.3102 time: 5.8318s\n",
      " Epoch: 0099 loss_train: 1.2964 loss_kl_train: 76.5734 loss_x_train: 1.2887 loss_val: 1.3102 time: 5.7795s\n",
      " Epoch: 0100 loss_train: 1.2962 loss_kl_train: 76.1044 loss_x_train: 1.2886 loss_val: 1.3102 time: 5.7453s\n",
      " Epoch: 0101 loss_train: 1.2961 loss_kl_train: 75.7453 loss_x_train: 1.2886 loss_val: 1.3098 time: 5.7966s\n",
      " Epoch: 0102 loss_train: 1.2960 loss_kl_train: 75.3933 loss_x_train: 1.2885 loss_val: 1.3096 time: 5.7717s\n",
      " Epoch: 0103 loss_train: 1.2959 loss_kl_train: 75.0442 loss_x_train: 1.2884 loss_val: 1.3094 time: 5.7974s\n",
      " Epoch: 0104 loss_train: 1.2958 loss_kl_train: 74.7543 loss_x_train: 1.2883 loss_val: 1.3094 time: 5.8090s\n",
      " Epoch: 0105 loss_train: 1.2956 loss_kl_train: 74.3675 loss_x_train: 1.2882 loss_val: 1.3091 time: 5.7456s\n",
      " Epoch: 0106 loss_train: 1.2956 loss_kl_train: 74.2425 loss_x_train: 1.2882 loss_val: 1.3089 time: 5.7892s\n",
      " Epoch: 0107 loss_train: 1.2955 loss_kl_train: 73.7457 loss_x_train: 1.2881 loss_val: 1.3091 time: 5.8103s\n",
      " Epoch: 0108 loss_train: 1.2954 loss_kl_train: 73.2400 loss_x_train: 1.2880 loss_val: 1.3091 time: 5.8252s\n",
      " Epoch: 0109 loss_train: 1.2953 loss_kl_train: 73.1318 loss_x_train: 1.2880 loss_val: 1.3092 time: 5.8377s\n",
      " Epoch: 0110 loss_train: 1.2952 loss_kl_train: 72.7793 loss_x_train: 1.2879 loss_val: 1.3091 time: 5.7488s\n",
      " Epoch: 0111 loss_train: 1.2950 loss_kl_train: 72.3630 loss_x_train: 1.2877 loss_val: 1.3092 time: 5.7880s\n",
      " Epoch: 0112 loss_train: 1.2949 loss_kl_train: 72.0315 loss_x_train: 1.2877 loss_val: 1.3094 time: 5.8114s\n",
      " Epoch: 0113 loss_train: 1.2947 loss_kl_train: 71.7280 loss_x_train: 1.2876 loss_val: 1.3098 time: 5.8357s\n",
      " Epoch: 0114 loss_train: 1.2947 loss_kl_train: 71.4390 loss_x_train: 1.2875 loss_val: 1.3097 time: 5.8324s\n",
      " Epoch: 0115 loss_train: 1.2946 loss_kl_train: 71.2096 loss_x_train: 1.2875 loss_val: 1.3097 time: 5.8246s\n",
      " Epoch: 0116 loss_train: 1.2945 loss_kl_train: 70.7912 loss_x_train: 1.2874 loss_val: 1.3091 time: 5.7933s\n",
      " Epoch: 0117 loss_train: 1.2944 loss_kl_train: 70.5465 loss_x_train: 1.2873 loss_val: 1.3086 time: 5.8182s\n",
      " Epoch: 0118 loss_train: 1.2943 loss_kl_train: 70.3421 loss_x_train: 1.2873 loss_val: 1.3084 time: 5.8165s\n",
      " Epoch: 0119 loss_train: 1.2942 loss_kl_train: 70.0039 loss_x_train: 1.2872 loss_val: 1.3081 time: 5.8021s\n",
      " Epoch: 0120 loss_train: 1.2941 loss_kl_train: 69.6855 loss_x_train: 1.2871 loss_val: 1.3082 time: 5.7738s\n",
      " Epoch: 0121 loss_train: 1.2940 loss_kl_train: 69.4580 loss_x_train: 1.2871 loss_val: 1.3083 time: 5.7804s\n",
      " Epoch: 0122 loss_train: 1.2940 loss_kl_train: 69.2446 loss_x_train: 1.2871 loss_val: 1.3084 time: 5.8471s\n",
      " Epoch: 0123 loss_train: 1.2940 loss_kl_train: 68.9214 loss_x_train: 1.2871 loss_val: 1.3084 time: 5.7371s\n",
      " Epoch: 0124 loss_train: 1.2939 loss_kl_train: 68.5580 loss_x_train: 1.2871 loss_val: 1.3080 time: 5.7657s\n",
      " Epoch: 0125 loss_train: 1.2938 loss_kl_train: 68.3063 loss_x_train: 1.2870 loss_val: 1.3079 time: 5.7947s\n",
      " Epoch: 0126 loss_train: 1.2935 loss_kl_train: 67.9735 loss_x_train: 1.2867 loss_val: 1.3084 time: 5.7728s\n",
      " Epoch: 0127 loss_train: 1.2934 loss_kl_train: 67.8428 loss_x_train: 1.2867 loss_val: 1.3084 time: 5.8028s\n",
      " Epoch: 0128 loss_train: 1.2933 loss_kl_train: 67.4289 loss_x_train: 1.2866 loss_val: 1.3084 time: 5.8154s\n",
      " Epoch: 0129 loss_train: 1.2933 loss_kl_train: 67.1780 loss_x_train: 1.2866 loss_val: 1.3087 time: 5.8060s\n",
      " Epoch: 0130 loss_train: 1.2932 loss_kl_train: 67.0143 loss_x_train: 1.2865 loss_val: 1.3091 time: 5.8022s\n",
      " Epoch: 0131 loss_train: 1.2932 loss_kl_train: 66.8231 loss_x_train: 1.2865 loss_val: 1.3091 time: 5.7998s\n",
      " Epoch: 0132 loss_train: 1.2930 loss_kl_train: 66.4703 loss_x_train: 1.2864 loss_val: 1.3083 time: 5.8435s\n",
      " Epoch: 0133 loss_train: 1.2930 loss_kl_train: 66.3374 loss_x_train: 1.2863 loss_val: 1.3076 time: 5.8492s\n",
      " Epoch: 0134 loss_train: 1.2928 loss_kl_train: 65.9750 loss_x_train: 1.2862 loss_val: 1.3073 time: 5.7761s\n",
      " Epoch: 0135 loss_train: 1.2928 loss_kl_train: 65.8306 loss_x_train: 1.2862 loss_val: 1.3073 time: 5.7185s\n",
      " Epoch: 0136 loss_train: 1.2927 loss_kl_train: 65.5757 loss_x_train: 1.2862 loss_val: 1.3074 time: 5.7859s\n",
      " Epoch: 0137 loss_train: 1.2927 loss_kl_train: 65.2780 loss_x_train: 1.2861 loss_val: 1.3077 time: 5.7815s\n",
      " Epoch: 0138 loss_train: 1.2927 loss_kl_train: 65.0540 loss_x_train: 1.2862 loss_val: 1.3076 time: 5.8572s\n",
      " Epoch: 0139 loss_train: 1.2927 loss_kl_train: 64.8620 loss_x_train: 1.2862 loss_val: 1.3074 time: 5.8228s\n",
      " Epoch: 0140 loss_train: 1.2926 loss_kl_train: 64.6042 loss_x_train: 1.2861 loss_val: 1.3072 time: 5.7811s\n",
      " Epoch: 0141 loss_train: 1.2925 loss_kl_train: 64.3179 loss_x_train: 1.2860 loss_val: 1.3069 time: 5.7928s\n",
      " Epoch: 0142 loss_train: 1.2923 loss_kl_train: 64.1951 loss_x_train: 1.2859 loss_val: 1.3074 time: 5.7504s\n",
      " Epoch: 0143 loss_train: 1.2922 loss_kl_train: 63.8611 loss_x_train: 1.2858 loss_val: 1.3087 time: 5.8153s\n",
      " Epoch: 0144 loss_train: 1.2921 loss_kl_train: 63.6733 loss_x_train: 1.2858 loss_val: 1.3092 time: 5.7934s\n",
      " Epoch: 0145 loss_train: 1.2922 loss_kl_train: 63.3987 loss_x_train: 1.2858 loss_val: 1.3097 time: 5.8119s\n",
      " Epoch: 0146 loss_train: 1.2923 loss_kl_train: 63.2528 loss_x_train: 1.2859 loss_val: 1.3092 time: 5.8377s\n",
      " Epoch: 0147 loss_train: 1.2922 loss_kl_train: 63.0975 loss_x_train: 1.2859 loss_val: 1.3084 time: 5.7646s\n",
      " Epoch: 0148 loss_train: 1.2921 loss_kl_train: 62.8972 loss_x_train: 1.2859 loss_val: 1.3072 time: 5.7741s\n",
      " Epoch: 0149 loss_train: 1.2920 loss_kl_train: 62.7339 loss_x_train: 1.2857 loss_val: 1.3064 time: 5.7961s\n",
      " Epoch: 0150 loss_train: 1.2917 loss_kl_train: 62.4256 loss_x_train: 1.2855 loss_val: 1.3068 time: 5.8197s\n",
      " Epoch: 0151 loss_train: 1.2917 loss_kl_train: 62.2694 loss_x_train: 1.2855 loss_val: 1.3073 time: 5.8501s\n",
      " Epoch: 0152 loss_train: 1.2917 loss_kl_train: 62.0569 loss_x_train: 1.2855 loss_val: 1.3079 time: 5.8250s\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " Epoch: 0153 loss_train: 1.2919 loss_kl_train: 61.8252 loss_x_train: 1.2857 loss_val: 1.3077 time: 5.8833s\n",
      " Epoch: 0154 loss_train: 1.2920 loss_kl_train: 61.6415 loss_x_train: 1.2858 loss_val: 1.3070 time: 5.8709s\n",
      " Epoch: 0155 loss_train: 1.2918 loss_kl_train: 61.4580 loss_x_train: 1.2856 loss_val: 1.3068 time: 5.8465s\n",
      " Epoch: 0156 loss_train: 1.2915 loss_kl_train: 61.2386 loss_x_train: 1.2854 loss_val: 1.3078 time: 5.8892s\n",
      " Epoch: 0157 loss_train: 1.2913 loss_kl_train: 60.9787 loss_x_train: 1.2852 loss_val: 1.3090 time: 5.8245s\n",
      " Epoch: 0158 loss_train: 1.2912 loss_kl_train: 60.8216 loss_x_train: 1.2852 loss_val: 1.3092 time: 5.8018s\n",
      " Epoch: 0159 loss_train: 1.2914 loss_kl_train: 60.6049 loss_x_train: 1.2853 loss_val: 1.3091 time: 5.7530s\n",
      " Epoch: 0160 loss_train: 1.2915 loss_kl_train: 60.4296 loss_x_train: 1.2854 loss_val: 1.3080 time: 5.7873s\n",
      " Epoch: 0161 loss_train: 1.2913 loss_kl_train: 60.2614 loss_x_train: 1.2853 loss_val: 1.3065 time: 5.7683s\n",
      " Epoch: 0162 loss_train: 1.2911 loss_kl_train: 60.1373 loss_x_train: 1.2851 loss_val: 1.3061 time: 5.6896s\n",
      " Epoch: 0163 loss_train: 1.2909 loss_kl_train: 59.9083 loss_x_train: 1.2849 loss_val: 1.3070 time: 5.6628s\n",
      " Epoch: 0164 loss_train: 1.2910 loss_kl_train: 59.7238 loss_x_train: 1.2850 loss_val: 1.3075 time: 5.7513s\n",
      " Epoch: 0165 loss_train: 1.2910 loss_kl_train: 59.5251 loss_x_train: 1.2851 loss_val: 1.3071 time: 5.8185s\n",
      " Epoch: 0166 loss_train: 1.2911 loss_kl_train: 59.2874 loss_x_train: 1.2851 loss_val: 1.3064 time: 5.7348s\n",
      " Epoch: 0167 loss_train: 1.2909 loss_kl_train: 59.1201 loss_x_train: 1.2850 loss_val: 1.3065 time: 5.7325s\n",
      " Epoch: 0168 loss_train: 1.2906 loss_kl_train: 58.8793 loss_x_train: 1.2847 loss_val: 1.3077 time: 5.7717s\n",
      " Epoch: 0169 loss_train: 1.2905 loss_kl_train: 58.7229 loss_x_train: 1.2846 loss_val: 1.3085 time: 5.8067s\n",
      " Epoch: 0170 loss_train: 1.2905 loss_kl_train: 58.6265 loss_x_train: 1.2846 loss_val: 1.3082 time: 5.8460s\n",
      " Epoch: 0171 loss_train: 1.2906 loss_kl_train: 58.4047 loss_x_train: 1.2847 loss_val: 1.3070 time: 5.8136s\n",
      " Epoch: 0172 loss_train: 1.2905 loss_kl_train: 58.2937 loss_x_train: 1.2847 loss_val: 1.3060 time: 5.8164s\n",
      " Epoch: 0173 loss_train: 1.2903 loss_kl_train: 58.0641 loss_x_train: 1.2845 loss_val: 1.3055 time: 5.7634s\n",
      " Epoch: 0174 loss_train: 1.2901 loss_kl_train: 57.9659 loss_x_train: 1.2844 loss_val: 1.3063 time: 5.7231s\n",
      " Epoch: 0175 loss_train: 1.2902 loss_kl_train: 57.7669 loss_x_train: 1.2844 loss_val: 1.3066 time: 5.7940s\n",
      " Epoch: 0176 loss_train: 1.2903 loss_kl_train: 57.6221 loss_x_train: 1.2845 loss_val: 1.3063 time: 5.8242s\n",
      " Epoch: 0177 loss_train: 1.2903 loss_kl_train: 57.4209 loss_x_train: 1.2845 loss_val: 1.3059 time: 5.8321s\n",
      " Epoch: 0178 loss_train: 1.2901 loss_kl_train: 57.2087 loss_x_train: 1.2844 loss_val: 1.3059 time: 5.8331s\n",
      " Epoch: 0179 loss_train: 1.2899 loss_kl_train: 57.1195 loss_x_train: 1.2842 loss_val: 1.3070 time: 5.8136s\n",
      " Epoch: 0180 loss_train: 1.2898 loss_kl_train: 56.8220 loss_x_train: 1.2841 loss_val: 1.3083 time: 5.7458s\n",
      " Epoch: 0181 loss_train: 1.2899 loss_kl_train: 56.8279 loss_x_train: 1.2842 loss_val: 1.3089 time: 5.8037s\n",
      " Epoch: 0182 loss_train: 1.2901 loss_kl_train: 56.6438 loss_x_train: 1.2845 loss_val: 1.3077 time: 5.8059s\n",
      " Epoch: 0183 loss_train: 1.2900 loss_kl_train: 56.5231 loss_x_train: 1.2843 loss_val: 1.3060 time: 5.8323s\n",
      " Epoch: 0184 loss_train: 1.2898 loss_kl_train: 56.4078 loss_x_train: 1.2841 loss_val: 1.3053 time: 5.8070s\n",
      " Epoch: 0185 loss_train: 1.2895 loss_kl_train: 56.1509 loss_x_train: 1.2839 loss_val: 1.3059 time: 5.7987s\n",
      " Epoch: 0186 loss_train: 1.2896 loss_kl_train: 56.0539 loss_x_train: 1.2840 loss_val: 1.3071 time: 5.7299s\n",
      " Epoch: 0187 loss_train: 1.2898 loss_kl_train: 55.8838 loss_x_train: 1.2842 loss_val: 1.3068 time: 5.7854s\n",
      " Epoch: 0188 loss_train: 1.2898 loss_kl_train: 55.7000 loss_x_train: 1.2843 loss_val: 1.3056 time: 5.7907s\n",
      " Epoch: 0189 loss_train: 1.2897 loss_kl_train: 55.5336 loss_x_train: 1.2841 loss_val: 1.3058 time: 5.8149s\n",
      " Epoch: 0190 loss_train: 1.2893 loss_kl_train: 55.3194 loss_x_train: 1.2838 loss_val: 1.3071 time: 5.8207s\n",
      " Epoch: 0191 loss_train: 1.2892 loss_kl_train: 55.2068 loss_x_train: 1.2837 loss_val: 1.3085 time: 5.7730s\n",
      " Epoch: 0192 loss_train: 1.2895 loss_kl_train: 55.0485 loss_x_train: 1.2840 loss_val: 1.3086 time: 5.7737s\n",
      " Epoch: 0193 loss_train: 1.2896 loss_kl_train: 54.9158 loss_x_train: 1.2842 loss_val: 1.3067 time: 5.8720s\n",
      " Epoch: 0194 loss_train: 1.2894 loss_kl_train: 54.7941 loss_x_train: 1.2839 loss_val: 1.3054 time: 5.8043s\n",
      " Epoch: 0195 loss_train: 1.2892 loss_kl_train: 54.7007 loss_x_train: 1.2837 loss_val: 1.3054 time: 5.8264s\n",
      " Epoch: 0196 loss_train: 1.2890 loss_kl_train: 54.5643 loss_x_train: 1.2836 loss_val: 1.3067 time: 5.9277s\n",
      " Epoch: 0197 loss_train: 1.2892 loss_kl_train: 54.4441 loss_x_train: 1.2838 loss_val: 1.3072 time: 5.7532s\n",
      " Epoch: 0198 loss_train: 1.2895 loss_kl_train: 54.2241 loss_x_train: 1.2841 loss_val: 1.3056 time: 5.7806s\n",
      " Epoch: 0199 loss_train: 1.2894 loss_kl_train: 54.0857 loss_x_train: 1.2840 loss_val: 1.3055 time: 5.8666s\n",
      " Epoch: 0200 loss_train: 1.2889 loss_kl_train: 53.9266 loss_x_train: 1.2836 loss_val: 1.3075 time: 5.8292s\n",
      " Epoch: 0201 loss_train: 1.2889 loss_kl_train: 53.7611 loss_x_train: 1.2835 loss_val: 1.3087 time: 5.8354s\n",
      " Epoch: 0202 loss_train: 1.2890 loss_kl_train: 53.6287 loss_x_train: 1.2837 loss_val: 1.3084 time: 5.8058s\n",
      " Epoch: 0203 loss_train: 1.2891 loss_kl_train: 53.5623 loss_x_train: 1.2838 loss_val: 1.3061 time: 5.7886s\n",
      " Epoch: 0204 loss_train: 1.2890 loss_kl_train: 53.4766 loss_x_train: 1.2836 loss_val: 1.3048 time: 5.8080s\n",
      " Epoch: 0205 loss_train: 1.2886 loss_kl_train: 53.2814 loss_x_train: 1.2833 loss_val: 1.3053 time: 5.7455s\n",
      " Epoch: 0206 loss_train: 1.2885 loss_kl_train: 53.2027 loss_x_train: 1.2832 loss_val: 1.3062 time: 5.7863s\n",
      " Epoch: 0207 loss_train: 1.2886 loss_kl_train: 53.0465 loss_x_train: 1.2833 loss_val: 1.3059 time: 5.8036s\n",
      " Epoch: 0208 loss_train: 1.2887 loss_kl_train: 52.9179 loss_x_train: 1.2834 loss_val: 1.3050 time: 5.8092s\n",
      " Epoch: 0209 loss_train: 1.2885 loss_kl_train: 52.7988 loss_x_train: 1.2832 loss_val: 1.3052 time: 5.8512s\n",
      " Epoch: 0210 loss_train: 1.2882 loss_kl_train: 52.5576 loss_x_train: 1.2830 loss_val: 1.3063 time: 5.7128s\n",
      " Epoch: 0211 loss_train: 1.2882 loss_kl_train: 52.4115 loss_x_train: 1.2830 loss_val: 1.3072 time: 5.6923s\n",
      " Epoch: 0212 loss_train: 1.2883 loss_kl_train: 52.3854 loss_x_train: 1.2831 loss_val: 1.3064 time: 5.6937s\n",
      " Epoch: 0213 loss_train: 1.2883 loss_kl_train: 52.2425 loss_x_train: 1.2830 loss_val: 1.3054 time: 5.7963s\n",
      " Epoch: 0214 loss_train: 1.2881 loss_kl_train: 52.1339 loss_x_train: 1.2829 loss_val: 1.3048 time: 5.7877s\n",
      " Epoch: 0215 loss_train: 1.2880 loss_kl_train: 52.0495 loss_x_train: 1.2828 loss_val: 1.3048 time: 5.7727s\n",
      " Epoch: 0216 loss_train: 1.2878 loss_kl_train: 51.8253 loss_x_train: 1.2827 loss_val: 1.3051 time: 5.7077s\n",
      " Epoch: 0217 loss_train: 1.2879 loss_kl_train: 51.6990 loss_x_train: 1.2827 loss_val: 1.3048 time: 5.7126s\n",
      " Epoch: 0218 loss_train: 1.2878 loss_kl_train: 51.6039 loss_x_train: 1.2826 loss_val: 1.3047 time: 5.7139s\n",
      " Epoch: 0219 loss_train: 1.2877 loss_kl_train: 51.4158 loss_x_train: 1.2826 loss_val: 1.3048 time: 5.7495s\n",
      " Epoch: 0220 loss_train: 1.2876 loss_kl_train: 51.3350 loss_x_train: 1.2825 loss_val: 1.3053 time: 5.7696s\n",
      " Epoch: 0221 loss_train: 1.2876 loss_kl_train: 51.1955 loss_x_train: 1.2825 loss_val: 1.3058 time: 5.7444s\n",
      " Epoch: 0222 loss_train: 1.2876 loss_kl_train: 51.1282 loss_x_train: 1.2825 loss_val: 1.3064 time: 5.7600s\n",
      " Epoch: 0223 loss_train: 1.2878 loss_kl_train: 50.9696 loss_x_train: 1.2827 loss_val: 1.3062 time: 5.7874s\n",
      " Epoch: 0224 loss_train: 1.2878 loss_kl_train: 50.8635 loss_x_train: 1.2827 loss_val: 1.3050 time: 5.8047s\n",
      " Epoch: 0225 loss_train: 1.2876 loss_kl_train: 50.7529 loss_x_train: 1.2825 loss_val: 1.3043 time: 5.7772s\n",
      " Epoch: 0226 loss_train: 1.2874 loss_kl_train: 50.6657 loss_x_train: 1.2823 loss_val: 1.3049 time: 5.7682s\n",
      " Epoch: 0227 loss_train: 1.2874 loss_kl_train: 50.5487 loss_x_train: 1.2824 loss_val: 1.3055 time: 5.7733s\n",
      " Epoch: 0228 loss_train: 1.2876 loss_kl_train: 50.4589 loss_x_train: 1.2825 loss_val: 1.3056 time: 5.7493s\n",
      " Epoch: 0229 loss_train: 1.2877 loss_kl_train: 50.2369 loss_x_train: 1.2827 loss_val: 1.3054 time: 5.7818s\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " Epoch: 0230 loss_train: 1.2876 loss_kl_train: 50.1425 loss_x_train: 1.2826 loss_val: 1.3047 time: 5.8357s\n",
      " Epoch: 0231 loss_train: 1.2874 loss_kl_train: 49.9958 loss_x_train: 1.2824 loss_val: 1.3050 time: 5.8305s\n",
      " Epoch: 0232 loss_train: 1.2873 loss_kl_train: 49.9012 loss_x_train: 1.2823 loss_val: 1.3063 time: 5.8613s\n",
      " Epoch: 0233 loss_train: 1.2873 loss_kl_train: 49.7833 loss_x_train: 1.2823 loss_val: 1.3084 time: 5.7581s\n",
      " Epoch: 0234 loss_train: 1.2876 loss_kl_train: 49.6595 loss_x_train: 1.2826 loss_val: 1.3074 time: 5.7056s\n",
      " Epoch: 0235 loss_train: 1.2877 loss_kl_train: 49.6346 loss_x_train: 1.2828 loss_val: 1.3056 time: 5.8376s\n",
      " Epoch: 0236 loss_train: 1.2875 loss_kl_train: 49.5237 loss_x_train: 1.2826 loss_val: 1.3045 time: 5.8224s\n",
      " Epoch: 0237 loss_train: 1.2873 loss_kl_train: 49.4566 loss_x_train: 1.2823 loss_val: 1.3048 time: 5.8897s\n",
      " Epoch: 0238 loss_train: 1.2871 loss_kl_train: 49.2505 loss_x_train: 1.2822 loss_val: 1.3062 time: 5.8648s\n",
      " Epoch: 0239 loss_train: 1.2874 loss_kl_train: 49.2529 loss_x_train: 1.2825 loss_val: 1.3079 time: 5.7466s\n",
      " Epoch: 0240 loss_train: 1.2878 loss_kl_train: 49.0527 loss_x_train: 1.2829 loss_val: 1.3063 time: 5.6923s\n",
      " Epoch: 0241 loss_train: 1.2877 loss_kl_train: 48.9279 loss_x_train: 1.2828 loss_val: 1.3047 time: 5.6807s\n",
      " Epoch: 0242 loss_train: 1.2873 loss_kl_train: 48.7640 loss_x_train: 1.2824 loss_val: 1.3057 time: 5.7036s\n",
      " Epoch: 0243 loss_train: 1.2871 loss_kl_train: 48.6219 loss_x_train: 1.2822 loss_val: 1.3088 time: 5.7948s\n",
      " Epoch: 0244 loss_train: 1.2874 loss_kl_train: 48.5723 loss_x_train: 1.2826 loss_val: 1.3092 time: 5.7753s\n",
      " Epoch: 0245 loss_train: 1.2879 loss_kl_train: 48.4914 loss_x_train: 1.2830 loss_val: 1.3078 time: 5.7976s\n",
      " Epoch: 0246 loss_train: 1.2878 loss_kl_train: 48.3891 loss_x_train: 1.2829 loss_val: 1.3050 time: 5.7894s\n",
      " Epoch: 0247 loss_train: 1.2873 loss_kl_train: 48.3422 loss_x_train: 1.2824 loss_val: 1.3046 time: 5.8430s\n",
      " Epoch: 0248 loss_train: 1.2871 loss_kl_train: 48.3266 loss_x_train: 1.2822 loss_val: 1.3063 time: 5.7912s\n",
      " Epoch: 0249 loss_train: 1.2872 loss_kl_train: 48.1541 loss_x_train: 1.2823 loss_val: 1.3072 time: 5.8174s\n",
      " Epoch: 0250 loss_train: 1.2874 loss_kl_train: 47.9985 loss_x_train: 1.2826 loss_val: 1.3052 time: 5.8589s\n",
      " Epoch: 0251 loss_train: 1.2872 loss_kl_train: 47.8968 loss_x_train: 1.2825 loss_val: 1.3045 time: 5.7718s\n",
      " Epoch: 0252 loss_train: 1.2869 loss_kl_train: 47.7004 loss_x_train: 1.2821 loss_val: 1.3063 time: 5.7405s\n",
      " Epoch: 0253 loss_train: 1.2868 loss_kl_train: 47.6221 loss_x_train: 1.2820 loss_val: 1.3084 time: 5.7723s\n",
      " Epoch: 0254 loss_train: 1.2870 loss_kl_train: 47.4778 loss_x_train: 1.2823 loss_val: 1.3066 time: 5.7847s\n",
      " Epoch: 0255 loss_train: 1.2870 loss_kl_train: 47.5287 loss_x_train: 1.2822 loss_val: 1.3048 time: 5.8105s\n",
      " Epoch: 0256 loss_train: 1.2867 loss_kl_train: 47.4037 loss_x_train: 1.2820 loss_val: 1.3043 time: 5.7704s\n",
      " Epoch: 0257 loss_train: 1.2865 loss_kl_train: 47.2811 loss_x_train: 1.2817 loss_val: 1.3049 time: 5.7586s\n",
      " Epoch: 0258 loss_train: 1.2865 loss_kl_train: 47.2589 loss_x_train: 1.2818 loss_val: 1.3053 time: 5.8122s\n",
      " Epoch: 0259 loss_train: 1.2866 loss_kl_train: 47.1251 loss_x_train: 1.2818 loss_val: 1.3044 time: 5.8102s\n",
      " Epoch: 0260 loss_train: 1.2865 loss_kl_train: 47.0151 loss_x_train: 1.2818 loss_val: 1.3039 time: 5.7778s\n",
      " Epoch: 0261 loss_train: 1.2862 loss_kl_train: 46.8889 loss_x_train: 1.2815 loss_val: 1.3051 time: 5.7982s\n",
      " Epoch: 0262 loss_train: 1.2862 loss_kl_train: 46.7256 loss_x_train: 1.2815 loss_val: 1.3058 time: 5.7674s\n",
      " Epoch: 0263 loss_train: 1.2862 loss_kl_train: 46.6673 loss_x_train: 1.2815 loss_val: 1.3053 time: 5.7989s\n",
      " Epoch: 0264 loss_train: 1.2862 loss_kl_train: 46.5599 loss_x_train: 1.2815 loss_val: 1.3041 time: 5.7918s\n",
      " Epoch: 0265 loss_train: 1.2861 loss_kl_train: 46.4545 loss_x_train: 1.2814 loss_val: 1.3038 time: 5.9459s\n",
      " Epoch: 0266 loss_train: 1.2859 loss_kl_train: 46.3750 loss_x_train: 1.2812 loss_val: 1.3040 time: 5.9070s\n",
      " Epoch: 0267 loss_train: 1.2859 loss_kl_train: 46.3441 loss_x_train: 1.2812 loss_val: 1.3045 time: 5.8378s\n",
      " Epoch: 0268 loss_train: 1.2859 loss_kl_train: 46.2158 loss_x_train: 1.2813 loss_val: 1.3040 time: 5.8238s\n",
      " Epoch: 0269 loss_train: 1.2859 loss_kl_train: 46.1297 loss_x_train: 1.2813 loss_val: 1.3037 time: 5.8696s\n",
      " Epoch: 0270 loss_train: 1.2858 loss_kl_train: 45.9612 loss_x_train: 1.2812 loss_val: 1.3040 time: 5.8067s\n",
      " Epoch: 0271 loss_train: 1.2857 loss_kl_train: 45.8965 loss_x_train: 1.2811 loss_val: 1.3046 time: 5.6966s\n",
      " Epoch: 0272 loss_train: 1.2856 loss_kl_train: 45.7869 loss_x_train: 1.2811 loss_val: 1.3053 time: 5.8067s\n",
      " Epoch: 0273 loss_train: 1.2858 loss_kl_train: 45.7533 loss_x_train: 1.2812 loss_val: 1.3050 time: 5.8474s\n",
      " Epoch: 0274 loss_train: 1.2858 loss_kl_train: 45.6253 loss_x_train: 1.2812 loss_val: 1.3041 time: 5.8091s\n",
      " Epoch: 0275 loss_train: 1.2857 loss_kl_train: 45.4957 loss_x_train: 1.2811 loss_val: 1.3039 time: 5.8121s\n",
      " Epoch: 0276 loss_train: 1.2856 loss_kl_train: 45.5361 loss_x_train: 1.2810 loss_val: 1.3038 time: 5.8209s\n",
      " Epoch: 0277 loss_train: 1.2856 loss_kl_train: 45.4356 loss_x_train: 1.2811 loss_val: 1.3057 time: 5.8599s\n",
      " Epoch: 0278 loss_train: 1.2859 loss_kl_train: 45.2725 loss_x_train: 1.2814 loss_val: 1.3053 time: 5.8746s\n",
      " Epoch: 0279 loss_train: 1.2860 loss_kl_train: 45.1257 loss_x_train: 1.2815 loss_val: 1.3042 time: 5.8356s\n",
      " Epoch: 0280 loss_train: 1.2859 loss_kl_train: 45.0496 loss_x_train: 1.2814 loss_val: 1.3042 time: 5.7634s\n",
      " Epoch: 0281 loss_train: 1.2857 loss_kl_train: 45.0461 loss_x_train: 1.2812 loss_val: 1.3056 time: 5.7905s\n",
      " Epoch: 0282 loss_train: 1.2857 loss_kl_train: 44.8969 loss_x_train: 1.2812 loss_val: 1.3077 time: 5.7837s\n",
      " Epoch: 0283 loss_train: 1.2858 loss_kl_train: 44.8032 loss_x_train: 1.2813 loss_val: 1.3083 time: 5.8288s\n",
      " Epoch: 0284 loss_train: 1.2861 loss_kl_train: 44.6730 loss_x_train: 1.2816 loss_val: 1.3059 time: 5.8260s\n",
      " Epoch: 0285 loss_train: 1.2866 loss_kl_train: 48.4344 loss_x_train: 1.2818 loss_val: 1.3037 time: 5.8378s\n",
      " Epoch: 0286 loss_train: 1.2888 loss_kl_train: 71.7460 loss_x_train: 1.2817 loss_val: 1.3041 time: 5.7141s\n",
      " Epoch: 0287 loss_train: 1.2915 loss_kl_train: 81.4791 loss_x_train: 1.2834 loss_val: 1.3097 time: 5.7787s\n",
      " Epoch: 0288 loss_train: 1.3178 loss_kl_train: 351.1531 loss_x_train: 1.2827 loss_val: 1.3069 time: 5.8702s\n",
      " Epoch: 0289 loss_train: 1.2869 loss_kl_train: 45.0529 loss_x_train: 1.2824 loss_val: 1.3053 time: 5.8069s\n",
      " Epoch: 0290 loss_train: 1.2901 loss_kl_train: 80.3450 loss_x_train: 1.2820 loss_val: 1.3054 time: 5.8459s\n",
      " Epoch: 0291 loss_train: 1.2918 loss_kl_train: 87.2041 loss_x_train: 1.2831 loss_val: 1.3083 time: 5.8124s\n",
      " Epoch: 0292 loss_train: 1.3012 loss_kl_train: 185.0117 loss_x_train: 1.2827 loss_val: 1.3094 time: 5.8315s\n",
      " Epoch: 0293 loss_train: 1.7511 loss_kl_train: 4452.0127 loss_x_train: 1.3059 loss_val: 1.3338 time: 5.8038s\n",
      " Epoch: 0294 loss_train: 1.3298 loss_kl_train: 233.6661 loss_x_train: 1.3064 loss_val: 1.3186 time: 5.8283s\n",
      " Epoch: 0295 loss_train: 1.2998 loss_kl_train: 67.3282 loss_x_train: 1.2931 loss_val: 1.3499 time: 5.8377s\n",
      " Epoch: 0296 loss_train: 2.0019 loss_kl_train: 6728.6093 loss_x_train: 1.3290 loss_val: 1.3338 time: 5.7963s\n",
      " Epoch: 0297 loss_train: 3.2064 loss_kl_train: 18772.5287 loss_x_train: 1.3291 loss_val: 1.3597 time: 5.8473s\n",
      " Epoch: 0298 loss_train: 1.3292 loss_kl_train: 161.3756 loss_x_train: 1.3130 loss_val: 1.3372 time: 5.8034s\n",
      " Epoch: 0299 loss_train: 1.3902 loss_kl_train: 834.2346 loss_x_train: 1.3068 loss_val: 1.3308 time: 5.7631s\n",
      " Epoch: 0300 loss_train: 1.4130 loss_kl_train: 1097.8954 loss_x_train: 1.3032 loss_val: 1.3223 time: 5.7951s\n",
      " Epoch: 0301 loss_train: 5.9815 loss_kl_train: 41194.5885 loss_x_train: 1.8620 loss_val: 1.4394 time: 5.7899s\n",
      " Epoch: 0302 loss_train: 3.8926 loss_kl_train: 25100.3829 loss_x_train: 1.3826 loss_val: 1.4958 time: 5.7757s\n",
      " Epoch: 0303 loss_train: 1.4569 loss_kl_train: 1192.6810 loss_x_train: 1.3376 loss_val: 1.8655 time: 5.7794s\n",
      " Epoch: 0304 loss_train: 1.3962 loss_kl_train: 766.2633 loss_x_train: 1.3195 loss_val: 1.6577 time: 5.8255s\n",
      " Epoch: 0305 loss_train: 1.8404 loss_kl_train: 5207.1821 loss_x_train: 1.3197 loss_val: 1.4131 time: 5.8159s\n",
      " Epoch: 0306 loss_train: 1.3496 loss_kl_train: 350.5021 loss_x_train: 1.3145 loss_val: 1.3661 time: 5.8119s\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " Epoch: 0307 loss_train: 1.5102 loss_kl_train: 1974.0967 loss_x_train: 1.3128 loss_val: 1.3285 time: 5.8429s\n",
      " Epoch: 0308 loss_train: 2.3431 loss_kl_train: 9679.3044 loss_x_train: 1.3752 loss_val: 1.3616 time: 5.7701s\n",
      " Epoch: 0309 loss_train: 4.8837 loss_kl_train: 8846.0609 loss_x_train: 3.9991 loss_val: 1.4836 time: 5.7330s\n",
      " Epoch: 0310 loss_train: 13.1742 loss_kl_train: 76219.1052 loss_x_train: 5.5523 loss_val: 1.6024 time: 5.8124s\n",
      " Epoch: 0311 loss_train: 26.4357 loss_kl_train: 245405.2128 loss_x_train: 1.8952 loss_val: 1.8267 time: 5.8534s\n",
      " Epoch: 0312 loss_train: 3.4633 loss_kl_train: 16580.1079 loss_x_train: 1.8053 loss_val: 1.8121 time: 5.8299s\n",
      " Epoch: 0313 loss_train: 1.9184 loss_kl_train: 2141.6075 loss_x_train: 1.7043 loss_val: 1.7146 time: 5.8653s\n",
      " Epoch: 0314 loss_train: 1.6950 loss_kl_train: 928.6013 loss_x_train: 1.6021 loss_val: 1.6831 time: 5.7937s\n",
      " Epoch: 0315 loss_train: 1.6551 loss_kl_train: 1012.4858 loss_x_train: 1.5539 loss_val: 1.5480 time: 5.7689s\n",
      " Epoch: 0316 loss_train: 1.6260 loss_kl_train: 997.0048 loss_x_train: 1.5263 loss_val: 1.5182 time: 5.8277s\n",
      " Epoch: 0317 loss_train: 1.6134 loss_kl_train: 833.8883 loss_x_train: 1.5300 loss_val: 1.5194 time: 5.8558s\n",
      " Epoch: 0318 loss_train: nan loss_kl_train: nan loss_x_train: nan loss_val: nan time: 5.4396s\n",
      " Epoch: 0319 loss_train: nan loss_kl_train: nan loss_x_train: nan loss_val: nan time: 5.1976s\n",
      " Epoch: 0320 loss_train: nan loss_kl_train: nan loss_x_train: nan loss_val: nan time: 5.1519s\n",
      " Epoch: 0321 loss_train: nan loss_kl_train: nan loss_x_train: nan loss_val: nan time: 5.1576s\n",
      " Epoch: 0322 loss_train: nan loss_kl_train: nan loss_x_train: nan loss_val: nan time: 5.2061s\n",
      " Epoch: 0323 loss_train: nan loss_kl_train: nan loss_x_train: nan loss_val: nan time: 5.1893s\n",
      " Epoch: 0324 loss_train: nan loss_kl_train: nan loss_x_train: nan loss_val: nan time: 5.1643s\n",
      " Epoch: 0325 loss_train: nan loss_kl_train: nan loss_x_train: nan loss_val: nan time: 5.1530s\n",
      " Epoch: 0326 loss_train: nan loss_kl_train: nan loss_x_train: nan loss_val: nan time: 5.1674s\n",
      " Epoch: 0327 loss_train: nan loss_kl_train: nan loss_x_train: nan loss_val: nan time: 5.1804s\n",
      " Epoch: 0328 loss_train: nan loss_kl_train: nan loss_x_train: nan loss_val: nan time: 5.1867s\n",
      " Epoch: 0329 loss_train: nan loss_kl_train: nan loss_x_train: nan loss_val: nan time: 5.2031s\n",
      " Epoch: 0330 loss_train: nan loss_kl_train: nan loss_x_train: nan loss_val: nan time: 5.2120s\n",
      " Epoch: 0331 loss_train: nan loss_kl_train: nan loss_x_train: nan loss_val: nan time: 5.2507s\n",
      " Epoch: 0332 loss_train: nan loss_kl_train: nan loss_x_train: nan loss_val: nan time: 5.1513s\n",
      " Epoch: 0333 loss_train: nan loss_kl_train: nan loss_x_train: nan loss_val: nan time: 5.1521s\n",
      " Epoch: 0334 loss_train: nan loss_kl_train: nan loss_x_train: nan loss_val: nan time: 5.2428s\n",
      " Epoch: 0335 loss_train: nan loss_kl_train: nan loss_x_train: nan loss_val: nan time: 5.1902s\n",
      " Epoch: 0336 loss_train: nan loss_kl_train: nan loss_x_train: nan loss_val: nan time: 5.2532s\n",
      " Epoch: 0337 loss_train: nan loss_kl_train: nan loss_x_train: nan loss_val: nan time: 5.2434s\n",
      " Epoch: 0338 loss_train: nan loss_kl_train: nan loss_x_train: nan loss_val: nan time: 5.2144s\n",
      " Epoch: 0339 loss_train: nan loss_kl_train: nan loss_x_train: nan loss_val: nan time: 5.1301s\n",
      " Epoch: 0340 loss_train: nan loss_kl_train: nan loss_x_train: nan loss_val: nan time: 5.1731s\n",
      " Epoch: 0341 loss_train: nan loss_kl_train: nan loss_x_train: nan loss_val: nan time: 5.1872s\n",
      " Epoch: 0342 loss_train: nan loss_kl_train: nan loss_x_train: nan loss_val: nan time: 5.0977s\n",
      " Epoch: 0343 loss_train: nan loss_kl_train: nan loss_x_train: nan loss_val: nan time: 5.0826s\n",
      " Epoch: 0344 loss_train: nan loss_kl_train: nan loss_x_train: nan loss_val: nan time: 5.1249s\n",
      " Epoch: 0345 loss_train: nan loss_kl_train: nan loss_x_train: nan loss_val: nan time: 5.1854s\n",
      " Epoch: 0346 loss_train: nan loss_kl_train: nan loss_x_train: nan loss_val: nan time: 5.1480s\n",
      " Epoch: 0347 loss_train: nan loss_kl_train: nan loss_x_train: nan loss_val: nan time: 5.2066s\n",
      " Epoch: 0348 loss_train: nan loss_kl_train: nan loss_x_train: nan loss_val: nan time: 5.2032s\n",
      " Epoch: 0349 loss_train: nan loss_kl_train: nan loss_x_train: nan loss_val: nan time: 5.1854s\n",
      " Epoch: 0350 loss_train: nan loss_kl_train: nan loss_x_train: nan loss_val: nan time: 5.2086s\n",
      " Epoch: 0351 loss_train: nan loss_kl_train: nan loss_x_train: nan loss_val: nan time: 5.1889s\n",
      " Epoch: 0352 loss_train: nan loss_kl_train: nan loss_x_train: nan loss_val: nan time: 5.2045s\n",
      " Epoch: 0353 loss_train: nan loss_kl_train: nan loss_x_train: nan loss_val: nan time: 5.2466s\n",
      " Epoch: 0354 loss_train: nan loss_kl_train: nan loss_x_train: nan loss_val: nan time: 5.2023s\n",
      " Epoch: 0355 loss_train: nan loss_kl_train: nan loss_x_train: nan loss_val: nan time: 5.1549s\n",
      " Epoch: 0356 loss_train: nan loss_kl_train: nan loss_x_train: nan loss_val: nan time: 5.2330s\n",
      " Epoch: 0357 loss_train: nan loss_kl_train: nan loss_x_train: nan loss_val: nan time: 5.1194s\n",
      " Epoch: 0358 loss_train: nan loss_kl_train: nan loss_x_train: nan loss_val: nan time: 5.1309s\n",
      " Epoch: 0359 loss_train: nan loss_kl_train: nan loss_x_train: nan loss_val: nan time: 5.1955s\n",
      " Epoch: 0360 loss_train: nan loss_kl_train: nan loss_x_train: nan loss_val: nan time: 5.2494s\n",
      " Epoch: 0361 loss_train: nan loss_kl_train: nan loss_x_train: nan loss_val: nan time: 5.2243s\n",
      " Epoch: 0362 loss_train: nan loss_kl_train: nan loss_x_train: nan loss_val: nan time: 5.1845s\n",
      " Epoch: 0363 loss_train: nan loss_kl_train: nan loss_x_train: nan loss_val: nan time: 5.2238s\n",
      " Epoch: 0364 loss_train: nan loss_kl_train: nan loss_x_train: nan loss_val: nan time: 5.2185s\n",
      " Epoch: 0365 loss_train: nan loss_kl_train: nan loss_x_train: nan loss_val: nan time: 5.2245s\n",
      " Epoch: 0366 loss_train: nan loss_kl_train: nan loss_x_train: nan loss_val: nan time: 5.2152s\n",
      " Epoch: 0367 loss_train: nan loss_kl_train: nan loss_x_train: nan loss_val: nan time: 5.1743s\n",
      " Epoch: 0368 loss_train: nan loss_kl_train: nan loss_x_train: nan loss_val: nan time: 5.2167s\n",
      " Epoch: 0369 loss_train: nan loss_kl_train: nan loss_x_train: nan loss_val: nan time: 5.2038s\n",
      " Epoch: 0370 loss_train: nan loss_kl_train: nan loss_x_train: nan loss_val: nan time: 5.1972s\n",
      " Epoch: 0371 loss_train: nan loss_kl_train: nan loss_x_train: nan loss_val: nan time: 5.2397s\n",
      " Epoch: 0372 loss_train: nan loss_kl_train: nan loss_x_train: nan loss_val: nan time: 5.2127s\n",
      " Epoch: 0373 loss_train: nan loss_kl_train: nan loss_x_train: nan loss_val: nan time: 5.1370s\n",
      " Epoch: 0374 loss_train: nan loss_kl_train: nan loss_x_train: nan loss_val: nan time: 5.2103s\n",
      " Epoch: 0375 loss_train: nan loss_kl_train: nan loss_x_train: nan loss_val: nan time: 5.2398s\n",
      " Epoch: 0376 loss_train: nan loss_kl_train: nan loss_x_train: nan loss_val: nan time: 5.2666s\n",
      " Epoch: 0377 loss_train: nan loss_kl_train: nan loss_x_train: nan loss_val: nan time: 5.2192s\n",
      " Epoch: 0378 loss_train: nan loss_kl_train: nan loss_x_train: nan loss_val: nan time: 5.2508s\n",
      " Epoch: 0379 loss_train: nan loss_kl_train: nan loss_x_train: nan loss_val: nan time: 5.2320s\n",
      " Epoch: 0380 loss_train: nan loss_kl_train: nan loss_x_train: nan loss_val: nan time: 5.1992s\n",
      " Epoch: 0381 loss_train: nan loss_kl_train: nan loss_x_train: nan loss_val: nan time: 5.1601s\n",
      " Epoch: 0382 loss_train: nan loss_kl_train: nan loss_x_train: nan loss_val: nan time: 5.2524s\n",
      " Epoch: 0383 loss_train: nan loss_kl_train: nan loss_x_train: nan loss_val: nan time: 5.2309s\n",
      " Epoch: 0384 loss_train: nan loss_kl_train: nan loss_x_train: nan loss_val: nan time: 5.1430s\n",
      " Epoch: 0385 loss_train: nan loss_kl_train: nan loss_x_train: nan loss_val: nan time: 5.1894s\n",
      " Epoch: 0386 loss_train: nan loss_kl_train: nan loss_x_train: nan loss_val: nan time: 5.2254s\n",
      " Epoch: 0387 loss_train: nan loss_kl_train: nan loss_x_train: nan loss_val: nan time: 5.2081s\n",
      " Epoch: 0388 loss_train: nan loss_kl_train: nan loss_x_train: nan loss_val: nan time: 5.2231s\n",
      " Epoch: 0389 loss_train: nan loss_kl_train: nan loss_x_train: nan loss_val: nan time: 5.2472s\n",
      " Epoch: 0390 loss_train: nan loss_kl_train: nan loss_x_train: nan loss_val: nan time: 5.2085s\n",
      " Epoch: 0391 loss_train: nan loss_kl_train: nan loss_x_train: nan loss_val: nan time: 5.1148s\n",
      " Epoch: 0392 loss_train: nan loss_kl_train: nan loss_x_train: nan loss_val: nan time: 5.1637s\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " Epoch: 0393 loss_train: nan loss_kl_train: nan loss_x_train: nan loss_val: nan time: 5.2414s\n",
      " Epoch: 0394 loss_train: nan loss_kl_train: nan loss_x_train: nan loss_val: nan time: 5.2098s\n",
      " Epoch: 0395 loss_train: nan loss_kl_train: nan loss_x_train: nan loss_val: nan time: 5.2142s\n",
      " Epoch: 0396 loss_train: nan loss_kl_train: nan loss_x_train: nan loss_val: nan time: 5.2641s\n",
      " Epoch: 0397 loss_train: nan loss_kl_train: nan loss_x_train: nan loss_val: nan time: 5.2642s\n",
      " Epoch: 0398 loss_train: nan loss_kl_train: nan loss_x_train: nan loss_val: nan time: 5.1533s\n",
      " Epoch: 0399 loss_train: nan loss_kl_train: nan loss_x_train: nan loss_val: nan time: 5.1668s\n",
      " Epoch: 0400 loss_train: nan loss_kl_train: nan loss_x_train: nan loss_val: nan time: 5.1698s\n",
      " Epoch: 0401 loss_train: nan loss_kl_train: nan loss_x_train: nan loss_val: nan time: 5.1616s\n",
      " Epoch: 0402 loss_train: nan loss_kl_train: nan loss_x_train: nan loss_val: nan time: 5.1361s\n",
      " Epoch: 0403 loss_train: nan loss_kl_train: nan loss_x_train: nan loss_val: nan time: 5.2461s\n",
      " Epoch: 0404 loss_train: nan loss_kl_train: nan loss_x_train: nan loss_val: nan time: 5.2145s\n",
      " Epoch: 0405 loss_train: nan loss_kl_train: nan loss_x_train: nan loss_val: nan time: 5.1969s\n",
      " Epoch: 0406 loss_train: nan loss_kl_train: nan loss_x_train: nan loss_val: nan time: 5.2075s\n",
      " Epoch: 0407 loss_train: nan loss_kl_train: nan loss_x_train: nan loss_val: nan time: 5.2296s\n",
      " Epoch: 0408 loss_train: nan loss_kl_train: nan loss_x_train: nan loss_val: nan time: 5.1395s\n",
      " Epoch: 0409 loss_train: nan loss_kl_train: nan loss_x_train: nan loss_val: nan time: 5.1861s\n",
      " Epoch: 0410 loss_train: nan loss_kl_train: nan loss_x_train: nan loss_val: nan time: 5.2697s\n",
      " Epoch: 0411 loss_train: nan loss_kl_train: nan loss_x_train: nan loss_val: nan time: 5.2193s\n",
      " Epoch: 0412 loss_train: nan loss_kl_train: nan loss_x_train: nan loss_val: nan time: 5.2043s\n",
      " Epoch: 0413 loss_train: nan loss_kl_train: nan loss_x_train: nan loss_val: nan time: 5.2327s\n",
      " Epoch: 0414 loss_train: nan loss_kl_train: nan loss_x_train: nan loss_val: nan time: 5.2158s\n",
      " Epoch: 0415 loss_train: nan loss_kl_train: nan loss_x_train: nan loss_val: nan time: 5.2798s\n",
      " Epoch: 0416 loss_train: nan loss_kl_train: nan loss_x_train: nan loss_val: nan time: 5.1740s\n",
      " Epoch: 0417 loss_train: nan loss_kl_train: nan loss_x_train: nan loss_val: nan time: 5.2208s\n",
      " Epoch: 0418 loss_train: nan loss_kl_train: nan loss_x_train: nan loss_val: nan time: 5.2555s\n",
      " Epoch: 0419 loss_train: nan loss_kl_train: nan loss_x_train: nan loss_val: nan time: 5.2292s\n",
      " Epoch: 0420 loss_train: nan loss_kl_train: nan loss_x_train: nan loss_val: nan time: 5.1445s\n",
      " Epoch: 0421 loss_train: nan loss_kl_train: nan loss_x_train: nan loss_val: nan time: 5.2081s\n",
      " Epoch: 0422 loss_train: nan loss_kl_train: nan loss_x_train: nan loss_val: nan time: 5.2179s\n",
      " Epoch: 0423 loss_train: nan loss_kl_train: nan loss_x_train: nan loss_val: nan time: 5.1646s\n",
      " Epoch: 0424 loss_train: nan loss_kl_train: nan loss_x_train: nan loss_val: nan time: 5.1666s\n",
      " Epoch: 0425 loss_train: nan loss_kl_train: nan loss_x_train: nan loss_val: nan time: 5.1767s\n",
      " Epoch: 0426 loss_train: nan loss_kl_train: nan loss_x_train: nan loss_val: nan time: 5.1737s\n",
      " Epoch: 0427 loss_train: nan loss_kl_train: nan loss_x_train: nan loss_val: nan time: 5.2374s\n",
      " Epoch: 0428 loss_train: nan loss_kl_train: nan loss_x_train: nan loss_val: nan time: 5.2567s\n",
      " Epoch: 0429 loss_train: nan loss_kl_train: nan loss_x_train: nan loss_val: nan time: 5.2669s\n",
      " Epoch: 0430 loss_train: nan loss_kl_train: nan loss_x_train: nan loss_val: nan time: 5.2206s\n",
      " Epoch: 0431 loss_train: nan loss_kl_train: nan loss_x_train: nan loss_val: nan time: 5.1510s\n",
      " Epoch: 0432 loss_train: nan loss_kl_train: nan loss_x_train: nan loss_val: nan time: 5.1675s\n",
      " Epoch: 0433 loss_train: nan loss_kl_train: nan loss_x_train: nan loss_val: nan time: 5.1975s\n",
      " Epoch: 0434 loss_train: nan loss_kl_train: nan loss_x_train: nan loss_val: nan time: 5.2091s\n",
      " Epoch: 0435 loss_train: nan loss_kl_train: nan loss_x_train: nan loss_val: nan time: 5.2203s\n",
      " Epoch: 0436 loss_train: nan loss_kl_train: nan loss_x_train: nan loss_val: nan time: 5.2275s\n",
      " Epoch: 0437 loss_train: nan loss_kl_train: nan loss_x_train: nan loss_val: nan time: 5.2210s\n",
      " Epoch: 0438 loss_train: nan loss_kl_train: nan loss_x_train: nan loss_val: nan time: 5.1820s\n",
      " Epoch: 0439 loss_train: nan loss_kl_train: nan loss_x_train: nan loss_val: nan time: 5.2087s\n",
      " Epoch: 0440 loss_train: nan loss_kl_train: nan loss_x_train: nan loss_val: nan time: 5.1526s\n",
      " Epoch: 0441 loss_train: nan loss_kl_train: nan loss_x_train: nan loss_val: nan time: 5.1900s\n",
      " Epoch: 0442 loss_train: nan loss_kl_train: nan loss_x_train: nan loss_val: nan time: 5.1790s\n"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-12-79b0f854e412>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m    103\u001b[0m \u001b[0;31m# for ep in range(10000,20000):\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    104\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 105\u001b[0;31m     \u001b[0mtrain_loss_ep\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mep\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mtrain_loss_kl_ep\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mep\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mtrain_loss_x_ep\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mep\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mval_loss_ep\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mep\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mtrain\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mep\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    106\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    107\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m<ipython-input-12-79b0f854e412>\u001b[0m in \u001b[0;36mtrain\u001b[0;34m(epoch)\u001b[0m\n\u001b[1;32m     46\u001b[0m         \u001b[0mloss\u001b[0m\u001b[0;34m=\u001b[0m \u001b[0mloss_kl_train\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0;36m0.0001\u001b[0m\u001b[0;34m+\u001b[0m\u001b[0mloss_x_train\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     47\u001b[0m         \u001b[0mloss\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mbackward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 48\u001b[0;31m         \u001b[0moptimizerVAEXA\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstep\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     49\u001b[0m         \u001b[0mloss_all\u001b[0m\u001b[0;34m+=\u001b[0m\u001b[0mloss\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mitem\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     50\u001b[0m         \u001b[0mloss_x_all\u001b[0m\u001b[0;34m+=\u001b[0m\u001b[0mloss_x_train\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mitem\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/lib/python3/dist-packages/torch/optim/optimizer.py\u001b[0m in \u001b[0;36mwrapper\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m     86\u001b[0m                 \u001b[0mprofile_name\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m\"Optimizer.step#{}.step\"\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mformat\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mobj\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m__class__\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m__name__\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     87\u001b[0m                 \u001b[0;32mwith\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mautograd\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mprofiler\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mrecord_function\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mprofile_name\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 88\u001b[0;31m                     \u001b[0;32mreturn\u001b[0m \u001b[0mfunc\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     89\u001b[0m             \u001b[0;32mreturn\u001b[0m \u001b[0mwrapper\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     90\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/lib/python3/dist-packages/torch/autograd/grad_mode.py\u001b[0m in \u001b[0;36mdecorate_context\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m     26\u001b[0m         \u001b[0;32mdef\u001b[0m \u001b[0mdecorate_context\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     27\u001b[0m             \u001b[0;32mwith\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m__class__\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 28\u001b[0;31m                 \u001b[0;32mreturn\u001b[0m \u001b[0mfunc\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     29\u001b[0m         \u001b[0;32mreturn\u001b[0m \u001b[0mcast\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mF\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdecorate_context\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     30\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/lib/python3/dist-packages/torch/optim/adam.py\u001b[0m in \u001b[0;36mstep\u001b[0;34m(self, closure)\u001b[0m\n\u001b[1;32m    105\u001b[0m                     \u001b[0mstate_steps\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mappend\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mstate\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m'step'\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    106\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 107\u001b[0;31m             F.adam(params_with_grad,\n\u001b[0m\u001b[1;32m    108\u001b[0m                    \u001b[0mgrads\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    109\u001b[0m                    \u001b[0mexp_avgs\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/lib/python3/dist-packages/torch/optim/_functional.py\u001b[0m in \u001b[0;36madam\u001b[0;34m(params, grads, exp_avgs, exp_avg_sqs, max_exp_avg_sqs, state_steps, amsgrad, beta1, beta2, lr, weight_decay, eps)\u001b[0m\n\u001b[1;32m     92\u001b[0m             \u001b[0mdenom\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0mmax_exp_avg_sqs\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mi\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msqrt\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m/\u001b[0m \u001b[0mmath\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msqrt\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mbias_correction2\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0madd_\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0meps\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     93\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 94\u001b[0;31m             \u001b[0mdenom\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0mexp_avg_sq\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msqrt\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m/\u001b[0m \u001b[0mmath\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msqrt\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mbias_correction2\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0madd_\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0meps\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     95\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     96\u001b[0m         \u001b[0mstep_size\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mlr\u001b[0m \u001b[0;34m/\u001b[0m \u001b[0mbias_correction1\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "def train(epoch):\n",
    "    t = time.time()\n",
    "    model.train()\n",
    "    \n",
    "    loss_all=0\n",
    "    loss_x_all=0\n",
    "    loss_kl_all=0\n",
    "    for i in range(int(np.ceil(trainIdx.shape[0]/batchsize))):\n",
    "        trainIdx_i=trainIdx[i*batchsize:min((i+1)*batchsize,trainIdx.shape[0])]\n",
    "        \n",
    "    \n",
    "        optimizerVAEXA.zero_grad()\n",
    "\n",
    "        features=data_train[trainIdx_i].cuda(0).float()\n",
    "        if 'dca' in model_str:\n",
    "            features_raw=data_raw[trainIdx_i].cuda(0)\n",
    "\n",
    "        if use_cuda:\n",
    "            train_nodes_idx=torch.tensor(np.repeat(True,trainIdx_i.size)).cuda(0)\n",
    "\n",
    "        adj_recon,mu,logvar,z,features_recon = model(features, None)\n",
    "\n",
    "        adj_recon=adj_recon.cuda(0)\n",
    "        mu=mu.cuda(0)\n",
    "        logvar=logvar.cuda(0)\n",
    "        z=z.cuda(0)\n",
    "\n",
    "        loss_kl_train=loss_kl(mu, logvar, train_nodes_idx)\n",
    "\n",
    "        if 'dca' in model_str:\n",
    "            if 'NB' in name:\n",
    "                loss_x_train=loss_x(features_recon, features,train_nodes_idx,XreconWeight)\n",
    "            else:\n",
    "                loss_x_train=loss_x(features_recon, features,train_nodes_idx,XreconWeight,ridgeL,features_raw)\n",
    "        else:\n",
    "            loss_x_train=loss_x(features_recon, features,train_nodes_idx,XreconWeight,mse)\n",
    "    \n",
    "    \n",
    "        loss= loss_kl_train*0.0001+loss_x_train\n",
    "        loss.backward()\n",
    "        optimizerVAEXA.step()\n",
    "        loss_all+=loss.item()\n",
    "        loss_x_all+=loss_x_train.item()\n",
    "        loss_kl_all+=loss_kl_train.item()\n",
    "        \n",
    "    loss_all=loss_all/int(np.ceil(trainIdx.shape[0]/batchsize))\n",
    "    loss_x_all=loss_x_all/int(np.ceil(trainIdx.shape[0]/batchsize))\n",
    "    loss_kl_all=loss_kl_all/int(np.ceil(trainIdx.shape[0]/batchsize))\n",
    "\n",
    "    with torch.no_grad():\n",
    "        model.eval()\n",
    "        loss_val_all=0\n",
    "        nvalBatches=int(np.ceil(valIdx.shape[0]/batchsize))\n",
    "        for i in range(nvalBatches):\n",
    "            valIdx_i=valIdx[i*batchsize:min((i+1)*batchsize,valIdx.shape[0])]\n",
    "            \n",
    "            features=data_train[valIdx_i].cuda(0).float()\n",
    "            if 'dca' in model_str:\n",
    "                features_raw=data_raw[valIdx_i].cuda(0)\n",
    "            adj_recon,mu,logvar,z, features_recon = model(features,None)\n",
    "            adj_recon=adj_recon.cuda(0)\n",
    "            mu=mu.cuda(0)\n",
    "            logvar=logvar.cuda(0)\n",
    "            z=z.cuda(0)\n",
    "    \n",
    "            if use_cuda:\n",
    "                val_nodes_idx=torch.tensor(np.repeat(True,valIdx_i.size)).cuda(0)\n",
    "            if 'dca' in model_str:\n",
    "                if 'NB' in name:\n",
    "                    loss_x_val=loss_x(features_recon, features,val_nodes_idx,XreconWeight)\n",
    "                else:\n",
    "                    loss_x_val=loss_x(features_recon, features,val_nodes_idx,XreconWeight,ridgeL,features_raw)\n",
    "            else:\n",
    "                loss_x_val=loss_x(features_recon, features,val_nodes_idx,XreconWeight,mse)\n",
    "\n",
    "\n",
    "            loss_val_all += loss_x_val.item()\n",
    "        loss_val_all=loss_val_all/nvalBatches\n",
    "    print(' Epoch: {:04d}'.format(epoch),\n",
    "          'loss_train: {:.4f}'.format(loss_all),\n",
    "          'loss_kl_train: {:.4f}'.format(loss_kl_all),\n",
    "          'loss_x_train: {:.4f}'.format(loss_x_all),\n",
    "          'loss_val: {:.4f}'.format(loss_val_all),\n",
    "          'time: {:.4f}s'.format(time.time() - t))\n",
    "    return loss_all,loss_kl_all,loss_x_all,loss_val_all\n",
    "\n",
    "    \n",
    "# print('cross-validation ',seti)\n",
    "train_loss_ep=[None]*epochs\n",
    "train_loss_kl_ep=[None]*epochs\n",
    "train_loss_x_ep=[None]*epochs\n",
    "val_loss_ep=[None]*epochs\n",
    "t_ep=time.time()\n",
    "\n",
    "for ep in range(epochs):\n",
    "# for ep in range(10000,20000):\n",
    "    \n",
    "    train_loss_ep[ep],train_loss_kl_ep[ep],train_loss_x_ep[ep],val_loss_ep[ep]=train(ep)\n",
    "\n",
    "        \n",
    "    if ep%saveFreq == (saveFreq-1):\n",
    "        torch.save(model.state_dict(), os.path.join(modelsavepath,str(ep)+'.pt'))\n",
    "    if use_cuda:\n",
    "#         model.cuda()\n",
    "        torch.cuda.empty_cache()\n",
    "print(' total time: {:.4f}s'.format(time.time() - t_ep))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.save(model.state_dict(), os.path.join(modelsavepath,str(ep)+'.pt'))\n",
    "torch.cuda.empty_cache()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'os' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-1-14845424754c>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m     22\u001b[0m \u001b[0;31m#     val_loss_advD_ep[:5911]=pickle.load(input)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     23\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 24\u001b[0;31m \u001b[0;32mwith\u001b[0m \u001b[0mopen\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mos\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mpath\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mjoin\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mlogsavepath\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m'train_loss'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m'wb'\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mas\u001b[0m \u001b[0moutput\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     25\u001b[0m     \u001b[0mpickle\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdump\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtrain_loss_ep\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0moutput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mpickle\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mHIGHEST_PROTOCOL\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     26\u001b[0m \u001b[0;32mwith\u001b[0m \u001b[0mopen\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mos\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mpath\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mjoin\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mlogsavepath\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m'train_loss_kl'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m'wb'\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mas\u001b[0m \u001b[0moutput\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mNameError\u001b[0m: name 'os' is not defined"
     ]
    }
   ],
   "source": [
    "with open(os.path.join(logsavepath,'train_loss'), 'wb') as output:\n",
    "    pickle.dump(train_loss_ep, output, pickle.HIGHEST_PROTOCOL)\n",
    "with open(os.path.join(logsavepath,'train_loss_kl'), 'wb') as output:\n",
    "    pickle.dump(train_loss_kl_ep, output, pickle.HIGHEST_PROTOCOL)\n",
    "with open(os.path.join(logsavepath,'train_loss_x'), 'wb') as output:\n",
    "    pickle.dump(train_loss_x_ep, output, pickle.HIGHEST_PROTOCOL)\n",
    "with open(os.path.join(logsavepath,'val_loss'), 'wb') as output:\n",
    "    pickle.dump(val_loss_ep, output, pickle.HIGHEST_PROTOCOL)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "plt.plot(np.arange(epochs),train_loss_x_ep)\n",
    "plt.plot(np.arange(epochs),val_loss_ep)\n",
    "plt.plot(np.arange(epochs),train_loss_kl_ep)\n",
    "\n",
    "plt.ylim((0,5))\n",
    "# plt.xlim((0,3000))\n",
    "plt.legend(['training x recon loss','validation x recon loss','training kl loss'],loc='upper right')\n",
    "plt.savefig(os.path.join(plotsavepath,'loss_seed6_zoom.jpg'))\n",
    "plt.show()\n"
   ]
  }
 ],
 "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.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}