train_model.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import numpy as np\n",
    "import torch\n",
    "from torch import nn\n",
    "from torch.optim import Adam\n",
    "from torch.utils.data import DataLoader\n",
    "from torch.optim.lr_scheduler import ReduceLROnPlateau\n",
    "from scipy.spatial.transform import Rotation as R\n",
    "import time\n",
    "import pickle\n",
    "import argparse \n",
    "from utils import *\n",
    "\n",
    "from deepVCP import DeepVCP\n",
    "from ModelNet40Dataset import ModelNet40Dataset\n",
    "from KITTIDataset import KITTIDataset\n",
    "from deepVCP_loss import deepVCP_loss\n",
    "\n",
    "import matplotlib\n",
    "matplotlib.use('Agg')\n",
    "from matplotlib import pyplot as plt\n",
    "  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = \"modelnet\"\n",
    "retrain_path = False\n",
    "model_path = \"final_model.pt\"\n",
    "full_dataset = True "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Params: epochs: 10, batch: 1, lr: 0.001, alpha: 0.5\n",
      "\n",
      "device: cuda\n",
      "# Total clouds 2\n",
      "# Total clouds 2\n",
      "Train dataset size:  2\n",
      "Test dataset size:  2\n"
     ]
    }
   ],
   "source": [
    "num_epochs = 10\n",
    "batch_size = 1\n",
    "lr = 0.001\n",
    "# loss balancing factor \n",
    "alpha = 0.5\n",
    "\n",
    "print(f\"Params: epochs: {num_epochs}, batch: {batch_size}, lr: {lr}, alpha: {alpha}\\n\")\n",
    "\n",
    "# check if cuda is available\n",
    "device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
    "print(f\"device: {device}\")\n",
    "\n",
    "# dataset \n",
    "\n",
    "root = 'modelnet40_normal_resampled/'\n",
    "shape_names = np.loadtxt(root+\"modelnet10_shape_names.txt\", dtype=\"str\")\n",
    "train_data= ModelNet40Dataset(root=root, augment=True, full_dataset=full_dataset, split='train')\n",
    "test_data = ModelNet40Dataset(root=root, augment=True, full_dataset=full_dataset,  split='test')\n",
    "\n",
    "\n",
    "train_loader = DataLoader(dataset=train_data, batch_size=batch_size, shuffle=False)\n",
    "test_loader = DataLoader(dataset=test_data, batch_size=batch_size, shuffle=False)\n",
    "\n",
    "\n",
    "num_train = len(train_data)\n",
    "num_test = len(test_data)\n",
    "print('Train dataset size: ', num_train)\n",
    "print('Test dataset size: ', num_test)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "No retrain\n",
      "epoch #0\n",
      "Processing file: bathtub_0001\n"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "Cell \u001b[1;32mIn[4], line 35\u001b[0m\n\u001b[0;32m     33\u001b[0m src, target, R_gt, t_gt \u001b[38;5;241m=\u001b[39m src\u001b[38;5;241m.\u001b[39mto(device), target\u001b[38;5;241m.\u001b[39mto(device), R_gt\u001b[38;5;241m.\u001b[39mto(device), t_gt\u001b[38;5;241m.\u001b[39mto(device)\n\u001b[0;32m     34\u001b[0m t_init \u001b[38;5;241m=\u001b[39m torch\u001b[38;5;241m.\u001b[39mzeros(\u001b[38;5;241m1\u001b[39m, \u001b[38;5;241m3\u001b[39m)\n\u001b[1;32m---> 35\u001b[0m src_keypts, target_vcp \u001b[38;5;241m=\u001b[39m \u001b[43mmodel\u001b[49m\u001b[43m(\u001b[49m\u001b[43msrc\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mtarget\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mR_gt\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mt_init\u001b[49m\u001b[43m)\u001b[49m\n\u001b[0;32m     36\u001b[0m \u001b[38;5;66;03m# print('src_keypts shape', src_keypts.shape)\u001b[39;00m\n\u001b[0;32m     37\u001b[0m \u001b[38;5;66;03m# print('target_vcp shape', target_vcp.shape)\u001b[39;00m\n\u001b[0;32m     38\u001b[0m \u001b[38;5;66;03m# zero gradient \u001b[39;00m\n\u001b[0;32m     39\u001b[0m optim\u001b[38;5;241m.\u001b[39mzero_grad()\n",
      "File \u001b[1;32mc:\\Users\\bingwang\\.conda\\envs\\InverseCW4\\lib\\site-packages\\torch\\nn\\modules\\module.py:1501\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[1;34m(self, *args, **kwargs)\u001b[0m\n\u001b[0;32m   1496\u001b[0m \u001b[39m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[0;32m   1497\u001b[0m \u001b[39m# this function, and just call forward.\u001b[39;00m\n\u001b[0;32m   1498\u001b[0m \u001b[39mif\u001b[39;00m \u001b[39mnot\u001b[39;00m (\u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_backward_hooks \u001b[39mor\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_backward_pre_hooks \u001b[39mor\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_forward_hooks \u001b[39mor\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_forward_pre_hooks\n\u001b[0;32m   1499\u001b[0m         \u001b[39mor\u001b[39;00m _global_backward_pre_hooks \u001b[39mor\u001b[39;00m _global_backward_hooks\n\u001b[0;32m   1500\u001b[0m         \u001b[39mor\u001b[39;00m _global_forward_hooks \u001b[39mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[1;32m-> 1501\u001b[0m     \u001b[39mreturn\u001b[39;00m forward_call(\u001b[39m*\u001b[39margs, \u001b[39m*\u001b[39m\u001b[39m*\u001b[39mkwargs)\n\u001b[0;32m   1502\u001b[0m \u001b[39m# Do not call functions when jit is used\u001b[39;00m\n\u001b[0;32m   1503\u001b[0m full_backward_hooks, non_full_backward_hooks \u001b[39m=\u001b[39m [], []\n",
      "File \u001b[1;32m\\\\evs2.cs.ucl.ac.uk\\student-msc\\cgvi\\2022\\bingwang\\GitHub\\DeepVCP-Pointcloud-Registration\\deepVCP.py:30\u001b[0m, in \u001b[0;36mDeepVCP.forward\u001b[1;34m(self, src_pts, tgt_pts, R_init, t_init)\u001b[0m\n\u001b[0;32m     28\u001b[0m \u001b[39m# deep features extracted from FE layer: B x N x 32\u001b[39;00m\n\u001b[0;32m     29\u001b[0m fe_start_time \u001b[39m=\u001b[39m time\u001b[39m.\u001b[39mtime()\n\u001b[1;32m---> 30\u001b[0m src_deep_feat_xyz, src_deep_feat_pts \u001b[39m=\u001b[39m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49mFE1(src_pts)\n\u001b[0;32m     31\u001b[0m \u001b[39mprint\u001b[39m(\u001b[39m\"\u001b[39m\u001b[39mfeature extraction time: \u001b[39m\u001b[39m\"\u001b[39m, time\u001b[39m.\u001b[39mtime() \u001b[39m-\u001b[39m fe_start_time)\n\u001b[0;32m     33\u001b[0m \u001b[39m# obtain the top k indices for src point clouds\u001b[39;00m\n",
      "File \u001b[1;32mc:\\Users\\bingwang\\.conda\\envs\\InverseCW4\\lib\\site-packages\\torch\\nn\\modules\\module.py:1501\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[1;34m(self, *args, **kwargs)\u001b[0m\n\u001b[0;32m   1496\u001b[0m \u001b[39m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[0;32m   1497\u001b[0m \u001b[39m# this function, and just call forward.\u001b[39;00m\n\u001b[0;32m   1498\u001b[0m \u001b[39mif\u001b[39;00m \u001b[39mnot\u001b[39;00m (\u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_backward_hooks \u001b[39mor\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_backward_pre_hooks \u001b[39mor\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_forward_hooks \u001b[39mor\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_forward_pre_hooks\n\u001b[0;32m   1499\u001b[0m         \u001b[39mor\u001b[39;00m _global_backward_pre_hooks \u001b[39mor\u001b[39;00m _global_backward_hooks\n\u001b[0;32m   1500\u001b[0m         \u001b[39mor\u001b[39;00m _global_forward_hooks \u001b[39mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[1;32m-> 1501\u001b[0m     \u001b[39mreturn\u001b[39;00m forward_call(\u001b[39m*\u001b[39margs, \u001b[39m*\u001b[39m\u001b[39m*\u001b[39mkwargs)\n\u001b[0;32m   1502\u001b[0m \u001b[39m# Do not call functions when jit is used\u001b[39;00m\n\u001b[0;32m   1503\u001b[0m full_backward_hooks, non_full_backward_hooks \u001b[39m=\u001b[39m [], []\n",
      "File \u001b[1;32m\\\\evs2.cs.ucl.ac.uk\\student-msc\\cgvi\\2022\\bingwang\\GitHub\\DeepVCP-Pointcloud-Registration\\deep_feat_extraction.py:28\u001b[0m, in \u001b[0;36mfeat_extraction_layer.forward\u001b[1;34m(self, pts)\u001b[0m\n\u001b[0;32m     26\u001b[0m output_xyz, output_pts \u001b[39m=\u001b[39m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39msa1(xyz, normal)\n\u001b[0;32m     27\u001b[0m output_xyz, output_pts \u001b[39m=\u001b[39m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39msa2(output_xyz, normal)\n\u001b[1;32m---> 28\u001b[0m output_xyz, output_pts \u001b[39m=\u001b[39m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49msa3(output_xyz, normal)\n\u001b[0;32m     29\u001b[0m output_xyz \u001b[39m=\u001b[39m output_xyz\u001b[39m.\u001b[39mpermute(\u001b[39m0\u001b[39m, \u001b[39m2\u001b[39m, \u001b[39m1\u001b[39m)\n\u001b[0;32m     30\u001b[0m output_pts \u001b[39m=\u001b[39m output_pts\u001b[39m.\u001b[39mpermute(\u001b[39m0\u001b[39m, \u001b[39m2\u001b[39m, \u001b[39m1\u001b[39m)\n",
      "File \u001b[1;32mc:\\Users\\bingwang\\.conda\\envs\\InverseCW4\\lib\\site-packages\\torch\\nn\\modules\\module.py:1501\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[1;34m(self, *args, **kwargs)\u001b[0m\n\u001b[0;32m   1496\u001b[0m \u001b[39m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[0;32m   1497\u001b[0m \u001b[39m# this function, and just call forward.\u001b[39;00m\n\u001b[0;32m   1498\u001b[0m \u001b[39mif\u001b[39;00m \u001b[39mnot\u001b[39;00m (\u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_backward_hooks \u001b[39mor\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_backward_pre_hooks \u001b[39mor\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_forward_hooks \u001b[39mor\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_forward_pre_hooks\n\u001b[0;32m   1499\u001b[0m         \u001b[39mor\u001b[39;00m _global_backward_pre_hooks \u001b[39mor\u001b[39;00m _global_backward_hooks\n\u001b[0;32m   1500\u001b[0m         \u001b[39mor\u001b[39;00m _global_forward_hooks \u001b[39mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[1;32m-> 1501\u001b[0m     \u001b[39mreturn\u001b[39;00m forward_call(\u001b[39m*\u001b[39margs, \u001b[39m*\u001b[39m\u001b[39m*\u001b[39mkwargs)\n\u001b[0;32m   1502\u001b[0m \u001b[39m# Do not call functions when jit is used\u001b[39;00m\n\u001b[0;32m   1503\u001b[0m full_backward_hooks, non_full_backward_hooks \u001b[39m=\u001b[39m [], []\n",
      "File \u001b[1;32m\\\\evs2.cs.ucl.ac.uk\\student-msc\\cgvi\\2022\\bingwang\\GitHub\\DeepVCP-Pointcloud-Registration\\pointnet2_utils.py:192\u001b[0m, in \u001b[0;36mPointNetSetAbstraction.forward\u001b[1;34m(self, xyz, points)\u001b[0m\n\u001b[0;32m    190\u001b[0m     new_xyz, new_points \u001b[39m=\u001b[39m sample_and_group_all(xyz, points)\n\u001b[0;32m    191\u001b[0m \u001b[39melse\u001b[39;00m:\n\u001b[1;32m--> 192\u001b[0m     new_xyz, new_points \u001b[39m=\u001b[39m sample_and_group(\u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49mnpoint, \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49mradius, \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49mnsample, xyz, points)\n\u001b[0;32m    193\u001b[0m \u001b[39m# new_xyz: sampled points position data, [B, npoint, C]\u001b[39;00m\n\u001b[0;32m    194\u001b[0m \u001b[39m# new_points: sampled points data, [B, npoint, nsample, C+D]\u001b[39;00m\n\u001b[0;32m    195\u001b[0m new_points \u001b[39m=\u001b[39m new_points\u001b[39m.\u001b[39mpermute(\u001b[39m0\u001b[39m, \u001b[39m3\u001b[39m, \u001b[39m2\u001b[39m, \u001b[39m1\u001b[39m) \u001b[39m# [B, C+D, nsample,npoint]\u001b[39;00m\n",
      "File \u001b[1;32m\\\\evs2.cs.ucl.ac.uk\\student-msc\\cgvi\\2022\\bingwang\\GitHub\\DeepVCP-Pointcloud-Registration\\pointnet2_utils.py:124\u001b[0m, in \u001b[0;36msample_and_group\u001b[1;34m(npoint, radius, nsample, xyz, points, returnidx)\u001b[0m\n\u001b[0;32m    122\u001b[0m B, N, C \u001b[39m=\u001b[39m xyz\u001b[39m.\u001b[39mshape\n\u001b[0;32m    123\u001b[0m S \u001b[39m=\u001b[39m npoint\n\u001b[1;32m--> 124\u001b[0m fps_idx \u001b[39m=\u001b[39m farthest_point_sample(xyz, npoint) \u001b[39m# [B, npoint, C]\u001b[39;00m\n\u001b[0;32m    125\u001b[0m new_xyz \u001b[39m=\u001b[39m index_points(xyz, fps_idx)\n\u001b[0;32m    126\u001b[0m idx \u001b[39m=\u001b[39m query_ball_point(radius, nsample, xyz, new_xyz)\n",
      "File \u001b[1;32m\\\\evs2.cs.ucl.ac.uk\\student-msc\\cgvi\\2022\\bingwang\\GitHub\\DeepVCP-Pointcloud-Registration\\pointnet2_utils.py:82\u001b[0m, in \u001b[0;36mfarthest_point_sample\u001b[1;34m(xyz, npoint)\u001b[0m\n\u001b[0;32m     80\u001b[0m     dist \u001b[39m=\u001b[39m torch\u001b[39m.\u001b[39msum((xyz \u001b[39m-\u001b[39m centroid) \u001b[39m*\u001b[39m\u001b[39m*\u001b[39m \u001b[39m2\u001b[39m, \u001b[39m-\u001b[39m\u001b[39m1\u001b[39m)\n\u001b[0;32m     81\u001b[0m     mask \u001b[39m=\u001b[39m dist \u001b[39m<\u001b[39m distance\n\u001b[1;32m---> 82\u001b[0m     distance[mask] \u001b[39m=\u001b[39m dist[mask]\u001b[39m.\u001b[39;49mfloat()\n\u001b[0;32m     83\u001b[0m     farthest \u001b[39m=\u001b[39m torch\u001b[39m.\u001b[39mmax(distance, \u001b[39m-\u001b[39m\u001b[39m1\u001b[39m)[\u001b[39m1\u001b[39m]\n\u001b[0;32m     84\u001b[0m \u001b[39mreturn\u001b[39;00m centroids\n",
      "\u001b[1;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "use_normal = False if dataset == \"kitti\" else True\n",
    "\n",
    "# Initialize the model\n",
    "model = DeepVCP(use_normal=use_normal)\n",
    "if torch.cuda.device_count() > 1:\n",
    "    print(\"Let's use\", torch.cuda.device_count(), \"GPUs!\")\n",
    "    # dim = 0 [30, xxx] -> [10, ...], [10, ...], [10, ...] on 3 GPUs\n",
    "    model = nn.DataParallel(model)\n",
    "\n",
    "model.to(device)\n",
    "\n",
    "# Retrain\n",
    "if retrain_path:\n",
    "    print(\"Retrain on \", retrain_path)\n",
    "    model.load_state_dict(torch.load(retrain_path))\n",
    "else:\n",
    "    print(\"No retrain\")\n",
    "\n",
    "# Define the optimizer\n",
    "optim = Adam(model.parameters(), lr=lr)\n",
    "\n",
    "# begin train \n",
    "model.train()\n",
    "loss_epoch_avg = []\n",
    "for epoch in range(num_epochs):\n",
    "    print(f\"epoch #{epoch}\")\n",
    "    loss_epoch = []\n",
    "    running_loss = 0.0\n",
    "    \n",
    "    for n_batch, (src, target, R_gt, t_gt, ) in enumerate(train_loader):\n",
    "        start_time = time.time()\n",
    "        # mini batch\n",
    "        src, target, R_gt, t_gt = src.to(device), target.to(device), R_gt.to(device), t_gt.to(device)\n",
    "        t_init = torch.zeros(1, 3)\n",
    "        src_keypts, target_vcp = model(src, target, R_gt, t_init)\n",
    "        # print('src_keypts shape', src_keypts.shape)\n",
    "        # print('target_vcp shape', target_vcp.shape)\n",
    "        # zero gradient \n",
    "        optim.zero_grad()\n",
    "        loss, R_pred, t_pred = deepVCP_loss(src_keypts, target_vcp, R_gt, t_gt, alpha=0.5)\n",
    "\n",
    "        # error metric for rigid body transformation\n",
    "        r_pred = R.from_matrix(R_pred.squeeze(0).cpu().detach().numpy())\n",
    "        r_pred_arr = torch.tensor(r_pred.as_euler('xyz', degrees=True)).reshape(1, 3)\n",
    "        r_gt = R.from_matrix(R_gt.squeeze(0).cpu().detach().numpy())\n",
    "        r_gt_arr = torch.tensor(r_gt.as_euler('xyz', degrees=True)).reshape(1, 3)\n",
    "        pdist = nn.PairwiseDistance(p = 2)\n",
    "        \n",
    "        t_gt = t_gt.reshape(1,3)\n",
    "        t_pred = t_pred.reshape(1,3)\n",
    "        print(\"rotation error: \", pdist(r_pred_arr, r_gt_arr).item())\n",
    "        print(\"translation error: \", pdist(t_pred, t_gt).item())\n",
    "\n",
    "        # backward pass\n",
    "        loss.backward()\n",
    "        # update parameters \n",
    "        optim.step()\n",
    "\n",
    "        running_loss += loss.item()\n",
    "        loss_epoch += [loss.item()]\n",
    "        print(\"--- %s seconds ---\" % (time.time() - start_time))\n",
    "        if (n_batch + 1) % 5 == 0:\n",
    "            print(\"Epoch: [{}/{}], Batch: {}, Loss: {}\".format(\n",
    "                epoch, num_epochs, n_batch, loss.item()))\n",
    "            running_loss = 0.0\n",
    "    \n",
    "    torch.save(model.state_dict(), \"epoch_\" + str(epoch) + \"_model.pt\")\n",
    "    loss_epoch_avg += [sum(loss_epoch) / len(loss_epoch)]\n",
    "    with open(\"training_loss_\" + str(epoch) + \".txt\", \"wb\") as fp:   #Pickling\n",
    "        pickle.dump(loss_epoch, fp)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.10.10 ('InverseCW4')",
   "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.10"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "709ee83345f2a73a4810e0cadc6e4be23f2f850d1ef51d9b25c99df9df8dd8b0"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}