Pred_Traj_2EncDecTransf_CA_eth.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import sqlite3\n",
    "import time\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "import torchvision\n",
    "from torchvision.utils import save_image\n",
    "from torch.distributions.normal import Normal\n",
    "\n",
    "from PIL import Image\n",
    "\n",
    "import math\n",
    "import os\n",
    "import numpy as np\n",
    "import random\n",
    "import matplotlib.pyplot as plt\n",
    "from skimage.io import imread, imshow\n",
    "import datetime\n",
    "from skimage.util import img_as_ubyte\n",
    "import seaborn as sns\n",
    "import warnings\n",
    "warnings.filterwarnings('ignore')\n",
    "import argparse\n",
    "\n",
    "from datetime import datetime\n",
    "\n",
    "from models.indiv_crossAttention_seq2seq_code_size import crossAttention\n",
    "\n",
    "from IPython import display\n",
    "\n",
    "from torch.utils.tensorboard import SummaryWriter\n",
    "\n",
    "stochastic_mode         = 1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Make folder for outputs and logs\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Dataset name\n",
    "dataset_name = \"eth\" # dataset options: 'university', 'zara_01', 'zara_02', 'eth', 'hotel'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "__file__=os.getcwd()\n",
    "print(__file__)\n",
    "\n",
    "now = datetime.now() # current date and time\n",
    "current_time_date = now.strftime(\"%d_%m_%y_%H_%M_%S\")\n",
    "run_folder  = \"Outputs/traj_pred_\"+ dataset_name + \"_\" + str(os.path.basename(__file__)) + str(current_time_date)\n",
    "os.makedirs(run_folder)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Skip to left side bar\n",
    ">\n",
    "/\n",
    "Name\n",
    "Last Modified\n",
    "\n",
    "# Make log folder for tensorboard"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "SummaryWriter_path = \"2Encod_Transf_decoder_eth_ca/\"\n",
    "os.makedirs(SummaryWriter_path)   \n",
    "writer = SummaryWriter(SummaryWriter_path,comment=\"ADE_FDE_Train\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Make image folder to save outputs\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "image_path  = run_folder + \"/Visual_Prediction\"\n",
    "os.makedirs(image_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# GPU\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "device = torch.device('cuda:0')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "os.environ['CUDA_VISIBLE_DEVICES'] = '0'\n",
    "torch.cuda.device_count()\n",
    "torch.cuda.current_device()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Variables\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#DataBase Variables\n",
    "image_folder_path = 'data/data_trajpred/'+dataset_name\n",
    "DB_PATH_train = \"./data/data_trajpred/\"+dataset_name+\"/pos_data_train.db\"\n",
    "cnx_train    = sqlite3.connect(DB_PATH_train)\n",
    "DB_PATH_val  = \"./data/data_trajpred/\"+dataset_name+\"/pos_data_val.db\"\n",
    "cnx_val      = sqlite3.connect(DB_PATH_val)\n",
    "DB_DIR       = run_folder + '/database'\n",
    "os.makedirs( DB_DIR )\n",
    "DB_PATH2     = DB_DIR+'/db_one_ped_delta_coordinates_results.db'\n",
    "cnx2         = sqlite3.connect(DB_PATH2)\n",
    "\n",
    "#Other variables\n",
    "T_obs                   = 8\n",
    "T_pred                  = 12\n",
    "T_total                 = T_obs + T_pred #8+12=20\n",
    "data_id                 = 0 \n",
    "batch_size              = 40\n",
    "chunk_size              = batch_size * T_total # Chunksize should be multiple of T_total\n",
    "in_size                 = 2\n",
    "stochastic_out_size     = in_size * 2\n",
    "hidden_size             = 256 \n",
    "embed_size              = 64 \n",
    "global dropout_val\n",
    "dropout_val             = 0.2 #0.5\n",
    "teacher_forcing_ratio   = 0.7 # 0.9\n",
    "regularization_factor   = 0.5 # 0.001\n",
    "avg_n_path_eval         = 20\n",
    "bst_n_path_eval         = 20\n",
    "path_mode               = \"top5\" #\"avg\",\"bst\",\"single\",\"top5\"\n",
    "regularization_mode     = \"regular\" #\"weighted\",\"e_weighted\", \"regular\"\n",
    "startpoint_mode         = \"on\" #\"on\",\"off\"\n",
    "enc_out                 = \"on\" #\"on\",\"off\"\n",
    "biased_loss_mode        = 0 # 0 , 1\n",
    "\n",
    "table_out   = \"results_delta\"\n",
    "table       = \"dataset_T_length_20delta_coordinates\"\n",
    "df_id       = pd.read_sql_query(\"SELECT data_id FROM \"+table, cnx_train)\n",
    "data_size   = df_id.data_id.max() * T_total\n",
    "epoch_num   = 400\n",
    "from_epoch  = 0\n",
    "\n",
    "#Visual Variables\n",
    "image_size              = 256  \n",
    "image_dimension         = 3\n",
    "mask_size               = 16\n",
    "visual_features_size    = 128 \n",
    "visual_embed_size       = 64\n",
    "vsn_module_out_size    = 256\n",
    "to_pil = torchvision.transforms.ToPILImage()\n",
    "\n",
    "#Model Path\n",
    "model_path = run_folder + \"/NNmodel\" \n",
    "os.makedirs(model_path)   \n",
    "model_path = model_path + str(\"/model\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Handle Sequential Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class TrajectoryPredictionDataset(torch.utils.data.Dataset): \n",
    "    def __init__(self, ROOT_DIR, DB_PATH, cnx):\n",
    "        self.pos_df    = pd.read_sql_query(\"SELECT * FROM \"+str(table), cnx)\n",
    "        self.root_dir  = ROOT_DIR+'/visual_data'\n",
    "        self.transform = torchvision.transforms.Compose([torchvision.transforms.Resize((image_size,image_size)), \\\n",
    "                                                         torchvision.transforms.ToTensor(), \\\n",
    "                                                         torchvision.transforms.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225))])\n",
    "        self.visual_data = []\n",
    "        # read sorted frames\n",
    "        for img in sorted(os.listdir(self.root_dir)): \n",
    "            self.visual_data.append(self.transform( Image.open(os.path.join(self.root_dir)+\"/\"+img) ))\n",
    "        self.visual_data = torch.stack(self.visual_data)  \n",
    "\n",
    "    def __len__(self):\n",
    "        return self.pos_df.data_id.max() #data_id maximum dans dataset\n",
    "    \n",
    "    def __getitem__(self, idx):\n",
    "        if torch.is_tensor(idx):\n",
    "            idx = idx.tolist() \n",
    "        \n",
    "        extracted_df     = self.pos_df[ self.pos_df[\"data_id\"] == idx ]\n",
    "        tensor           = torch.tensor(extracted_df[['pos_x_delta','pos_y_delta']].values).reshape(-1,T_total,in_size)\n",
    "        obs, pred        = torch.split(tensor,[T_obs,T_pred],dim=1)\n",
    "        start_frames     = (extracted_df.groupby('data_id').frame_num.min().values/10).astype('int')\n",
    "        for i in start_frames:            \n",
    "            extracted_frames.append(self.visual_data[i:i+T_obs])\n",
    "        frames = torch.stack(extracted_frames) #stack concatenates a sequence of tensors along a new dimension.\n",
    "        start_frames = torch.tensor(start_frames)\n",
    "        return obs, pred, frames, start_frames"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Initialize random weights for NN models\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def init_weights(m):\n",
    "    for name, param in m.named_parameters():\n",
    "        nn.init.uniform_(param.data, -0.2, 0.2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Regularizer loss\n",
    "The distance_from_line_regularizer function is implemented as a regularization method to enforce the predicted trajectory to stay close to the observed trajectory. The method is used to minimize the Euclidean distances between each step of the predicted trajectory and a line fitted to the observed trajectory.\n",
    "It takes 2 inputs:\n",
    "-input_tensor : The observed trajectory in the form of a tensor of size (batch_size, T_obs, 2)\n",
    "-prediction : The predicted trajectory in the form of a tensor of size (batch_size, T_pred, 2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sum_sigma_distance  = torch.zeros(1)\n",
    "\n",
    "def distance_from_line_regularizer(input_tensor,prediction):\n",
    "    global sum_sigma_distance\n",
    "    #the regularization is defined as the sum of Euclidean distances between each step of the predicted trajectory Tf , and a line fitted to the observed trajectory To.\n",
    "    # Fit a line to observation points over batch \n",
    "    input_tensor    = input_tensor.double()\n",
    "    prediction      = prediction.double()\n",
    "    input_tensor    = input_tensor.cumsum(dim=1).double()\n",
    "    X               = torch.ones_like(input_tensor).double().to('cuda', non_blocking=True)\n",
    "    X[:,:,0]        = input_tensor[:,:,0]\n",
    "    Y               = (input_tensor[:,:,1]).unsqueeze(-1).double()\n",
    "    try:\n",
    "        try:\n",
    "            XTX_1       = torch.matmul( X.transpose(-1,-2), X).double().inverse()\n",
    "        except:\n",
    "            XTX_1       = torch.matmul( X.transpose(-1,-2), X).double().pinverse()\n",
    "        XTY             = torch.matmul( X.transpose(-1,-2), Y)\n",
    "        theta           = torch.matmul( XTX_1.double(), XTY.double())\n",
    "        # Calculate real values of prediction instead of delta\n",
    "        prediction[:,:,0] = prediction[:,:,0] + input_tensor[:,-1,0].unsqueeze(-1) \n",
    "        prediction[:,:,1] = prediction[:,:,1] + input_tensor[:,-1,1].unsqueeze(-1)\n",
    "        \n",
    "        # Calculate distance ( predicted_points , observation_fitted_line ) over batch\n",
    "        theta0x0        = theta[:,0,:].double() * prediction[:,:,0].double()\n",
    "        denominator     = torch.sqrt( theta[:,0,:].double() * theta[:,0,:].double() + 1 )\n",
    "        nominator       = theta0x0 + theta[:,1,:] - prediction[:,:,1].double()\n",
    "        distance        = nominator.abs() / denominator\n",
    "        if regularization_mode =='weighted':\n",
    "            weight              = torch.flip( torch.arange(1,T_pred+1).cuda().float(),[0])\n",
    "            weight              = (weight / T_pred).repeat(distance.size(0)).view(-1,T_pred)\n",
    "            weighted_distance   = weight * distance\n",
    "\n",
    "        elif regularization_mode =='e_weighted':\n",
    "            weight              = torch.flip( torch.arange(1,T_pred+1).cuda().float(),[0])\n",
    "            weight              = (weight / T_pred).repeat(distance.size(0)).view(distance.size(0),T_pred)\n",
    "            weight              = torch.exp(weight)\n",
    "            weighted_distance   = weight*distance\n",
    "\n",
    "        else:\n",
    "            weighted_distance = distance\n",
    "        sigma_distance  = torch.mean(weighted_distance,1)\n",
    "        sum_sigma_distance  = torch.mean(sigma_distance)\n",
    "        return sum_sigma_distance\n",
    "    except:\n",
    "        print(\"SINGULAR VALUE\")\n",
    "        sum_sigma_distance = torch.zeros(1).to('cuda', non_blocking=True) + 20\n",
    "        return sum_sigma_distance"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Vision Transformer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Spatial Features Extraction (Resnet 18)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.utils.data import DataLoader \n",
    "from torchvision import datasets \n",
    "from torchvision.transforms import ToTensor \n",
    "from torchsummary import summary"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def Resnet(pretrain=True,layers_to_unfreeze=8,layers_to_delete=2,in_planes=3):\n",
    "    \"\"\"\n",
    "    param:\n",
    "        pretrain: Define if we load a pretrained model from ImageNet\n",
    "        layers_to_unfreeze: Define the number of layers that we want to train at the end of the Resnet\n",
    "        layers_to_delete: Define the numbers of layers that we want to delete\n",
    "        in_planes: Define the numbers of input channels of images (supported values: 1,2 or 3)\n",
    "    return: The Resnet model\n",
    "    \"\"\"\n",
    "    resnet = torchvision.models.resnet18(pretrained=pretrain)\n",
    "    # Create a new model cause we don't want the pooling operation at the end and the classifier\n",
    "    model = nn.Sequential()\n",
    "    number_of_layers = len(list(resnet.children())) - layers_to_delete # In practice it remove the pooling operation and the classifier\n",
    "\n",
    "    if number_of_layers<layers_to_unfreeze:\n",
    "        layers_to_unfreeze = number_of_layers\n",
    "    layers_to_freeze = number_of_layers - layers_to_unfreeze\n",
    "    i=0\n",
    "    for child in resnet.children():\n",
    "        # For the first layers we create a new weight if in_planes is not 3 cause ResNet is pretrain on image with 3 channels there is no version for 1 channel\n",
    "        if i==0 and in_planes<3:\n",
    "            if i<layers_to_freeze: # Define if we freeze this layer or no\n",
    "                for param in child.parameters():\n",
    "                    param.requires_grad = False # Freeze the layers by passing requires_grad attribute to False\n",
    "            w = child._parameters['weight'].data # Get the weight for 3 channels data\n",
    "            child._modules['0'] = nn.Conv2d(in_planes, 64, kernel_size=3, padding=1) # Define the new conv layer\n",
    "            if in_planes == 1:\n",
    "                child._parameters['weight'].data = w.mean(dim=1, keepdim=True) # If the number of channels is 1 we made the mean of channels to set the new weight\n",
    "            else:\n",
    "                child._parameters['weight'].data = w[:, :-1] * 1.5\n",
    "\n",
    "        if i<layers_to_freeze: # Define if we freeze this layer or no\n",
    "            for param in child.parameters():\n",
    "                param.requires_grad = False # Freeze the layers by passing requires_grad attribute to False\n",
    "        if i<number_of_layers: # To define if we keep this layer or not\n",
    "            model.append(child) \n",
    "        i+=1\n",
    "    return model\n",
    "\n",
    "\n",
    "class features_extraction(nn.Module):\n",
    "    \"\"\"\n",
    "    param:\n",
    "    conv_model: The convolution model used before capsules for the moment only ResNet is supported\n",
    "    in_planes: Numbers of channels for the image\n",
    "    \"\"\"\n",
    "    def __init__(self,conv_model,in_planes: int):\n",
    "        super().__init__()\n",
    "        self.conv_model = conv_model\n",
    "        self.in_planes = in_planes\n",
    "        self.pooling = nn.AdaptiveAvgPool2d((1,1))\n",
    "        \n",
    "    def forward(self,input):\n",
    "        shape = input.size()\n",
    "        x = input.view(-1,self.in_planes,shape[-2],shape[-1])\n",
    "        x = self.conv_model(x)\n",
    "        x = self.pooling(x)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Transformer : EncoderSelfAttention"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def position_embedding(input, d_model):\n",
    "    input = input.view(-1, 1)\n",
    "    dim = torch.arange(d_model // 2, dtype=torch.float32, device=input.device).view(1, -1)\n",
    "    sin = torch.sin(input / 10000 ** (2 * dim / d_model))\n",
    "    cos = torch.cos(input / 10000 ** (2 * dim / d_model))\n",
    "\n",
    "    out = torch.zeros((input.shape[0], d_model), device=input.device)\n",
    "    out[:, ::2] = sin\n",
    "    out[:, 1::2] = cos\n",
    "    return out\n",
    "\n",
    "def sinusoid_encoding_table(max_len, d_model):\n",
    "    pos = torch.arange(max_len, dtype=torch.float32)\n",
    "    out = position_embedding(pos, d_model)\n",
    "    return out\n",
    "\n",
    "class ScaledDotProductAttention(nn.Module):\n",
    "    def __init__(self, d_model, d_k, d_v, h):\n",
    "        \"\"\"\n",
    "        param:\n",
    "        d_model: Output dimensionality of the model\n",
    "        d_k: Dimensionality of queries and keys\n",
    "        d_v: Dimensionality of values\n",
    "        h: Number of heads\n",
    "        \"\"\"\n",
    "        super(ScaledDotProductAttention, self).__init__()\n",
    "        self.fc_q = nn.Linear(d_model, h * d_k)\n",
    "        self.fc_k = nn.Linear(d_model, h * d_k)\n",
    "        self.fc_v = nn.Linear(d_model, h * d_v)\n",
    "        self.fc_o = nn.Linear(h * d_v, d_model)\n",
    "\n",
    "        self.d_model = d_model\n",
    "        self.d_k = d_k\n",
    "        self.d_v = d_v\n",
    "        self.h = h\n",
    "\n",
    "        self.init_weights(gain=1.0)\n",
    "\n",
    "    def init_weights(self, gain=1.0):\n",
    "        nn.init.xavier_normal_(self.fc_q.weight, gain=gain)\n",
    "        nn.init.xavier_normal_(self.fc_k.weight, gain=gain)\n",
    "        nn.init.xavier_normal_(self.fc_v.weight, gain=gain)\n",
    "        nn.init.xavier_normal_(self.fc_o.weight, gain=gain)\n",
    "        nn.init.constant_(self.fc_q.bias, 0)\n",
    "        nn.init.constant_(self.fc_k.bias, 0)\n",
    "        nn.init.constant_(self.fc_v.bias, 0)\n",
    "        nn.init.constant_(self.fc_o.bias, 0)\n",
    "\n",
    "    def forward(self, queries, keys, values):\n",
    "        \"\"\"\n",
    "        Computes\n",
    "        :param queries: Queries (b_s, nq, d_model)\n",
    "        :param keys: Keys (b_s, nk, d_model)\n",
    "        :param values: Values (b_s, nk, d_model)\n",
    "        :return:\n",
    "        \"\"\"\n",
    "        b_s, nq = queries.shape[:2]\n",
    "        nk = keys.shape[1]\n",
    "        q = self.fc_q(queries).view(b_s, nq, self.h, self.d_k).permute(0, 2, 1, 3)  # (b_s, h, nq, d_k)\n",
    "        k = self.fc_k(keys).view(b_s, nk, self.h, self.d_k).permute(0, 2, 3, 1)  # (b_s, h, d_k, nk)\n",
    "        v = self.fc_v(values).view(b_s, nk, self.h, self.d_v).permute(0, 2, 1, 3)  # (b_s, h, nk, d_v)\n",
    "\n",
    "        att = torch.matmul(q, k) / np.sqrt(self.d_k)  # (b_s, h, nq, nk)\n",
    "\n",
    "        att = torch.softmax(att, -1)\n",
    "\n",
    "        out = torch.matmul(att, v).permute(0, 2, 1, 3).contiguous().view(b_s, nq, self.h * self.d_v)  # (b_s, nq, h*d_v)\n",
    "        out = self.fc_o(out)  # (b_s, nq, d_model)\n",
    "        return out\n",
    "    \n",
    "class MultiHeadAttention(nn.Module):\n",
    "    def __init__(self, d_model, d_k, d_v, h, dff=2048, dropout=.1):\n",
    "        super(MultiHeadAttention, self).__init__()\n",
    "\n",
    "        self.attention = ScaledDotProductAttention(d_model=d_model, d_k=d_k, d_v=d_v, h=h)\n",
    "        self.dropout = nn.Dropout(p=dropout)\n",
    "        self.layer_norm = nn.LayerNorm(d_model)\n",
    "        self.fc = nn.Sequential(*[nn.Linear(d_model, dff), nn.ReLU(inplace=True), nn.Dropout(p=dropout),nn.Linear(dff, d_model)])\n",
    "\n",
    "    def forward(self, queries, keys, values):\n",
    "        att = self.attention(queries, keys, values)\n",
    "        att = self.dropout(att)\n",
    "        att = self.fc(att)\n",
    "        att = self.dropout(att)\n",
    "        return self.layer_norm(queries + att)\n",
    "    \n",
    "class EncoderSelfAttention(nn.Module):\n",
    "    def __init__(self, device, d_model, d_k, d_v, n_head, dff=2048, dropout_transformer=.1, n_module=6):\n",
    "        super(EncoderSelfAttention, self).__init__()\n",
    "        self.encoder = nn.ModuleList([MultiHeadAttention(d_model, d_k, d_v, n_head, dff, dropout_transformer) for _ in range(n_module)])\n",
    "        self.device = device\n",
    "    \n",
    "    def forward(self, x):        \n",
    "        in_encoder = x + sinusoid_encoding_table(x.shape[1], x.shape[2]).expand(x.shape).to(self.device)\n",
    "        for l in self.encoder:\n",
    "            in_encoder = l(in_encoder, in_encoder, in_encoder)\n",
    "            \n",
    "        return in_encoder"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Partie vision : Resnet + Transformer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "class _GestureTransformer(nn.Module):\n",
    "    \"\"\"Multi-Modal model on 3 or 1 channel\"\"\"\n",
    "    def __init__(self,device,backbone=\"resnet\",in_planes=3,pretrained= True,input_dim=512,layers_to_unfreeze=8,layers_to_delete=2,n_head=8,n_module=6,ff_size=1024,dropout1d=0.5):\n",
    "        super(_GestureTransformer, self).__init__()\n",
    "\n",
    "        self.in_planes = in_planes\n",
    "        self.device = device\n",
    "        self.conv_name = backbone\n",
    "        self.conv_model = None\n",
    "        \n",
    "        if self.conv_name.lower()==\"resnet\":\n",
    "            self.conv_model = Resnet(pretrained,layers_to_unfreeze,layers_to_delete,in_planes)\n",
    "        else:\n",
    "            raise NotImplementedError(\"The model {} is not supported!\".format(self.conv_name))\n",
    "            \n",
    "        self.conv_model.to(device)\n",
    "        self.features = features_extraction(self.conv_model,in_planes)\n",
    "        self.self_attention = EncoderSelfAttention(device,input_dim,64,64,n_head=n_head,dff=ff_size,dropout_transformer=dropout1d,n_module=n_module)\n",
    "\n",
    "    def forward(self, x):\n",
    "        \n",
    "        shape = x.shape        \n",
    "        x = self.features(x)\n",
    "        x = x.view(shape[0],shape[1],-1)\n",
    "        x = self.self_attention(x)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# CoordinatesTransformer (partie cinématique)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class CoordinatesTransformer(nn.Module):\n",
    "    \"\"\"Multi-Modal model on 3 or 1 channel\"\"\"\n",
    "    def __init__(self,device,pretrained= True,input_dim=512,layers_to_unfreeze=8,layers_to_delete=2,n_head=8,n_module=6,ff_size=1024,dropout1d=0.5):\n",
    "        super(CoordinatesTransformer, self).__init__()\n",
    "        \n",
    "        self.linear_mapper= torch.nn.Sequential(\n",
    "                              torch.nn.Linear(2, 32),\n",
    "                              torch.nn.ReLU(),\n",
    "                              torch.nn.Linear(32, 64),\n",
    "                              torch.nn.ReLU(),\n",
    "                              torch.nn.Linear(64, 128),\n",
    "                              torch.nn.ReLU(),\n",
    "                              torch.nn.Linear(128, 256),\n",
    "                              torch.nn.ReLU(),)\n",
    "        \n",
    "        \n",
    "        self.self_attention = EncoderSelfAttention(device,256,64,64,n_head=n_head,dff=ff_size,dropout_transformer=dropout1d,n_module=n_module)\n",
    "        self.pool = nn.AdaptiveAvgPool2d((1,256))\n",
    "        self.fc_out = nn.Sequential(nn.Linear(256,256),nn.ReLU(),nn.Dropout(p=dropout_val))\n",
    "        self.embedder_out = nn.Sequential(nn.Linear(8*256, 256),nn.ReLU(),nn.Dropout(p=dropout_val),nn.Linear(256, 256),nn.ReLU())\n",
    "\n",
    "    def forward(self, x):\n",
    "        shape = x.shape\n",
    "        x=self.linear_mapper(x)\n",
    "        x = x.view(shape[0],shape[1],-1)\n",
    "        x = self.self_attention(x)\n",
    "        x = self.fc_out(x)\n",
    "        return x\n",
    "    \n",
    "    def emb_out(self,input):\n",
    "        out= self.embedder_out(input)\n",
    "        return out"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Decoder Transformer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class DecoderLayer(nn.Module):\n",
    "    def __init__(self, d_model, nhead):\n",
    "        super(DecoderLayer, self).__init__()\n",
    "        self.self_attn = nn.MultiheadAttention(d_model, nhead)\n",
    "        self.norm1 = nn.LayerNorm(d_model)\n",
    "        self.norm2 = nn.LayerNorm(d_model)\n",
    "        self.feed_forward = nn.Linear(d_model, d_model)\n",
    "\n",
    "    def forward(self, x, e_output):\n",
    "        # Multihead self-attention\n",
    "        attn_output, _ = self.self_attn(x, x, x)\n",
    "        x = x + self.norm1(attn_output)\n",
    "        # Feed forward\n",
    "        ff_output = self.feed_forward(x)\n",
    "        x = x + self.norm2(ff_output)\n",
    "        return x\n",
    "\n",
    "    \n",
    "class DecoderTransformer(nn.Module):\n",
    "    def __init__(self, in_size, embed_size, hidden_size, d_model=512, dropout_val=dropout_val, batch_size=1, nhead=8, num_layers=6):\n",
    "        super(DecoderTransformer, self).__init__()\n",
    "                \n",
    "        self.in_size                = in_size #2\n",
    "        self.stochastic_out_size    = stochastic_out_size #2*2=4\n",
    "        self.hidden_size            = hidden_size #256\n",
    "        self.batch_size             = batch_size\n",
    "        self.embed_size             = embed_size #64\n",
    "        self.seq_length             = T_pred #12\n",
    "        self.dropout_val            = dropout_val #0.2\n",
    "        self.visual_embed_size      = visual_embed_size #64\n",
    "        self.visual_embed_size      = visual_embed_size\n",
    "        self.visual_size            = image_dimension * image_size * image_size #3*256*256\n",
    "        \n",
    "        self.d_model=d_model\n",
    "        self.nhead=nhead\n",
    "        self.num_layers=num_layers\n",
    "        \n",
    "        self.embedder_rho = nn.Linear(self.in_size, self.embed_size) #(2,64)\n",
    "        self.fC_mu = nn.Sequential(nn.Linear(self.hidden_size + self.hidden_size + in_size, int(self.hidden_size/2), bias=True),nn.ReLU(),nn.Dropout(p=dropout_val),nn.Linear(int(self.hidden_size/2), self.stochastic_out_size, bias=True))\n",
    "        self.dropout = nn.Dropout(dropout_val)\n",
    "        self.reducted_size = int((self.hidden_size-1)/3)+1\n",
    "        self.reducted_size2 = int((self.hidden_size+in_size-1)/3)+1\n",
    "        self.FC_dim_red = nn.Sequential(nn.MaxPool2d(kernel_size=3, stride=3, padding=1),nn.Flatten(start_dim=1, end_dim=-1),nn.Linear(self.reducted_size*self.reducted_size2, 2*self.hidden_size+in_size, bias=True),nn.ReLU())\n",
    "        \n",
    "        self.embedding = nn.Linear(64, d_model)\n",
    "        self.layers = nn.ModuleList([DecoderLayer(d_model, nhead) for _ in range(num_layers)])\n",
    "        self.output_layer = nn.Linear(d_model, 514)\n",
    "                \n",
    "    def forward(self, x, encoder_outputs):\n",
    "        \n",
    "        # Coordination Embedding\n",
    "        embedding = self.embedder_rho(x.view(x.shape[0],-1,2))\n",
    "        embedding = F.relu(self.dropout(embedding))\n",
    "        \n",
    "        # Embed the decoder input\n",
    "        x = self.embedding(embedding)\n",
    "\n",
    "        for layer in self.layers:\n",
    "            x = layer(x, encoder_outputs)\n",
    "        output = self.output_layer(x)\n",
    "        \n",
    "        prediction = self.fC_mu(output.squeeze(0)) \n",
    "        \n",
    "        return prediction\n",
    "    \n",
    "    def dim_red(self, input):\n",
    "        output = self.FC_dim_red(input)\n",
    "        return output"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Seq2Seq(nn.Module):\n",
    "    def __init__(self, in_size, embed_size, hidden_size, batch_size=1, d_model=512, d_ff=2048, h=8, dropout_val=dropout_val, N=6, input_dim=512):\n",
    "        super(Seq2Seq, self).__init__()\n",
    "\n",
    "        torch.cuda.empty_cache()\n",
    "        \n",
    "        self.encoder = CoordinatesTransformer(device,dropout1d=dropout_val) #Encoder Transformer (partie cinématique)\n",
    "        self.encoder.apply(init_weights)\n",
    "                \n",
    "        self.decoder =  DecoderTransformer(in_size, embed_size, hidden_size, num_layers=6, nhead=8)\n",
    "        self.decoder.apply(init_weights)\n",
    "        \n",
    "        self.vsn_module = _GestureTransformer(device,dropout1d=dropout_val) #_GestureTransformer(partie vision)                   \n",
    "        self.vsn_module.apply(init_weights)\n",
    "        \n",
    "        self.pooling = nn.AdaptiveAvgPool1d((hidden_size))\n",
    "        \n",
    "        self.crossAttention = crossAttention(N=6,d_model=256, d_ff=2048, h=8, dropout=0.1)\n",
    "        \n",
    "        if device.type=='cuda':\n",
    "            self.encoder.cuda()\n",
    "            self.decoder.cuda()\n",
    "            self.vsn_module.cuda()\n",
    "\n",
    "    def forward(self, input_tensor, visual_input_tensor, output_tensor, batch_size, train_mode): \n",
    "        \n",
    "        batch_size      = int(input_tensor.size(0))        \n",
    "        \n",
    "        #encoder_outputs \n",
    "        encoder_outputs = torch.zeros(batch_size, T_obs, hidden_size).cuda()\n",
    "\n",
    "        start_point     = (input_tensor[:,0,:]).to(device).clone().detach()\n",
    "\n",
    "        if startpoint_mode==\"on\":\n",
    "            input_tensor[:,0,:]    = 0\n",
    "            \n",
    "        encoder_outputs                  = self.encoder(input_tensor.reshape(batch_size, -1, 2))\n",
    "            \n",
    "        # Encoder outputs :                     \n",
    "        visual_initial_vsn    = self.vsn_module(visual_input_tensor)\n",
    "        visual_initial_vsn    = self.pooling(visual_initial_vsn) #pooling qu'on a ajouté\n",
    "        \n",
    "        src_mask = None\n",
    "        obd_enc_mask = None\n",
    "        \n",
    "        cross_ouput = self.crossAttention(encoder_outputs, visual_initial_vsn, src_mask, obd_enc_mask)\n",
    "        e_outputss=cross_ouput\n",
    "        \n",
    "        visual_vsn_result   = visual_initial_vsn\n",
    "        \n",
    "        decoder_input = input_tensor[:,-1,:]\n",
    "        \n",
    "        # Tensor to store decoder outputs\n",
    "        outputs                         = torch.zeros(batch_size, T_pred , in_size).cuda() #torch.Size([32, 12, 2])\n",
    "        stochastic_outputs              = torch.zeros(batch_size, T_pred , stochastic_out_size).cuda() #torch.Size([32, 12, 4])\n",
    "        teacher_force                   = 1\n",
    "\n",
    "        epsilonX                        = Normal(torch.zeros(batch_size,1),torch.ones(batch_size,1))\n",
    "        epsilonY                        = Normal(torch.zeros(batch_size,1),torch.ones(batch_size,1))\n",
    "        teacher_force                   = int(random.random() < teacher_forcing_ratio) if train_mode else 0\n",
    "        \n",
    "        for t in range(0, T_pred):\n",
    "\n",
    "            stochastic_decoder_output = self.decoder(decoder_input, e_outputss)\n",
    "            \n",
    "            # Reparameterization Trick :)\n",
    "            decoder_output              = torch.zeros(batch_size,1,2).cuda()            \n",
    "\n",
    "            if stochastic_mode and path_mode=='single':\n",
    "                decoder_output[:,:,0]  = stochastic_decoder_output[:,:,0] + epsilonX.sample().cuda() * stochastic_decoder_output[:,:,1]\n",
    "                decoder_output[:,:,1]  = stochastic_decoder_output[:,:,2] + epsilonY.sample().cuda() * stochastic_decoder_output[:,:,3]\n",
    "            elif stochastic_mode and path_mode=='avg':\n",
    "                decoder_output[:,:,0]  = stochastic_decoder_output[:,:,0] + epsilonX.sample((avg_n_path_eval,1)).view(-1,avg_n_path_eval,1).mean(-2).cuda() * stochastic_decoder_output[:,:,1]\n",
    "                decoder_output[:,:,1]  = stochastic_decoder_output[:,:,2] + epsilonY.sample((avg_n_path_eval,1)).view(-1,avg_n_path_eval,1).mean(-2).cuda() * stochastic_decoder_output[:,:,3]\n",
    "            elif not(stochastic_mode):\n",
    "                decoder_output[:,:,0]  = stochastic_decoder_output[:,-1,0].clone().view(batch_size,-1) \n",
    "                decoder_output[:,:,1]  = stochastic_decoder_output[:,-1,2].clone().view(batch_size,-1) \n",
    "            elif stochastic_mode and path_mode == \"bst\":\n",
    "                epsilon_x               = torch.randn([batch_size,bst_n_path_eval,1], dtype=torch.float).cuda()\n",
    "                epsilon_y               = torch.randn([batch_size,bst_n_path_eval,1], dtype=torch.float).cuda()\n",
    "                multi_path_x            = stochastic_decoder_output[:,-1,0].clone().view(batch_size,1,-1) + epsilon_x * stochastic_decoder_output[:,-1,1].clone().view(batch_size,1,-1)\n",
    "                multi_path_y            = stochastic_decoder_output[:,-1,2].clone().view(batch_size,1,-1) + epsilon_y * stochastic_decoder_output[:,-1,3].clone().view(batch_size,1,-1)\n",
    "                ground_truth_x          = output_tensor[:,t,0].view(batch_size,1,1).cuda()\n",
    "                ground_truth_y          = output_tensor[:,t,1].view(batch_size,1,1).cuda()\n",
    "                diff_path_x             = multi_path_x - ground_truth_x\n",
    "                diff_path_y             = multi_path_y - ground_truth_y\n",
    "                diff_path               = (torch.sqrt( diff_path_x.pow(2) + diff_path_y.pow(2) )).sum(dim=-1)\n",
    "                idx                     = torch.arange(batch_size,dtype=torch.long).cuda()\n",
    "                min                     = torch.argmin(diff_path,dim=1).squeeze()\n",
    "                decoder_output[:,:,0]   = multi_path_x[idx,min,:].view(batch_size,1)\n",
    "                decoder_output[:,:,1]   = multi_path_y[idx,min,:].view(batch_size,1)\n",
    "            elif stochastic_mode and path_mode == \"top5\":\n",
    "                k = 5 #top k                \n",
    "                epsilon_x               = torch.randn([batch_size,bst_n_path_eval,1], dtype=torch.float).cuda()\n",
    "                epsilon_y               = torch.randn([batch_size,bst_n_path_eval,1], dtype=torch.float).cuda()\n",
    "                multi_path_x            = stochastic_decoder_output[:,-1,0].clone().view(batch_size,1,-1) + epsilon_x * stochastic_decoder_output[:,-1,1].clone().view(batch_size,1,-1)\n",
    "                multi_path_y            = stochastic_decoder_output[:,-1,2].clone().view(batch_size,1,-1) + epsilon_y * stochastic_decoder_output[:,-1,3].clone().view(batch_size,1,-1)\n",
    "                ground_truth_x          = output_tensor[:,t,0].view(batch_size,1,1).cuda()\n",
    "                ground_truth_y          = output_tensor[:,t,1].view(batch_size,1,1).cuda()\n",
    "                diff_path_x             = multi_path_x - ground_truth_x\n",
    "                diff_path_y             = multi_path_y - ground_truth_y\n",
    "                diff_path               = (torch.sqrt( diff_path_x.pow(2) + diff_path_y.pow(2) )).sum(dim=-1)\n",
    "                idx                     = torch.arange(batch_size,dtype=torch.long).repeat(k).view(k,-1).transpose(0,1).cuda()\n",
    "                min_val, min            = torch.topk(diff_path, k=k, dim=1,largest=False)\n",
    "                decoder_output[:,:,0]   = multi_path_x[idx,min,:].mean(dim=-2).view(batch_size,1)\n",
    "                decoder_output[:,:,1]   = multi_path_y[idx,min,:].mean(dim=-2).view(batch_size,1)\n",
    "                \n",
    "            outputs[:,t,:]                        = decoder_output.clone().squeeze() #+ decoder_input.squeeze()\n",
    "            stochastic_outputs[:,t,:]             = stochastic_decoder_output[:,-1,:].clone().squeeze()\n",
    "            decoder_input                         = output_tensor[:,:t+1,:].clone() if teacher_force else outputs[:,:t+1,:].clone()\n",
    "\n",
    "        return outputs, stochastic_outputs, visual_vsn_result"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Train\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import shutil\n",
    "\n",
    "def save_checkpoint(state, is_best, save_path, filename):\n",
    "    torch.save(state, os.path.join(save_path,filename))\n",
    "    if is_best:\n",
    "        shutil.copyfile(os.path.join(save_path,filename), os.path.join(save_path,'model_best.pth'))\n",
    "        \n",
    "def train(model, optimizer, scheduler, criterion, criterion_vision, clip,train_loader, validation_loader):\n",
    "    global batch_size\n",
    "    i               = None\n",
    "    checked_frame   = 0\n",
    "\n",
    "    losses = []\n",
    "    print(\"Data Size \",data_size,\"\\tChunk Size \",chunk_size)\n",
    "    global teacher_forcing_ratio\n",
    "    counter =0\n",
    "    best_val = float(\"inf\")\n",
    "    save_path = './save_modelseth1000ep'\n",
    "    for j in range(epoch_num):\n",
    "        model.train()\n",
    "        epoch_loss=0\n",
    "        if j%7 == 6:\n",
    "            teacher_forcing_ratio = (teacher_forcing_ratio - 0.2) if teacher_forcing_ratio>=0.1 else 0.0\n",
    "\n",
    "        # Update TeachForce ratio to gradually change during training\n",
    "        # global teacher_forcing_ratio\n",
    "        # teacher_forcing_ratio-= 1/epoch_num\n",
    "        print(\"TEACHER FORCE RATIO\\t\",teacher_forcing_ratio)\n",
    "\n",
    "        start_time = time.time()\n",
    "        if(j>=from_epoch):\n",
    "            optimizer.zero_grad()\n",
    "            \n",
    "            ADE = 0\n",
    "            FDE = 0\n",
    "            i   = 0\n",
    "            for i,data in enumerate(train_loader):\n",
    "                #print(\"\\n--------------- Batch %d/ %d ---------------\"%(j,i)) #(epoch/i)\n",
    "                # Forward\n",
    "                obs, pred, visual_obs, frame_tensor              = data\n",
    "                input_tensor, output_tensor                      = obs.float().squeeze().to('cuda', non_blocking=True), pred.float().squeeze().to('cuda', non_blocking=True)               #(obs.to(device), pred.to(device))\n",
    "                visual_input_tensor                              = visual_obs.squeeze().to('cuda', non_blocking=True)  #(visual_obs.to(device), visual_pred.to(device))\n",
    "                prediction, stochastic_prediction, visual_embedding = model(input_tensor,visual_input_tensor,output_tensor,batch_size,train_mode=1)\n",
    "\n",
    "                calculated_prediction = prediction.cumsum(axis=1)\n",
    "\n",
    "                loss_line_regularizer = distance_from_line_regularizer(input_tensor,calculated_prediction) #loss (regularisation term Lreg)\n",
    "                \n",
    "                if biased_loss_mode:\n",
    "                    weight  = torch.arange(1,2*T_pred+1,2).cuda().float()\n",
    "                    weight  = torch.exp(weight / T_pred).repeat(prediction.size(0)).view(prediction.size(0),T_pred,1)\n",
    "                    loss    = criterion( (calculated_prediction)*weight, torch.cumsum(output_tensor,dim=-2)*weight)\n",
    "                    #loss_vision = criterion_vision()\n",
    "                else:\n",
    "                    loss    = criterion( (calculated_prediction), torch.cumsum(output_tensor,dim=-2)) #mean squared error (lmse)\n",
    "                out_x       = output_tensor[:,:,0].cumsum(axis=1)\n",
    "                out_y       = output_tensor[:,:,1].cumsum(axis=1)\n",
    "                pred_x      = calculated_prediction[:,:,0]\n",
    "                pred_y      = calculated_prediction[:,:,1]\n",
    "                ADE         += ((out_x.sub(pred_x)**2).add((out_y.sub(pred_y)**2))**(1/2)).mean(0).mean(0)   \n",
    "                \n",
    "                # Backward Propagation\n",
    "                total_loss      = loss.double() + torch.tensor(regularization_factor).to('cuda', non_blocking=True) * loss_line_regularizer.double() #total loss\n",
    "                total_loss.backward()\n",
    "                torch.nn.utils.clip_grad_norm_(model.parameters(), clip)\n",
    "                optimizer.step()\n",
    "                #print(\"Total Loss\\t{:.2f}\".format(total_loss.item()))\n",
    "                epoch_loss += total_loss.item()\n",
    "                #print(\"Time\\t\\t{:.2f} sec \\n\".format(time.time() - start_time))\n",
    "                start_time = time.time()\n",
    "                torch.cuda.empty_cache()\n",
    "                writer.close()\n",
    "                count_div=i\n",
    "            \n",
    "            # tensorboard log\n",
    "            writer.add_scalar('ADE/train', ADE.item()/(count_div+1), counter)\n",
    "            # writer.add_scalar('FDE/train', FDE.item()/(count_div+1), counter)\n",
    "            # writer.add_scalar('LOSS/train', epoch_loss/(count_div+1), counter)\n",
    "            counter += 1\n",
    "\n",
    "        if scheduler.get_last_lr()[0]>0.001:\n",
    "            scheduler.step()\n",
    "        # validation(model, optimizer, criterion, criterion_vision, clip, validation_loader, j) \n",
    "        epoch_loss = epoch_loss / (int(data_size/chunk_size))\n",
    "        losses.append(epoch_loss)\n",
    "        display.clear_output(wait=True)\n",
    "        plt.plot(losses, '--ro', label='train loss')\n",
    "        plt.legend()\n",
    "        plt.title(f'epoch {j}')\n",
    "        plt.show()\n",
    "        print(\"Time\\t\\t{:.2f} sec \\n\".format(time.time() - start_time))\n",
    "        print(\"EPOCH \", j, \"\\tLOSS \", epoch_loss)\n",
    "        writer.add_scalar('epoch_loss/train', epoch_loss/ (int(data_size/chunk_size)), j ) #see how model performs on the training dataset\n",
    "        print(\"-----------------------------------------------\\n\"+\"-----------------------------------------------\")\n",
    "\n",
    "        # save checkpoint for each epoch and a fine called best_model so far \n",
    "        print(np.argmin(losses))\n",
    "        is_best = epoch_loss < best_val\n",
    "        best_val = min(epoch_loss, best_val)\n",
    "        save_checkpoint({'epoch': j+1,'state_dict': model.module.state_dict(),'optimizer': optimizer.state_dict(),'scheduler': scheduler.state_dict(),'best_loss': best_val}, is_best, save_path, 'epoch_{}.pth'.format(j+1))\n",
    "        \n",
    "    return epoch_loss / (int(data_size/chunk_size))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Evaluate\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def validation(model, optimizer, criterion, criterion_vision, clip, validation_loader, counter):\n",
    "    global batch_size\n",
    "    model.eval()\n",
    "    i           = None\n",
    "    ADEs        = 0\n",
    "    FDEs        = 0\n",
    "    epoch_loss  = 0\n",
    "    loss_line_regularizer = 0\n",
    "    loss = 0 \n",
    "    total_loss = 0\n",
    "    ADE  = 0\n",
    "    FDE  = 0\n",
    "    for i,data in enumerate(test_loader):\n",
    "        # Forward\n",
    "        obs, pred, visual_obs, frame_tensor = data\n",
    "        input_tensor, output_tensor         = obs.float().squeeze().to('cuda', non_blocking=True), pred.float().squeeze().to('cuda', non_blocking=True)\n",
    "        visual_input_tensor                 = visual_obs.squeeze().to('cuda', non_blocking=True) \n",
    "        prediction, stochastic_prediction, visual_embedding = model(input_tensor, visual_input_tensor, output_tensor, batch_size, train_mode=0)\n",
    "        \n",
    "        calculated_prediction = prediction.cumsum(axis=1) \n",
    "\n",
    "        loss_line_regularizer = distance_from_line_regularizer(input_tensor,calculated_prediction)\n",
    "        \n",
    "        if biased_loss_mode:\n",
    "            weight  = torch.arange(1,2*T_pred+1,2).cuda().float()\n",
    "            weight  = torch.exp(weight / T_pred).repeat(prediction.size(0)).view(prediction.size(0),T_pred,1)\n",
    "            loss    = criterion( (calculated_prediction)*weight, torch.cumsum(output_tensor,dim=-2)*weight)\n",
    "        else:\n",
    "            loss    = criterion( (calculated_prediction), torch.cumsum(output_tensor,dim=-2))\n",
    "        out_x       = output_tensor[:,:,0].cumsum(axis=1)\n",
    "        out_y       = output_tensor[:,:,1].cumsum(axis=1)\n",
    "        pred_x      = calculated_prediction[:,:,0]\n",
    "        pred_y      = calculated_prediction[:,:,1]\n",
    "        ADE         += ((out_x.sub(pred_x)**2).add((out_y.sub(pred_y)**2))**(1/2)).mean(0).mean(0)   \n",
    "        FDE         += ((out_x.sub(pred_x)**2).add((out_y.sub(pred_y)**2))**(1/2)).mean(0)[-1]\n",
    "        total_loss  += loss.double() + regularization_factor * loss_line_regularizer.double() \n",
    "        print(\"Total Loss\\t{:.2f}\".format(total_loss.item()))\n",
    "\n",
    "    writer.add_scalar('ADE/val_'+path_mode,             ADE.item()/(i+1),             counter)\n",
    "    writer.add_scalar('FDE/val_'+path_mode,             FDE.item()/(i+1),             counter)\n",
    "    writer.add_scalar('LOSS/val_'+path_mode,            total_loss.item()/(i+1)   ,   counter)\n",
    "    writer.add_scalar('LOSS_c/val_'+path_mode,          loss.item()/(i+1)        ,    counter)\n",
    "    writer.add_scalar('L-REGULARIZER/val_'+path_mode,   loss_line_regularizer.item()/(i+1), counter)\n",
    "    writer.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Evaluate\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def evaluate_eval(model, optimizer, criterion, criterion_vision, clip, five_fold_cross_validation):\n",
    "    global batch_size\n",
    "    model.eval()\n",
    "    i           = None\n",
    "    ADEs        = 0\n",
    "    FDEs        = 0\n",
    "    epoch_loss  = 0\n",
    "    list_x_obs          = ['x_obs_'+str(i)              for i in range(0,T_obs)] #x_obs_0 --> x_obs_7\n",
    "    list_y_obs          = ['y_obs_'+str(i)              for i in range(0,T_obs)] #y_obs_0 --> y_obs_7\n",
    "    list_x_pred         = ['x_pred_'+str(i)             for i in range(0,T_pred)] #x_pred_0 --> x_pred_11\n",
    "    list_y_pred         = ['y_pred_'+str(i)             for i in range(0,T_pred)] #y_pred_0 --> y_pred_11\n",
    "    list_x_stoch_pred_m = ['x_stoch_pred_m_'+str(i)     for i in range(0,T_pred)] #x_stoch_pred_m_0 --> x_stoch_pred_m_11\n",
    "    list_y_stoch_pred_m = ['y_stoch_pred_m_'+str(i)     for i in range(0,T_pred)] #y_stoch_pred_m_0 --> y_stoch_pred_m_11\n",
    "    list_x_stoch_pred_s = ['x_stoch_pred_s_'+str(i)     for i in range(0,T_pred)] #x_stoch_pred_s_0 --> x_stoch_pred_s_11\n",
    "    list_y_stoch_pred_s = ['y_stoch_pred_s_'+str(i)     for i in range(0,T_pred)] #y_stoch_pred_s_0 --> y_stoch_pred_s_11\n",
    "    list_x_out          = ['x_out_'+str(i)              for i in range(0,T_pred)] #x_out_0 --> x_out_11\n",
    "    list_y_out          = ['y_out_'+str(i)              for i in range(0,T_pred)] #y_out_0 --> y_out_11\n",
    "    list_vsn            = ['vsn_'+str(i)               for i in range(0,hidden_size)] #vsn_0 --> vsn_255\n",
    "    df_out              = pd.DataFrame(columns=list_x_obs + list_y_obs + list_x_out + list_y_out + list_x_pred + list_y_pred + list_x_stoch_pred_m + list_y_stoch_pred_m + list_x_stoch_pred_s + list_y_stoch_pred_s + list_vsn)# + list_vsn_visual + list_c_context + list_h_context)\n",
    "\n",
    "    for i,data in enumerate(test_loader):\n",
    "\n",
    "        start_time = time.time()\n",
    "        \n",
    "        # Forward\n",
    "        obs, pred, visual_obs, frame_tensor                 = data\n",
    "        input_tensor, output_tensor                         = obs.float().squeeze().to('cuda', non_blocking=True), pred.float().squeeze().to('cuda', non_blocking=True)               #(obs.to(device), pred.to(device))\n",
    "        visual_input_tensor                                 = visual_obs.squeeze().cuda()\n",
    "        if len(input_tensor.size()) == 2:\n",
    "            break\n",
    "        print(input_tensor.size())\n",
    "        prediction, stochastic_prediction, visual_embedding = model(input_tensor,visual_input_tensor,output_tensor,batch_size,train_mode=0)\n",
    "        \n",
    "        calculated_prediction =  prediction.cumsum(axis=1) \n",
    "\n",
    "        loss_line_regularizer = distance_from_line_regularizer(input_tensor,calculated_prediction) #lreg\n",
    "\n",
    "        if biased_loss_mode:\n",
    "            weight  = torch.arange(1,2*T_pred+1,2).cuda().float()\n",
    "            weight  = torch.exp(weight / T_pred).repeat(prediction.size(0)).view(prediction.size(0),T_pred,1)\n",
    "            loss    = criterion( (calculated_prediction)*weight, torch.cumsum(output_tensor,dim=-2)*weight)\n",
    "        else:\n",
    "            loss    = criterion( (calculated_prediction), torch.cumsum(output_tensor,dim=-2)) #lmse\n",
    "        out_x           = output_tensor[:,:,0].cumsum(axis=1)\n",
    "        out_y           = output_tensor[:,:,1].cumsum(axis=1)\n",
    "        pred_x          = calculated_prediction[:,:,0]\n",
    "        pred_y          = calculated_prediction[:,:,1]\n",
    "        ADE             = ((out_x.sub(pred_x)**2).add((out_y.sub(pred_y)**2))**(1/2)).mean(0).mean(0)   \n",
    "        FDE             = ((out_x.sub(pred_x)**2).add((out_y.sub(pred_y)**2))**(1/2)).mean(0)[-1]\n",
    "        total_loss      = loss.double() + regularization_factor * loss_line_regularizer.double() #loss\n",
    "        print(\"Total Loss\\t{:.2f}\".format(total_loss.item()))\n",
    "        epoch_loss += total_loss.item()\n",
    "        ADEs    += ADE.item()\n",
    "        FDEs    += FDE.item()\n",
    "        input_x_lin                 = input_tensor[:,:,0].view(-1, T_obs).cpu() #x_obs\n",
    "        input_y_lin                 = input_tensor[:,:,1].view(-1, T_obs).cpu() #y_obs\n",
    "        output_x_lin                = output_tensor[:,:,0].view(-1, T_pred).cpu() #x_out\n",
    "        output_y_lin                = output_tensor[:,:,1].view(-1, T_pred).cpu() #y_out\n",
    "        prediction_x_lin            = prediction[:,:,0].view(-1, T_pred).cpu() #x_pred\n",
    "        prediction_y_lin            = prediction[:,:,1].view(-1, T_pred).cpu() #y_pred\n",
    "        stoch_prediction_x_m        = stochastic_prediction[:,:,0].view(-1, T_pred).cpu() #x_stoch_pred_m\n",
    "        stoch_prediction_x_s        = stochastic_prediction[:,:,1].view(-1, T_pred).cpu() #x_stoch_pred_s\n",
    "        stoch_prediction_y_m        = stochastic_prediction[:,:,2].view(-1, T_pred).cpu() #y_stoch_pred_m\n",
    "        stoch_prediction_y_s        = stochastic_prediction[:,:,3].view(-1, T_pred).cpu() #y_stoch_pred_s\n",
    "        whole_data                  = torch.cat((input_x_lin, input_y_lin, output_x_lin, output_y_lin, prediction_x_lin, prediction_y_lin, stoch_prediction_x_m, stoch_prediction_y_m, stoch_prediction_x_s, stoch_prediction_y_s),1) #visual_embedding_weights), 1) #,context_c_lin, context_h_lin\n",
    "        temp                        = pd.DataFrame(whole_data.detach().cpu().numpy(), columns=list_x_obs + list_y_obs + list_x_out + list_y_out + list_x_pred + list_y_pred + list_x_stoch_pred_m + list_y_stoch_pred_m + list_x_stoch_pred_s + list_y_stoch_pred_s ) #list_vsn+ list_c_context + list_h_context\n",
    "        df_out                      = df_out.append(temp)\n",
    "        df_out.reset_index(drop=True,inplace=True)\n",
    "\n",
    "        print(\"Time\\t\\t{:.2f} sec \\n\".format(time.time() - start_time))\n",
    "\n",
    "    # ADE/FDE Report\n",
    "    out_x  = df_out[['x_out_' +str(i) for i in range(0,T_pred)]].cumsum(axis=1)\n",
    "    pred_x = df_out[['x_pred_'+str(i) for i in range(0,T_pred)]].cumsum(axis=1)\n",
    "    out_y  = df_out[['y_out_' +str(i) for i in range(0,T_pred)]].cumsum(axis=1)\n",
    "    pred_y = df_out[['y_pred_'+str(i) for i in range(0,T_pred)]].cumsum(axis=1)    \n",
    "    ADE = (out_x.sub(pred_x.values)**2).add((out_y.sub(pred_y.values)**2).values, axis=1)**(1/2)\n",
    "    df_out['ADE'] = ADE.mean(axis=1)\n",
    "    FDE = ADE.x_out_11\n",
    "    df_out['FDE'] = FDE\n",
    "    Mean_ADE = df_out.ADE.mean()\n",
    "    Mean_FDE = df_out.FDE.mean()\n",
    "    print(\"MEAN ADE/FDE\\t\",Mean_ADE,Mean_FDE)\n",
    "    writer.add_scalar(\"Final_Test/ADE_\"+path_mode, Mean_ADE, global_step=0)\n",
    "    writer.add_scalar(\"Final_Test/FDE_\"+path_mode, Mean_FDE, global_step=0)\n",
    "    df_out.to_sql(table_out+'_'+path_mode, cnx2, if_exists=\"replace\", index=False)\n",
    "    writer.close()\n",
    "    \n",
    "    return ADEs, FDEs, int(data_size/chunk_size)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# MAIN\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "model = Seq2Seq(in_size, embed_size, hidden_size, dropout_val=dropout_val, batch_size=batch_size)\n",
    "model = nn.DataParallel( model ).cuda()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"A summary of the model, where we can see the shape of each layer : \")\n",
    "print(model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "learning_step               = 40\n",
    "initial_learning_rate       = 0.01\n",
    "clip                        = 1\n",
    "\n",
    "# MSE loss\n",
    "criterion                   = nn.MSELoss(reduction='mean')#nn.NLLLoss()\n",
    "criterion_vision            = nn.MSELoss(reduction='sum')#nn.NLLLoss()\n",
    "\n",
    "# SGD optimizer\n",
    "optimizer                   = optim.SGD(model.parameters(), lr=initial_learning_rate, momentum=0.9, weight_decay=0.01) #SGD\n",
    "scheduler                   = torch.optim.lr_scheduler.StepLR(optimizer, step_size=learning_step, gamma=0.1)\n",
    "five_fold_cross_validation  = 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#train dataset and loader\n",
    "print(\"Initializing train dataset\")\n",
    "dataset_train = TrajectoryPredictionDataset(image_folder_path, DB_PATH_train, cnx_train)\n",
    "train_loader  = torch.utils.data.DataLoader(dataset_train, batch_size=batch_size, shuffle=True, num_workers=2, drop_last=True, pin_memory=True)\n",
    "validation_loader = None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"TRAIN\")\n",
    "model.train()\n",
    "print(\"path mode\\t\",path_mode)\n",
    "loss  = train(model, optimizer, scheduler, criterion, criterion_vision, clip, train_loader, validation_loader)\n",
    "print(\"LOSS \",loss)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_checkpoint(model, optimizer, scheduler, filename='checkpoint.pth.tar'):\n",
    "    start_epoch = 0\n",
    "    best_val=-1\n",
    "    if os.path.isfile(filename):\n",
    "        print(\"=> loading checkpoint '{}'\".format(filename))\n",
    "        checkpoint = torch.load(filename)\n",
    "        start_epoch = checkpoint['epoch']\n",
    "        model.load_state_dict(checkpoint['state_dict'])\n",
    "        optimizer.load_state_dict(checkpoint['optimizer'])\n",
    "        scheduler.load_state_dict(checkpoint['scheduler'])\n",
    "        try:\n",
    "            best_val=checkpoint['best_loss']\n",
    "        except:\n",
    "            best_val=-1\n",
    "        print(\"=> loaded checkpoint '{}' (epoch {})\".format(filename, checkpoint['epoch']))\n",
    "    else:\n",
    "        print(\"=> no checkpoint found at '{}'\".format(filename))\n",
    "\n",
    "    return model, optimizer, scheduler, start_epoch, best_val"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"LOAD MODEL\")\n",
    "# Change device to cpu\n",
    "#del model\n",
    "if torch.cuda.is_available():\n",
    "    torch.cuda.empty_cache()\n",
    "model      = Seq2Seq(in_size, embed_size, hidden_size, dropout_val=dropout_val, batch_size=batch_size)\n",
    "model, optimizer, scheduler, start_epoch, best_val = load_checkpoint(model, optimizer, scheduler, filename='./save_modelseth1000ep/model_best.pth')\n",
    "model.cuda()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#test dataset and loader\n",
    "print(\"Initializing val dataset\")\n",
    "dataset_val   = TrajectoryPredictionDataset(image_folder_path, DB_PATH_val, cnx_val)\n",
    "test_loader   = torch.utils.data.DataLoader(dataset_val, batch_size=batch_size, shuffle=False, num_workers=1, drop_last=False, pin_memory=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"EVALUATE bst\")\n",
    "\n",
    "model.eval()\n",
    "path_mode = 'bst'\n",
    "print(\"path mode\\t\",path_mode)\n",
    "evaluate_eval(model, optimizer, criterion, criterion_vision, clip, test_loader)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"EVALUATE top5\")\n",
    "model.eval()\n",
    "path_mode = 'top5'\n",
    "print(\"path mode\\t\",path_mode)\n",
    "evaluate_eval(model, optimizer, criterion, criterion_vision, clip, test_loader)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.10"
  },
  "vscode": {
   "interpreter": {
    "hash": "400dc80c028f38e78d0af9f2eedd78117085b98eeacd8fc6d80a8baae662b2d6"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}