TF4AML_practice.ipynb

{"cells":[{"cell_type":"markdown","metadata":{},"source":["# Using the Transformer Networks\n","\n","This notebook will guide you through the usage of a provided efficient implementation of Transformer Networks, to experiment with hyper-parameters and to perform ablation studies. This notebook will let you master accomplishing experiments with Transformer Networks and analyising the outcomes. Complete the code snippets where request and provide your observations. Feel free to refer the paper [Under the Hood of Transformer Networks for Trajectory Forecasting](https://arxiv.org/abs/2203.11878).\n","\n"," "]},{"cell_type":"markdown","metadata":{},"source":["# Initial setup"]},{"cell_type":"markdown","metadata":{},"source":["Only if you run from Google Colab run those 2 cells to sync with Google Drive."]},{"cell_type":"code","execution_count":null,"metadata":{},"outputs":[],"source":["from google.colab import drive\n","drive.mount('/content/drive', force_remount=True)"]},{"cell_type":"code","execution_count":null,"metadata":{"colab":{"base_uri":"https://localhost:8080/"},"executionInfo":{"elapsed":874,"status":"ok","timestamp":1665047302197,"user":{"displayName":"Alessandro Flaborea","userId":"12166823913657129807"},"user_tz":-120},"id":"-4LQv8CRbwnF","outputId":"eceb5573-7125-44e6-8626-886d49831ee5"},"outputs":[],"source":["%cd /content/drive/MyDrive/TF4AML/"]},{"cell_type":"markdown","metadata":{},"source":["Start with the import"]},{"cell_type":"code","execution_count":2,"metadata":{"id":"EDQ0I2ssSjDS"},"outputs":[],"source":["import torch\n","import torch.utils.data\n","import torch.nn as nn\n","import torch.nn.functional as F\n","import numpy as np\n","import scipy\n","import os\n","import time\n","\n","from transformer import baselineUtils\n","from transformer import individual_TF\n","from transformer.batch import subsequent_mask\n","from transformer.noam_opt import NoamOpt"]},{"cell_type":"code","execution_count":3,"metadata":{"colab":{"base_uri":"https://localhost:8080/"},"executionInfo":{"elapsed":471,"status":"ok","timestamp":1665047313035,"user":{"displayName":"Alessandro Flaborea","userId":"12166823913657129807"},"user_tz":-120},"id":"BegflNekUxqU","outputId":"ebe05426-4fea-441d-d700-dc4c06ed2f45"},"outputs":[{"name":"stdout","output_type":"stream","text":["Current device: cuda - Type: NVIDIA GeForce RTX 3090\n"]}],"source":["# Select GPU device for the training if available\n","if not torch.cuda.is_available():\n","    device=torch.device(\"cpu\")\n","    print(\"Current device:\", device)\n","else:\n","    device=torch.device(\"cuda\")\n","    print(\"Current device:\", device, \"- Type:\", torch.cuda.get_device_name(0))"]},{"cell_type":"markdown","metadata":{"id":"ZDYWfp9_Kp_1"},"source":["# Training and Testing"]},{"cell_type":"markdown","metadata":{"id":"17R3CEvylyj5"},"source":["## Data Loading (setup the dataset for train, validation and test)"]},{"cell_type":"markdown","metadata":{},"source":["The subdatasets are 5 (ETH, Hotel, Univ, Zara1 and Zara2) we will leave one of them out for testing and train on the other 4. \n","\n","I.e. choosing ```dataset_name = 'zara1'``` the training set is composed by ETH, Hotel, Univ and Zara2 and tested on Zara1.\n","\n","Moreover you can train and validate on a portion of the dataset setting percentage of the data (default is 50).\n","\n","------ \n","\n","Each sequence is composed by an observed part to train the Encoder and a part we are attempting to predict with the Decoder. \n","\n","Generally the standard setup plans to use the first 8 points for the observation and the following 12 for the prediction.\n","\n","------ \n","\n","Each created sequence has the shape (20, 4), where: \n","- $N_{obs}+N_{pred} = 8 + 12 = 20$;\n","- Positions + Speeds = ( $x_i,\\ y_i,\\ u_i,\\ v_i$) = ( $x_i,\\ y_i,\\ x_{i+1}-x_{i},\\ y_{i+1} - y_i$ )\n","\n","You can easily switch input type from position to speed setting the corresponding variable.\n","\n","Speeds $u_i, v_i$ are generally more robust input and allow to avoid problems with the reference system.\n","\n","------ \n","\n","Note: that $(u_0, v_0) = (0,0)$ and if speed are used the observed sequence has temporal length of $N_{obs} - 1$."]},{"cell_type":"code","execution_count":9,"metadata":{"colab":{"base_uri":"https://localhost:8080/"},"executionInfo":{"elapsed":47916,"status":"ok","timestamp":1665047365272,"user":{"displayName":"Alessandro Flaborea","userId":"12166823913657129807"},"user_tz":-120},"id":"WYJVV6RlUuYi","outputId":"658fe46f-d657-4019-a42f-9549ac1bc391"},"outputs":[{"name":"stdout","output_type":"stream","text":["start loading dataset\n","validation set size -> 0\n","001 / 007 - loading crowds_zara03_train.txt\n","002 / 007 - loading students003_train.txt\n","003 / 007 - loading uni_examples_train.txt\n","004 / 007 - loading biwi_eth_train.txt\n","005 / 007 - loading crowds_zara02_train.txt\n","006 / 007 - loading biwi_hotel_train.txt\n","007 / 007 - loading students001_train.txt\n","start loading dataset\n","validation set size -> 0\n","001 / 007 - loading biwi_eth_val.txt\n","002 / 007 - loading crowds_zara02_val.txt\n","003 / 007 - loading uni_examples_val.txt\n","004 / 007 - loading students001_val.txt\n","005 / 007 - loading crowds_zara03_val.txt\n","006 / 007 - loading students003_val.txt\n","007 / 007 - loading biwi_hotel_val.txt\n","start loading dataset\n","validation set size -> 0\n","001 / 001 - loading crowds_zara01.txt\n"]}],"source":["# Arguments to setup the datasets\n","dataset_name = 'zara1'\n","framework = 'regr'\n","obs_num = 8\n","preds_num = 12\n","\n","# We limit the number of samples to a fixed percentage for the sake of time\n","perc_data = 50\n","\n","# With predefined function we create dataset according to arguments\n","train_dataset,_ = baselineUtils.create_dataset('datasets', dataset_name, 0, obs_num, preds_num, delim='\\t', train=True, perc_data=perc_data, verbose=True)\n","val_dataset, _  = baselineUtils.create_dataset('datasets', dataset_name, 0, obs_num, preds_num, delim='\\t', train=False, perc_data=perc_data, verbose=True)\n","test_dataset, _ = baselineUtils.create_dataset('datasets', dataset_name, 0, obs_num, preds_num, delim='\\t', train=False, eval=True, verbose=True)\n","\n","# We create some folders to save model checkpoints\n","if not os.path.isdir(\"save_folder\"):\n","    os.mkdir(\"save_folder\")\n","if not os.path.isdir(\"save_folder/\"+framework):\n","    os.mkdir(\"save_folder/\"+framework)\n","\n","if not os.path.isdir(\"save_folder/\"+framework+\"/\"+dataset_name):\n","    os.mkdir(\"save_folder/\"+framework+\"/\"+dataset_name)"]},{"cell_type":"code","execution_count":10,"metadata":{},"outputs":[],"source":["input_type = 'speed'\n","\n","if input_type == 'speed':\n","    input_idx_1 = 2\n","    input_idx_2 = 4\n","    first_element = 1\n","    \n","elif input_type == 'position':\n","    input_idx_1 = 0\n","    input_idx_2 = 2\n","    first_element = 0"]},{"cell_type":"markdown","metadata":{"id":"zHh0nmOgpXk4"},"source":["We compute the mean and standard deviation of positions or speeds across the full training dataset and use those to normalize each entry in the sequence.\n","This normalization is beneficial prior to processing with neural networks."]},{"cell_type":"code","execution_count":11,"metadata":{"id":"8snhQNZ_UuQh"},"outputs":[],"source":["# After concatenating each observed and target sequence we compute the mean and std\n","mean = torch.cat((train_dataset[:]['src'][:, first_element:, input_idx_1:input_idx_2], train_dataset[:]['trg'][:, :, input_idx_1:input_idx_2]), 1).mean((0,1))\n","std  = torch.cat((train_dataset[:]['src'][:, first_element:, input_idx_1:input_idx_2], train_dataset[:]['trg'][:, :, input_idx_1:input_idx_2]), 1).std((0,1))"]},{"cell_type":"markdown","metadata":{"id":"7UIKe1xfojMb"},"source":["Following we create a torch dataloader that create the batches for each epoch."]},{"cell_type":"code","execution_count":12,"metadata":{"id":"FV4eX3VCXv9i"},"outputs":[],"source":["batch_size = 512\n","\n","tr_dl   = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=0)\n","val_dl  = torch.utils.data.DataLoader(val_dataset,   batch_size=batch_size, shuffle=True, num_workers=0)\n","test_dl = torch.utils.data.DataLoader(test_dataset,  batch_size=batch_size, shuffle=False, num_workers=0)"]},{"cell_type":"markdown","metadata":{"id":"Gl5VWs0Fo1EZ"},"source":["## Model instantiation"]},{"cell_type":"markdown","metadata":{},"source":["We create an instance of our transformer with the chosen configuration. \n","\n","Then we allocate it to the GPU for forward and backward accelerated computation."]},{"cell_type":"code","execution_count":13,"metadata":{"id":"xuDFofGtUuWA"},"outputs":[],"source":["# The input for the encoder are speeds (u,v) or positions (x,y)\n","enc_input_size = 2\n","# The input for the decoder are speeds (u,v) or positions (x,y) concatenated with mask array for start_of_sequence token [0, 0] \n","# Corresponding to start of sequence token the mask is 1 for the other speed input the mask is 0\n","dec_input_size = 3\n","# The output of the decoder are predicted speeds and corresponding mask that should be all zero (a loss for that is dedicated)\n","dec_output_size = 3\n","\n","emb_size = 512\n","ff_size = 1024\n","heads = 8\n","layers = 6\n","dropout = 0.1\n","\n","model = individual_TF.IndividualTF(enc_input_size, dec_input_size, dec_output_size, N=layers, d_model=emb_size, d_ff=ff_size, h=heads, dropout=dropout).to(device)"]},{"cell_type":"markdown","metadata":{"id":"WoeQw7ymp5dO"},"source":["## Training and Validation Step"]},{"cell_type":"markdown","metadata":{},"source":["Here we create two classes that define the single iteration function for train and validation."]},{"cell_type":"code","execution_count":14,"metadata":{"id":"n5NUO2pMUuOK"},"outputs":[],"source":["def train_step(model, batch, mean, std, device):\n","\n","    # If input type is speed then input (or source 'src') has shape (B, N_obs-1, 2) because the first one is (0,0).\n","    # Otherwise, if input type is position then input  has shape (B, N_obs, 2).\n","    # Note that the input of the decoder are only the first  N_pred-1  GT future value then target ('trg') has shape (B, N_pred-1, 2).\n","    inp    = (batch['src'][:,  first_element:, input_idx_1:input_idx_2].to(device) - mean.to(device)) / std.to(device)\n","    target = (batch['trg'][:, :-1, input_idx_1:input_idx_2].to(device) - mean.to(device)) / std.to(device)\n","\n","    # We create a third mask channel to append to the 2 speeds. \n","    # This helps the decoder differentiating between start of sequence token (with mask token 1) and target speeds (with mask token 0)\n","    # Summarizing: start_of_seq token is (0,0) and the mask token is 1 ---> [0, 0, 1]\n","    #              target inputs are (u_i, v_i) and the mask token is 0 ---> [u_i, v_i, 0]\n","    start_of_seq = torch.Tensor([0, 0, 1]).unsqueeze(0).unsqueeze(1).repeat(target.shape[0], 1, 1).to(device)\n","    target_c = torch.zeros((target.shape[0], target.shape[1], 1)).to(device)\n","    target = torch.cat((target, target_c), -1)\n","    # Final decoder input is the concatenation of them along temporal dimension\n","    dec_inp = torch.cat((start_of_seq, target), 1)\n","\n","    # Source attention is enabled between all the observed input (mask elements are setted to 1)\n","    src_att = torch.ones((inp.shape[0], 1, inp.shape[1])).to(device)\n","    # For the target attention we mask future elements to prevent model cheating (corresponding future mask elements are setted to False)\n","    # The mask is changed dinamically to use teacher forcing learning\n","    trg_att = subsequent_mask(dec_inp.shape[1]).repeat(dec_inp.shape[0], 1, 1).to(device)\n","    # Source, target and corresponding attention mask are passed to the model for the forward step\n","    pred = model(inp, dec_inp, src_att, trg_att)\n","\n","    return pred\n","\n","\n","def eval_step(model, batch, mean, std, device, preds=12):\n","\n","    # In the evaluation step we don't provide target to the decoder but we autoregressively input each prediction for the following one.\n","    inp = (batch['src'][:, first_element:, input_idx_1:input_idx_2].to(device) - mean.to(device)) / std.to(device)\n","\n","    # The decoder input is the only start of sequence token [0, 0, 1]\n","    # Please note that now model has to predict also the third channel mask (See loss2 in the main loop)\n","    src_att = torch.ones((inp.shape[0], 1, inp.shape[1])).to(device)\n","    start_of_seq = torch.Tensor([0, 0, 1]).unsqueeze(0).unsqueeze(1).repeat(inp.shape[0], 1, 1).to(device)\n","    dec_inp = start_of_seq\n","\n","    # We predict just one future speed and we append it to the decoder input for the next iteration (auto-regression)\n","    # At each step the target mask should be adapted\n","    for i in range(preds):\n","        trg_att = subsequent_mask(dec_inp.shape[1]).repeat(dec_inp.shape[0], 1, 1).to(device)\n","        out = model(inp, dec_inp, src_att, trg_att)\n","        dec_inp = torch.cat((dec_inp, out[:, -1:, :]), 1)\n","\n","    # Note at the each iteration of the loop we re-append the start of seq token, so after the last iteration we need to remove it\n","    return dec_inp[:, 1:, :]"]},{"cell_type":"markdown","metadata":{"id":"OJ0v-7sP5FUY"},"source":["## Optimizer"]},{"cell_type":"markdown","metadata":{},"source":["Here we select the **optimizer** proposed in the original Transformer Networks paper of Vaswani et al.\n","\n","It uses some initial warmup epochs, where the learning rate is increased. Then it slowly decreases according to a number of epoch and the chosen embedding size. The resulting formula is:\n","\n","LR = $\\frac{F}{\\sqrt{D}} min( \\frac{1}{\\sqrt{epoch}},\\ epoch \\cdot W^{-\\frac{3}{2}}) $\n","\n","where F is a scaling factor, D is the model embedding size, W is the number of warmup epochs."]},{"cell_type":"code","execution_count":null,"metadata":{"id":"xjgmsfmEUuTN"},"outputs":[],"source":["# Argument for the optimizer \n","factor = 1.\n","warmup = 5\n","\n","optim = NoamOpt(emb_size, factor, len(tr_dl)*warmup, torch.optim.Adam(model.parameters(), lr=0, betas=(0.9, 0.98), eps=1e-9))"]},{"cell_type":"markdown","metadata":{"id":"TZRUdtwQ5vTx"},"source":["## Main "]},{"cell_type":"markdown","metadata":{},"source":["\n","Then we can train, validate and test our transformer epoch by epoch.\n","\n","-------\n","\n","The **losses** used are 2:\n","\n","1.   $L_2$-loss distance between predicted $(\\hat{\\textbf{u}}, \\hat{\\textbf{v}})$ and GT $(\\textbf{u}, \\textbf{v})$ target speeds;\n","2.   $L_1$-loss for the target token mask. Note these should be all zero, so the loss is simply the mean.\n","\n","-------\n","\n","Moreover, the **metrics** used to validate the model goodness at Validation and Test time are the following:\n","\n","1.   Mean Average Displacement (MAD): $L_2$-distance between *all* the $N_pred$ GT and predicted future ***positions***;\n","2.   Final Average Displacement (FAD): $L_2$-distance between the *last* GT and predicted future ***positions***;\n","\n","-------\n","\n","Note: If you restart the training for any reason, remember to instanciate again model and optimizer in order to reset them.\n","\n","-------"]},{"cell_type":"code","execution_count":16,"metadata":{"colab":{"base_uri":"https://localhost:8080/","height":721},"executionInfo":{"elapsed":819,"status":"ok","timestamp":1665047397760,"user":{"displayName":"Alessandro Flaborea","userId":"12166823913657129807"},"user_tz":-120},"id":"3z4ndfv8UuLU","outputId":"748188e1-7e27-45f0-b8a6-bf881b8557b2"},"outputs":[{"name":"stdout","output_type":"stream","text":["Start Training...\n","\n","---> Epoch 010/100 <---  LR: 0.00222\n","Total Train Loss: 210.6400 - MAD:  0.2933 - FAD:  0.4105\n","Total Eval  Loss: 598.9922 - MAD:  1.0443 - FAD:  2.0041\n","Total Test  Loss: 695.5201 - MAD:  1.7321 - FAD:  2.8603 \n","\n","---> Epoch 020/100 <---  LR: 0.00179\n","Total Train Loss: 119.3469 - MAD:  0.1324 - FAD:  0.1960\n","Total Eval  Loss: 255.3613 - MAD:  0.4306 - FAD:  0.9609\n","Total Test  Loss: 267.4828 - MAD:  0.4663 - FAD:  1.0341 \n","\n","Saving checkpoint... \n"," \n","---> Epoch 030/100 <---  LR: 0.00145\n","Total Train Loss: 116.3218 - MAD:  0.1264 - FAD:  0.1885\n","Total Eval  Loss: 312.5784 - MAD:  0.5698 - FAD:  1.3008\n","Total Test  Loss: 344.2546 - MAD:  0.6154 - FAD:  1.4784 \n","\n","---> Epoch 040/100 <---  LR: 0.00125\n","Total Train Loss: 90.8324 - MAD:  0.1089 - FAD:  0.1587\n","Total Eval  Loss: 259.1963 - MAD:  0.4592 - FAD:  1.0382\n","Total Test  Loss: 283.3514 - MAD:  0.4906 - FAD:  1.1192 \n","\n","Saving checkpoint... \n"," \n"]}],"source":["# compute execution time of the cell\n","start_time = time.time()\n","\n","# Argument for the training \n","epoch = 0\n","max_epoch = 40          # Total number of epoch\n","ckp_save_step = 20      # Frequency for saving the model\n","log_step = 5           # Frequency for printing the loss\n","\n","\n","print(\"Start Training...\\n\")\n","\n","\n","for epoch in range(max_epoch):\n","\n","    if (epoch+1) % log_step == 0:\n","        print(\"---> Epoch %03i/%03i <---  LR: %7.5f\" % ((epoch+1), max_epoch, optim._rate))\n","\n","    ###### TRAIN ######\n","    model.train()\n","\n","    train_loss=0\n","    gt_posit = []\n","    pr_posit = []\n","    \n","    for id_b, batch in enumerate(tr_dl):\n","\n","        # All the gradients are resetted to zero before the training step\n","        optim.optimizer.zero_grad()\n","        \n","        # We predict target speeds and we save the corresponing GTs\n","        pred_speed = train_step(model, batch, mean, std, device)\n","        gt_speed = (batch['trg'][:, :, input_idx_1:input_idx_2].to(device) - mean.to(device)) / std.to(device)\n","\n","        # We compute the two losses, averaging on the batch\n","        loss1 = F.pairwise_distance(pred_speed[:, :, :2].contiguous().view(-1, 2), gt_speed.contiguous().view(-1, 2).to(device)).mean()\n","        loss2 = torch.abs(pred_speed[:, :, 2]).mean()\n","        loss = loss1 + loss2\n","\n","        # We accumulate and visualize the loss at the end of the epoch. \n","        # Note that here the loss is the mean on the batch but in the end we want the mean across the whole dataset.\n","        train_loss += loss.item() * batch['trg'].shape[0]\n","\n","        loss.backward()\n","        optim.step()\n","\n","        if input_type == 'speed':\n","            # If input type is speed, to compute MAD and FAD metrics we need to compute back positions from predicted speeds.\n","            # This is done easily adding cumulative (and denormalized) speeds to the last observed position. If last position in the input is (x_7, y_7) then:\n","            # (x_8, y_8)    =   (x_7, y_7) + (u_7, v_7)\n","            # (x_9, y_9)    =   (x_8, y_8) + (u_8, v_8)   =   (x_7, y_7) + (u_7, v_7) + (u_8, v_8)\n","            # (x_10, y_10)  =   (x_9, y_9) + (u_9, v_9)   =   (x_7, y_7) + (u_7, v_7) + (u_8, v_8) + (u_9, v_9)\n","            # We have always the last observed position (x_7, y_7) and we add progressively the \"cumulated\" speeds\n","            preds_tr_b = batch['src'][:, -1:, 0:2].cpu().numpy() + (pred_speed[:, :, 0:2].detach() * std.to(device) + mean.to(device)).cpu().numpy().cumsum(1)\n","\n","        elif input_type == 'position':\n","            # If input type is position, we simply append the output\n","            preds_tr_b = (pred_speed[:, :, 0:2].detach() * std.to(device) + mean.to(device)).cpu().numpy()\n","        \n","        # We store both predicted and GT positions\n","        pr_posit.append(preds_tr_b)\n","        gt_posit.append(batch['trg'][:, :, 0:2])\n","        \n","\n","    # After concatenation we compute MAD and FAD metrics\n","    gt_posit = np.concatenate(gt_posit, 0)\n","    pr_posit = np.concatenate(pr_posit, 0)\n","    mad, fad, errs = baselineUtils.distance_metrics(gt_posit, pr_posit)\n","\n","    if (epoch+1) % log_step == 0:\n","        print('Total Train Loss: %7.4f - MAD: %7.4f - FAD: %7.4f' % (train_loss/len(tr_dl), mad, fad))\n","\n","\n","\n","    ###### VALIDATION ######\n","    # Here is all the same exept for eval_step and computation of MAD and FAD metrics\n","    with torch.no_grad():\n","        model.eval()\n","\n","        val_loss = 0\n","        gt_posit = []\n","        pr_posit = []\n","\n","        for id_b, batch in enumerate(val_dl):\n","            \n","            pred_speed = eval_step(model, batch, mean, std, device, preds=preds_num)\n","            gt_speed = (batch['trg'][:, :, input_idx_1:input_idx_2].to(device) - mean.to(device)) / std.to(device)\n","\n","            loss1 = F.pairwise_distance(pred_speed[:, :, 0:2].contiguous().view(-1, 2), gt_speed.contiguous().view(-1, 2).to(device)).mean()\n","            loss2 = torch.abs(pred_speed[:, :, 2]).mean()\n","            loss = loss1 + loss2\n","            val_loss += loss.item() * batch['trg'].shape[0]\n","\n","            if input_type == 'speed':\n","                preds_tr_b = batch['src'][:, -1:, 0:2].cpu().numpy() + (pred_speed[:, :, 0:2] * std.to(device) + mean.to(device)).cpu().numpy().cumsum(1)\n","            elif input_type == 'position':\n","                preds_tr_b = (pred_speed[:, :, 0:2] * std.to(device) + mean.to(device)).cpu().numpy()\n","            \n","            # We store both predicted and GT positions\n","            pr_posit.append(preds_tr_b)\n","            gt_posit.append(batch['trg'][:, :, 0:2])\n","            \n","        # After concatenation we compute MAD and FAD metrics\n","        gt_posit = np.concatenate(gt_posit, 0)\n","        pr_posit = np.concatenate(pr_posit, 0)\n","        mad, fad, errs = baselineUtils.distance_metrics(gt_posit, pr_posit)\n","\n","        if (epoch+1) % log_step == 0:\n","            print('Total Eval  Loss: %7.4f - MAD: %7.4f - FAD: %7.4f' % (val_loss/len(val_dl), mad, fad))\n","\n","\n","\n","    ###### TEST ######\n","    # The test is same as eval \n","    with torch.no_grad():\n","        model.eval()\n","\n","        test_loss = 0\n","        gt = []\n","        pr = []\n","        \n","        for id_b, batch in enumerate(test_dl):\n","\n","            pred_speed = eval_step(model, batch, mean, std, device, preds=preds_num)\n","            gt_speed = (batch['trg'][:, :, input_idx_1:input_idx_2].to(device) - mean.to(device)) / std.to(device)\n","\n","            loss1 = F.pairwise_distance(pred_speed[:, :, 0:2].contiguous().view(-1, 2), gt_speed.contiguous().view(-1, 2).to(device)).mean()\n","            loss2 = torch.abs(pred_speed[:, :, 2]).mean()\n","            loss = loss1 + loss2\n","            test_loss += loss.item() * batch['trg'].shape[0]\n","\n","\n","            if input_type == 'speed':\n","                preds_tr_b = batch['src'][:, -1:, 0:2].cpu().numpy() + (pred_speed[:, :, 0:2] * std.to(device) + mean.to(device)).cpu().numpy().cumsum(1)\n","            elif input_type == 'position':\n","                preds_tr_b = (pred_speed[:, :, 0:2] * std.to(device) + mean.to(device)).cpu().numpy()\n","\n","            pr.append(preds_tr_b)\n","            gt.append(batch['trg'][:, :, 0:2])\n","\n","\n","        gt = np.concatenate(gt, 0)\n","        pr = np.concatenate(pr, 0)\n","        mad, fad, errs = baselineUtils.distance_metrics(gt, pr)\n","\n","        if (epoch+1) % log_step == 0:\n","            print('Total Test  Loss: %7.4f - MAD: %7.4f - FAD: %7.4f \\n'% (test_loss/len(test_dl), mad, fad))\n","\n","    # Here we save checkpoints to avoid repeated training\n","    if ((epoch+1) % (ckp_save_step) == 0):\n","        print(\"Saving checkpoint... \\n \")\n","        torch.save(model.state_dict(), f'save_folder/{framework}/{dataset_name}/{(epoch+1):05d}.pth')\n","\n","\n","\n","# print execution time\n","print(\"Total time: %s seconds\" % (time.time() - start_time))"]},{"cell_type":"markdown","metadata":{"id":"pt9drhMMFfuG"},"source":["## Load a model"]},{"cell_type":"markdown","metadata":{},"source":["\n","Here we leave a snippet of code to quickly load a model from a saved checkpoint. You can load model at specific epoch using this code before the main train-eval-test loop."]},{"cell_type":"code","execution_count":null,"metadata":{"id":"1AIk53ZHUuJW"},"outputs":[],"source":["# Instanciate a new model and loading its parameters\n","\n","# The input for the encoder are speeds (u,v) or positions (x,y)\n","enc_input_size = 2\n","# The input for the decoder are speeds (u,v) or positions (x,y) concatenated with mask array for start_of_sequence token [0, 0] \n","# Corresponding to start of sequence token the mask is 1 for the other speed input the mask is 0\n","dec_input_size = 3\n","# The output of the decoder are predicted speeds and corresponding mask that should be all zero (a loss for that is dedicated)\n","dec_output_size = 3\n","\n","emb_size = 512\n","ff_size = 2048\n","heads = 8\n","layers = 6\n","dropout = 0.1\n","\n","model = individual_TF.IndividualTF(enc_input_size, dec_input_size, dec_output_size, N=layers, d_model=emb_size, d_ff=ff_size, h=heads, dropout=dropout).to(device)\n","\n","\n","# Loading arguments\n","epoch = 50\n","dataset_name = 'zara1'\n","\n","path = f'save_folder/{dataset_name}/{(epoch):05d}.pth'\n","model.load_state_dict(torch.load(path))\n","\n","\n","# Setup correctly optimizer and its LR as well\n","factor = 1.\n","warmup = 10\n","\n","optim = NoamOpt(emb_size, factor, len(tr_dl)*warmup, torch.optim.Adam(model.parameters(), lr=0, betas=(0.9, 0.98), eps=1e-9))\n","optim._step = epoch-1"]},{"cell_type":"markdown","metadata":{},"source":["## Visualization  (***2 POINTS***)"]},{"cell_type":"markdown","metadata":{},"source":["Here you can implement some function to create qualitative plots.\n","\n","We recommend you the following:\n","\n","1. Loss plot for Train, Eval and Test;\n","2. MAD plot for Train, Eval and Test;\n","3. FAD plot for Train, Eval and Test;\n","4. Trajectory positions (observed points, GT target points and predicted target points)"]},{"cell_type":"code","execution_count":2,"metadata":{},"outputs":[],"source":["# Here your code "]},{"cell_type":"markdown","metadata":{},"source":["## Report  (***4 POINTS***)"]},{"cell_type":"markdown","metadata":{},"source":["Here you can report comments and results for the experiments up to this point.\n","\n","Perform experiments that improves the performances or that gives meaningfull insights.\n","\n","I.e. what happens if we change model hyperparamenters? What if we change learning rate?\n","\n","Please explain extensively the results and organize them clearly with tables, plots...\n","\n","----"]},{"cell_type":"markdown","metadata":{},"source":["Here your report"]},{"cell_type":"markdown","metadata":{},"source":["# Ablation Studies"]},{"cell_type":"markdown","metadata":{},"source":["\n","Here we ask you to change some settings in order to compare the benefit of some specific mechanism.\n","\n","Please follow the instructions and create a small report for each point adding your comments supported by plots, tables with results or whatever you think is usefull.\n","\n","Each extra study included to improve general performance or to draft a more complete analysis will be considered.\n","\n","---\n","\n","**Note:** to have a fair comparison we suggest to fix the setup (i.e. Regressive TF with speeds, obs=8, pred=12, ...) and change just the analysed module.\n","\n","---"]},{"cell_type":"markdown","metadata":{},"source":["## 1. Substitute for the Prediction Framework  (***6 POINTS***)"]},{"cell_type":"markdown","metadata":{},"source":["---\n","\n","The standard task is the regression of future speeds/positions. \n","\n","We propose to implement to different frameworks: Gaussian and Quantized.\n","\n","---"]},{"cell_type":"markdown","metadata":{},"source":["### a.  Gaussian"]},{"cell_type":"markdown","metadata":{},"source":["\n","Predicting normal distribution parameters mean vector $\\mu = (\\mu_x, \\mu_y)$ and covariance matrix $\\Sigma = \\biggl( \\begin{smallmatrix}\\sigma_x^2 & \\rho \\sigma_x \\sigma_y\\\\ \\rho \\sigma_x \\sigma_y & \\sigma_y^2 \\end{smallmatrix}\\biggr)$ of future predicition. \n","\n","Then the model output dimension is 5: 2 for mean parameters $\\mu_x, \\mu_y$ and 3 for the covariance parameters $\\sigma_x, \\sigma_y, \\rho$.\n","\n","---\n","\n","Note: consider carefully the following code snippet. In this way we force $\\sigma_x, \\sigma_y$ to be positive and $\\rho$ to be in $[-1, 1]$\n","\n","The following lines are meant to be a hint. Integrate those into the code of the previous cells."]},{"cell_type":"code","execution_count":null,"metadata":{},"outputs":[],"source":["# Arguments to setup the datasets\n","dataset_name = 'zara1'\n","framework = 'gauss'\n","obs_num = 8\n","preds_num = 12\n","\n","# We limit the number of samples to a fixed percentage for the sake of time\n","perc_data = 50\n","\n","# With predefined function we create dataset according to arguments\n","train_dataset,_ = baselineUtils.create_dataset('datasets', dataset_name, 0, obs_num, preds_num, delim='\\t', train=True, perc_data=perc_data, verbose=True)\n","val_dataset, _  = baselineUtils.create_dataset('datasets', dataset_name, 0, obs_num, preds_num, delim='\\t', train=False, perc_data=perc_data, verbose=True)\n","test_dataset, _ = baselineUtils.create_dataset('datasets', dataset_name, 0, obs_num, preds_num, delim='\\t', train=False, eval=True, verbose=True)\n","\n","# We create some folders to save model checkpoints\n","if not os.path.isdir(\"save_folder\"):\n","    os.mkdir(\"save_folder\")\n","if not os.path.isdir(\"save_folder/\"+framework):\n","    os.mkdir(\"save_folder/\"+framework)\n","if not os.path.isdir(\"save_folder/\"+framework+\"/\"+dataset_name):\n","    os.mkdir(\"save_folder/\"+framework+\"/\"+dataset_name)"]},{"cell_type":"code","execution_count":null,"metadata":{},"outputs":[],"source":["model = individual_TF.IndividualTF(enc_input_size, dec_input_size, dec_output_size, N=layers, d_model=emb_size, d_ff=ff_size, h=heads, dropout=dropout).to(device)\n","\n","output = model(inp, dec_inp, src_att, trg_att)\n","\n","mux = output[:, :, 0].unsqueeze(2)\n","muy = output[:, :, 1].unsqueeze(2)\n","sx = torch.exp(output[:, :, 2]).unsqueeze(2)\n","sy = torch.exp(output[:, :, 3]).unsqueeze(2)\n","corr = torch.tanh(output[:, :, 4]).unsqueeze(2)\n","\n","mean = torch.cat((mux, muy), dim=2).to(device)\n","cov = torch.cat((sx**2, corr*sx*sy, corr*sx*sy, sy**2), dim=2).view((-1, sx.size(1), 2, 2)).to(device)"]},{"cell_type":"markdown","metadata":{},"source":["Next prediction can be now sampled from the predicted distribution making the forecasting stochastic.\n","\n","The loss used in this case is the NLL.  \n","\n","Note: To relax the assumption you can also use predicted mean as input for following step (particularly in eval and test), avoiding the sampling and assuming identity as covariance matrix."]},{"cell_type":"code","execution_count":null,"metadata":{},"outputs":[],"source":["# Here your code "]},{"cell_type":"markdown","metadata":{},"source":["Here your report\n","\n","---"]},{"cell_type":"markdown","metadata":{},"source":["### b. Quantized"]},{"cell_type":"markdown","metadata":{},"source":["Transformer was originally introduced in the NLP i.e. for next word classification task.\n","    \n","To emulate this case we change dataset (clustering all possible speed in C classes) and model to classify the most likely one (with CE loss).\n","\n","Here we provide a script for the quantized dataset, so you may adapt the final part of the model to output probability score for each class (output_size=1000 + softmax) followed by CE loss.\n","\n","---\n","\n","Note: In the quantized framework the start of sequence token is adapted: \n","\n","The class indices spans from 0 to 999, so we add index 1000 to represent the start of sequence token\n","\n","The following lines are meant to be hint, wisely integrate them with the code in the previous cells.  "]},{"cell_type":"code","execution_count":null,"metadata":{},"outputs":[],"source":["# Arguments to setup the datasets\n","dataset_name = 'zara1'\n","framework = 'quant'\n","obs_num = 8\n","preds_num = 12\n","\n","# We limit the number of samples to a fixed percentage for the sake of time\n","perc_data = 50\n","\n","# With predefined function we create dataset according to arguments\n","train_dataset,_ = baselineUtils.create_dataset('datasets', dataset_name, 0, obs_num, preds_num, delim='\\t', train=True, perc_data=perc_data, verbose=True)\n","val_dataset, _  = baselineUtils.create_dataset('datasets', dataset_name, 0, obs_num, preds_num, delim='\\t', train=False, perc_data=perc_data, verbose=True)\n","test_dataset, _ = baselineUtils.create_dataset('datasets', dataset_name, 0, obs_num, preds_num, delim='\\t', train=False, eval=True, verbose=True)\n","\n","# Load precomputed clusters to quantize the data\n","mat = scipy.io.loadmat(os.path.join('datasets', dataset_name, \"clusters.mat\"))\n","clusters=mat['centroids']\n","num_classes = clusters.shape[0]\n","\n","# We create some folders to save model checkpoints\n","if not os.path.isdir(\"save_folder\"):\n","    os.mkdir(\"save_folder\")\n","if not os.path.isdir(\"save_folder/\"+framework):\n","    os.mkdir(\"save_folder/\"+framework)\n","if not os.path.isdir(\"save_folder/\"+framework+\"/\"+dataset_name):\n","    os.mkdir(\"save_folder/\"+framework+\"/\"+dataset_name)"]},{"cell_type":"code","execution_count":null,"metadata":{},"outputs":[],"source":["# The input for the encoder are speeds (u,v) or positions (x,y)\n","enc_input_size = num_classes\n","# The input for the decoder are speeds (u,v) or positions (x,y) concatenated with mask array for start_of_sequence token [0, 0] \n","# Corresponding to start of sequence token the mask is 1 for the other speed input the mask is 0\n","dec_input_size = num_classes+1\n","# The output of the decoder are predicted speeds and corresponding mask that should be all zero (a loss for that is dedicated)\n","dec_output_size = num_classes\n","\n","emb_size = 512\n","ff_size = 2048\n","heads = 8\n","layers = 6\n","dropout = 0.1\n","\n","\n","model = individual_TF.IndividualTF(enc_input_size, dec_input_size, dec_output_size, N=layers, d_model=emb_size, d_ff=ff_size, h=heads, dropout=dropout).to(device)\n","\n","\n","# Inside the train and eval step we need to convert speed/position to cluster index\n","batch_size = batch['src'].shape[0]\n","\n","# Associate the nearest class to each speed/position\n","speeds_inp=batch['src'][:,1:,2:4]\n","inp=torch.tensor(scipy.spatial.distance.cdist(speeds_inp.reshape(-1,2), clusters).argmin(axis=1).reshape(batch_size, -1)).to(device)\n","\n","speeds_trg = batch['trg'][:,:,2:4]\n","target = torch.tensor(scipy.spatial.distance.cdist(speeds_trg.reshape(-1, 2), clusters).argmin(axis=1).reshape(batch_size, -1)).to(device)\n","\n","\n","\n","\n","# Class are indices from 0 to 999. \n","# We add index 1000 to represent the start of sequence token\n","start_of_seq = torch.tensor([1000]).repeat(batch_size).unsqueeze(1).to(device)\n","\n","\n","\n","# We predict class indexes of future speeds/positions\n","output = model(inp, dec_inp, src_att, trg_att)\n","\n","loss = F.cross_entropy(output.view(-1, num_classes), target.view(-1), reduction='mean')\n","\n","\n","\n","# To compute metrics we need positions. Then we convert back each predicted index to the relative centroid speed/position values\n","preds_tr_b = batch['src'][:,-1:,0:2].cpu().numpy() + clusters[output.cpu().numpy()].cumsum(1)\n","pr.append(preds_tr_b)"]},{"cell_type":"code","execution_count":null,"metadata":{},"outputs":[],"source":["# Here your code "]},{"cell_type":"markdown","metadata":{},"source":["Here your report\n","\n","---"]},{"cell_type":"markdown","metadata":{},"source":["## 2. Increase Prediction Horizon (Short- or Long-term Forecasting)  (***3 POINTS***)"]},{"cell_type":"markdown","metadata":{},"source":["---\n","\n","You can easily increase/decrease the number of predictions (i.e. pred = 4, 8, 12, 20, 30, 50 ....) in the dataloader and see the effect on the MAD/FAD metric.\n","\n","Report your results in a table and/or plot and comment what you see.\n","\n","---"]},{"cell_type":"code","execution_count":null,"metadata":{},"outputs":[],"source":["# Here your code "]},{"cell_type":"markdown","metadata":{},"source":["Here your report\n","\n","---"]},{"cell_type":"markdown","metadata":{},"source":["## 3. Increasing Data Number  (***3 POINTS***)"]},{"cell_type":"markdown","metadata":{},"source":["Transformers are generally very large network and need a lot of data to perform well.\n","\n","Try to vary the percentage data variable (i.e. 10, 25, 50, 75, 100) and see how the performance changes.\n","\n","Please report here plots and/or tables for:\n","\n","1. MAD and FAD metrics \n","\n","2. Computational time"]},{"cell_type":"code","execution_count":null,"metadata":{},"outputs":[],"source":["# Here your code "]},{"cell_type":"markdown","metadata":{},"source":["Here your report\n","\n","---"]},{"cell_type":"markdown","metadata":{},"source":["## 4. Change input Type  (***2 POINTS*** - Bonus)"]},{"cell_type":"markdown","metadata":{},"source":["---\n","\n","What happens if we change the input form speed type (u,v) to position one (x,y)?\n","\n","Report then some quantitative results and plot trajectory predicted with both method to evaluate qualitative differences.\n","\n","---"]},{"cell_type":"code","execution_count":null,"metadata":{},"outputs":[],"source":["# Here your code "]},{"cell_type":"markdown","metadata":{},"source":["Here your report"]},{"cell_type":"markdown","metadata":{},"source":["## 5. Positional Encoding  (***3 POINTS*** - Bonus)"]},{"cell_type":"markdown","metadata":{},"source":["A number of positional encodings have been proposed. \n","\n","Implement the plain positional encoding [0,1,2,3,4,...] and report your comments and results.\n","\n","Change the commented class we prepared in the positional_encoding.py file and copy the class here. "]},{"cell_type":"code","execution_count":null,"metadata":{},"outputs":[],"source":["# Here your code "]},{"cell_type":"markdown","metadata":{},"source":["Here your report\n"]},{"cell_type":"markdown","metadata":{},"source":["---\n","\n","This notebook was created by Luca Franco and Alessandro Flaborea."]}],"metadata":{"accelerator":"GPU","colab":{"collapsed_sections":[],"provenance":[]},"kernelspec":{"display_name":"Python 3.10.6 64-bit","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.10.6"},"vscode":{"interpreter":{"hash":"31f2aee4e71d21fbe5cf8b01ff0e069b9275f58929596ceb00d14d90e3e16cd6"}}},"nbformat":4,"nbformat_minor":0}