training-motion-prediction-with-gnn.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Motion Prediction with Graph Neural Network"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Getting Started with this notebook:\n",
    "- You might need to add the utility script: kaggle_l5kit\n",
    "- Also activate the GPU as an accelerator on the right of your screen"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Summary"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "* optimizer: Adam\n",
    "* scheduler: CosineAnnealing with start lr of 1e-3 and min_lr of 1e-5, T_max = n_epochs\n",
    "* Batch_size : 150 (choose as great as possible)\n",
    "* Could train up to 25 epochs during 9h\n",
    "* Training module : We used Pytorch Lighning as well as torch_geometric, the geometric deep learning module of Pytorch. \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Looking in links: https://pytorch-geometric.com/whl/torch-1.6.0.html\n",
      "Collecting torch-scatter==latest+cu102\n",
      "  Using cached https://pytorch-geometric.com/whl/torch-1.6.0/torch_scatter-latest%2Bcu102-cp36-cp36m-linux_x86_64.whl (11.6 MB)\n",
      "\u001b[31mERROR: Error while checking for conflicts. Please file an issue on pip's issue tracker: https://github.com/pypa/pip/issues/new\n",
      "Traceback (most recent call last):\n",
      "  File \"/home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages/pip/_vendor/pkg_resources/__init__.py\", line 3021, in _dep_map\n",
      "    return self.__dep_map\n",
      "  File \"/home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages/pip/_vendor/pkg_resources/__init__.py\", line 2815, in __getattr__\n",
      "    raise AttributeError(attr)\n",
      "AttributeError: _DistInfoDistribution__dep_map\n",
      "\n",
      "During handling of the above exception, another exception occurred:\n",
      "\n",
      "Traceback (most recent call last):\n",
      "  File \"/home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages/pip/_vendor/pkg_resources/__init__.py\", line 3012, in _parsed_pkg_info\n",
      "    return self._pkg_info\n",
      "  File \"/home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages/pip/_vendor/pkg_resources/__init__.py\", line 2815, in __getattr__\n",
      "    raise AttributeError(attr)\n",
      "AttributeError: _pkg_info\n",
      "\n",
      "During handling of the above exception, another exception occurred:\n",
      "\n",
      "Traceback (most recent call last):\n",
      "  File \"/home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages/pip/_internal/commands/install.py\", line 520, in _determine_conflicts\n",
      "    return check_install_conflicts(to_install)\n",
      "  File \"/home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages/pip/_internal/operations/check.py\", line 108, in check_install_conflicts\n",
      "    package_set, _ = create_package_set_from_installed()\n",
      "  File \"/home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages/pip/_internal/operations/check.py\", line 50, in create_package_set_from_installed\n",
      "    package_set[name] = PackageDetails(dist.version, dist.requires())\n",
      "  File \"/home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages/pip/_vendor/pkg_resources/__init__.py\", line 2736, in requires\n",
      "    dm = self._dep_map\n",
      "  File \"/home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages/pip/_vendor/pkg_resources/__init__.py\", line 3023, in _dep_map\n",
      "    self.__dep_map = self._compute_dependencies()\n",
      "  File \"/home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages/pip/_vendor/pkg_resources/__init__.py\", line 3032, in _compute_dependencies\n",
      "    for req in self._parsed_pkg_info.get_all('Requires-Dist') or []:\n",
      "  File \"/home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages/pip/_vendor/pkg_resources/__init__.py\", line 3014, in _parsed_pkg_info\n",
      "    metadata = self.get_metadata(self.PKG_INFO)\n",
      "  File \"/home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages/pip/_vendor/pkg_resources/__init__.py\", line 1420, in get_metadata\n",
      "    value = self._get(path)\n",
      "  File \"/home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages/pip/_vendor/pkg_resources/__init__.py\", line 1616, in _get\n",
      "    with open(path, 'rb') as stream:\n",
      "FileNotFoundError: [Errno 2] No such file or directory: '/home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages/s3transfer-0.3.3.dist-info/METADATA'\u001b[0m\n",
      "Installing collected packages: torch-scatter\n",
      "  Attempting uninstall: torch-scatter\n",
      "    Found existing installation: torch-scatter 2.0.5\n",
      "    Uninstalling torch-scatter-2.0.5:\n",
      "      Successfully uninstalled torch-scatter-2.0.5\n",
      "Successfully installed torch-scatter-2.0.5\n",
      "\u001b[33mWARNING: You are using pip version 20.2.3; however, version 20.2.4 is available.\n",
      "You should consider upgrading via the '/home/ubuntu/anaconda3/envs/l5kit/bin/python -m pip install --upgrade pip' command.\u001b[0m\n",
      "Looking in links: https://pytorch-geometric.com/whl/torch-1.6.0.html\n",
      "Collecting torch-sparse==latest+cu102\n",
      "  Using cached https://pytorch-geometric.com/whl/torch-1.6.0/torch_sparse-latest%2Bcu102-cp36-cp36m-linux_x86_64.whl (23.0 MB)\n",
      "Requirement already satisfied: scipy in /home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages (from torch-sparse==latest+cu102) (1.1.0)\n",
      "\u001b[31mERROR: Error while checking for conflicts. Please file an issue on pip's issue tracker: https://github.com/pypa/pip/issues/new\n",
      "Traceback (most recent call last):\n",
      "  File \"/home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages/pip/_vendor/pkg_resources/__init__.py\", line 3021, in _dep_map\n",
      "    return self.__dep_map\n",
      "  File \"/home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages/pip/_vendor/pkg_resources/__init__.py\", line 2815, in __getattr__\n",
      "    raise AttributeError(attr)\n",
      "AttributeError: _DistInfoDistribution__dep_map\n",
      "\n",
      "During handling of the above exception, another exception occurred:\n",
      "\n",
      "Traceback (most recent call last):\n",
      "  File \"/home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages/pip/_vendor/pkg_resources/__init__.py\", line 3012, in _parsed_pkg_info\n",
      "    return self._pkg_info\n",
      "  File \"/home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages/pip/_vendor/pkg_resources/__init__.py\", line 2815, in __getattr__\n",
      "    raise AttributeError(attr)\n",
      "AttributeError: _pkg_info\n",
      "\n",
      "During handling of the above exception, another exception occurred:\n",
      "\n",
      "Traceback (most recent call last):\n",
      "  File \"/home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages/pip/_internal/commands/install.py\", line 520, in _determine_conflicts\n",
      "    return check_install_conflicts(to_install)\n",
      "  File \"/home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages/pip/_internal/operations/check.py\", line 108, in check_install_conflicts\n",
      "    package_set, _ = create_package_set_from_installed()\n",
      "  File \"/home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages/pip/_internal/operations/check.py\", line 50, in create_package_set_from_installed\n",
      "    package_set[name] = PackageDetails(dist.version, dist.requires())\n",
      "  File \"/home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages/pip/_vendor/pkg_resources/__init__.py\", line 2736, in requires\n",
      "    dm = self._dep_map\n",
      "  File \"/home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages/pip/_vendor/pkg_resources/__init__.py\", line 3023, in _dep_map\n",
      "    self.__dep_map = self._compute_dependencies()\n",
      "  File \"/home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages/pip/_vendor/pkg_resources/__init__.py\", line 3032, in _compute_dependencies\n",
      "    for req in self._parsed_pkg_info.get_all('Requires-Dist') or []:\n",
      "  File \"/home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages/pip/_vendor/pkg_resources/__init__.py\", line 3014, in _parsed_pkg_info\n",
      "    metadata = self.get_metadata(self.PKG_INFO)\n",
      "  File \"/home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages/pip/_vendor/pkg_resources/__init__.py\", line 1420, in get_metadata\n",
      "    value = self._get(path)\n",
      "  File \"/home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages/pip/_vendor/pkg_resources/__init__.py\", line 1616, in _get\n",
      "    with open(path, 'rb') as stream:\n",
      "FileNotFoundError: [Errno 2] No such file or directory: '/home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages/s3transfer-0.3.3.dist-info/METADATA'\u001b[0m\n",
      "Installing collected packages: torch-sparse\n",
      "  Attempting uninstall: torch-sparse\n",
      "    Found existing installation: torch-sparse 0.6.8\n",
      "    Uninstalling torch-sparse-0.6.8:\n",
      "      Successfully uninstalled torch-sparse-0.6.8\n",
      "Successfully installed torch-sparse-0.6.8\n",
      "\u001b[33mWARNING: You are using pip version 20.2.3; however, version 20.2.4 is available.\n",
      "You should consider upgrading via the '/home/ubuntu/anaconda3/envs/l5kit/bin/python -m pip install --upgrade pip' command.\u001b[0m\n",
      "Looking in links: https://pytorch-geometric.com/whl/torch-1.6.0.html\n",
      "Collecting torch-cluster==latest+cu102\n",
      "  Using cached https://pytorch-geometric.com/whl/torch-1.6.0/torch_cluster-latest%2Bcu102-cp36-cp36m-linux_x86_64.whl (20.4 MB)\n",
      "\u001b[31mERROR: Error while checking for conflicts. Please file an issue on pip's issue tracker: https://github.com/pypa/pip/issues/new\n",
      "Traceback (most recent call last):\n",
      "  File \"/home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages/pip/_vendor/pkg_resources/__init__.py\", line 3021, in _dep_map\n",
      "    return self.__dep_map\n",
      "  File \"/home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages/pip/_vendor/pkg_resources/__init__.py\", line 2815, in __getattr__\n",
      "    raise AttributeError(attr)\n",
      "AttributeError: _DistInfoDistribution__dep_map\n",
      "\n",
      "During handling of the above exception, another exception occurred:\n",
      "\n",
      "Traceback (most recent call last):\n",
      "  File \"/home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages/pip/_vendor/pkg_resources/__init__.py\", line 3012, in _parsed_pkg_info\n",
      "    return self._pkg_info\n",
      "  File \"/home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages/pip/_vendor/pkg_resources/__init__.py\", line 2815, in __getattr__\n",
      "    raise AttributeError(attr)\n",
      "AttributeError: _pkg_info\n",
      "\n",
      "During handling of the above exception, another exception occurred:\n",
      "\n",
      "Traceback (most recent call last):\n",
      "  File \"/home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages/pip/_internal/commands/install.py\", line 520, in _determine_conflicts\n",
      "    return check_install_conflicts(to_install)\n",
      "  File \"/home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages/pip/_internal/operations/check.py\", line 108, in check_install_conflicts\n",
      "    package_set, _ = create_package_set_from_installed()\n",
      "  File \"/home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages/pip/_internal/operations/check.py\", line 50, in create_package_set_from_installed\n",
      "    package_set[name] = PackageDetails(dist.version, dist.requires())\n",
      "  File \"/home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages/pip/_vendor/pkg_resources/__init__.py\", line 2736, in requires\n",
      "    dm = self._dep_map\n",
      "  File \"/home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages/pip/_vendor/pkg_resources/__init__.py\", line 3023, in _dep_map\n",
      "    self.__dep_map = self._compute_dependencies()\n",
      "  File \"/home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages/pip/_vendor/pkg_resources/__init__.py\", line 3032, in _compute_dependencies\n",
      "    for req in self._parsed_pkg_info.get_all('Requires-Dist') or []:\n",
      "  File \"/home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages/pip/_vendor/pkg_resources/__init__.py\", line 3014, in _parsed_pkg_info\n",
      "    metadata = self.get_metadata(self.PKG_INFO)\n",
      "  File \"/home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages/pip/_vendor/pkg_resources/__init__.py\", line 1420, in get_metadata\n",
      "    value = self._get(path)\n",
      "  File \"/home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages/pip/_vendor/pkg_resources/__init__.py\", line 1616, in _get\n",
      "    with open(path, 'rb') as stream:\n",
      "FileNotFoundError: [Errno 2] No such file or directory: '/home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages/s3transfer-0.3.3.dist-info/METADATA'\u001b[0m\n",
      "Installing collected packages: torch-cluster\n",
      "  Attempting uninstall: torch-cluster\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "    Found existing installation: torch-cluster 1.5.8\n",
      "    Uninstalling torch-cluster-1.5.8:\n",
      "      Successfully uninstalled torch-cluster-1.5.8\n",
      "Successfully installed torch-cluster-1.5.8\n",
      "\u001b[33mWARNING: You are using pip version 20.2.3; however, version 20.2.4 is available.\n",
      "You should consider upgrading via the '/home/ubuntu/anaconda3/envs/l5kit/bin/python -m pip install --upgrade pip' command.\u001b[0m\n",
      "Looking in links: https://pytorch-geometric.com/whl/torch-1.6.0.html\n",
      "Collecting torch-spline-conv==latest+cu102\n",
      "  Using cached https://pytorch-geometric.com/whl/torch-1.6.0/torch_spline_conv-latest%2Bcu102-cp36-cp36m-linux_x86_64.whl (6.1 MB)\n",
      "\u001b[31mERROR: Error while checking for conflicts. Please file an issue on pip's issue tracker: https://github.com/pypa/pip/issues/new\n",
      "Traceback (most recent call last):\n",
      "  File \"/home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages/pip/_vendor/pkg_resources/__init__.py\", line 3021, in _dep_map\n",
      "    return self.__dep_map\n",
      "  File \"/home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages/pip/_vendor/pkg_resources/__init__.py\", line 2815, in __getattr__\n",
      "    raise AttributeError(attr)\n",
      "AttributeError: _DistInfoDistribution__dep_map\n",
      "\n",
      "During handling of the above exception, another exception occurred:\n",
      "\n",
      "Traceback (most recent call last):\n",
      "  File \"/home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages/pip/_vendor/pkg_resources/__init__.py\", line 3012, in _parsed_pkg_info\n",
      "    return self._pkg_info\n",
      "  File \"/home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages/pip/_vendor/pkg_resources/__init__.py\", line 2815, in __getattr__\n",
      "    raise AttributeError(attr)\n",
      "AttributeError: _pkg_info\n",
      "\n",
      "During handling of the above exception, another exception occurred:\n",
      "\n",
      "Traceback (most recent call last):\n",
      "  File \"/home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages/pip/_internal/commands/install.py\", line 520, in _determine_conflicts\n",
      "    return check_install_conflicts(to_install)\n",
      "  File \"/home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages/pip/_internal/operations/check.py\", line 108, in check_install_conflicts\n",
      "    package_set, _ = create_package_set_from_installed()\n",
      "  File \"/home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages/pip/_internal/operations/check.py\", line 50, in create_package_set_from_installed\n",
      "    package_set[name] = PackageDetails(dist.version, dist.requires())\n",
      "  File \"/home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages/pip/_vendor/pkg_resources/__init__.py\", line 2736, in requires\n",
      "    dm = self._dep_map\n",
      "  File \"/home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages/pip/_vendor/pkg_resources/__init__.py\", line 3023, in _dep_map\n",
      "    self.__dep_map = self._compute_dependencies()\n",
      "  File \"/home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages/pip/_vendor/pkg_resources/__init__.py\", line 3032, in _compute_dependencies\n",
      "    for req in self._parsed_pkg_info.get_all('Requires-Dist') or []:\n",
      "  File \"/home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages/pip/_vendor/pkg_resources/__init__.py\", line 3014, in _parsed_pkg_info\n",
      "    metadata = self.get_metadata(self.PKG_INFO)\n",
      "  File \"/home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages/pip/_vendor/pkg_resources/__init__.py\", line 1420, in get_metadata\n",
      "    value = self._get(path)\n",
      "  File \"/home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages/pip/_vendor/pkg_resources/__init__.py\", line 1616, in _get\n",
      "    with open(path, 'rb') as stream:\n",
      "FileNotFoundError: [Errno 2] No such file or directory: '/home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages/s3transfer-0.3.3.dist-info/METADATA'\u001b[0m\n",
      "Installing collected packages: torch-spline-conv\n",
      "  Attempting uninstall: torch-spline-conv\n",
      "    Found existing installation: torch-spline-conv 1.2.0\n",
      "    Uninstalling torch-spline-conv-1.2.0:\n",
      "      Successfully uninstalled torch-spline-conv-1.2.0\n",
      "Successfully installed torch-spline-conv-1.2.0\n",
      "\u001b[33mWARNING: You are using pip version 20.2.3; however, version 20.2.4 is available.\n",
      "You should consider upgrading via the '/home/ubuntu/anaconda3/envs/l5kit/bin/python -m pip install --upgrade pip' command.\u001b[0m\n",
      "Requirement already satisfied: torch-geometric in /home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages (1.6.1)\n",
      "Requirement already satisfied: h5py in /home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages (from torch-geometric) (2.8.0)\n",
      "Requirement already satisfied: pandas in /home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages (from torch-geometric) (0.24.2)\n",
      "Requirement already satisfied: rdflib in /home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages (from torch-geometric) (5.0.0)\n",
      "Requirement already satisfied: googledrivedownloader in /home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages (from torch-geometric) (0.4)\n",
      "Requirement already satisfied: jinja2 in /home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages (from torch-geometric) (2.10)\n",
      "Requirement already satisfied: scikit-learn in /home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages (from torch-geometric) (0.20.3)\n",
      "Requirement already satisfied: numba in /home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages (from torch-geometric) (0.38.0)\n",
      "Requirement already satisfied: ase in /home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages (from torch-geometric) (3.20.1)\n",
      "Requirement already satisfied: networkx in /home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages (from torch-geometric) (2.1)\n",
      "Requirement already satisfied: torch in /home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages (from torch-geometric) (1.6.0)\n",
      "Requirement already satisfied: numpy in /home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages (from torch-geometric) (1.19.2)\n",
      "Requirement already satisfied: tqdm in /home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages (from torch-geometric) (4.50.2)\n",
      "Requirement already satisfied: scipy in /home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages (from torch-geometric) (1.1.0)\n",
      "Requirement already satisfied: requests in /home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages (from torch-geometric) (2.24.0)\n",
      "Requirement already satisfied: six in /home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages (from h5py->torch-geometric) (1.15.0)\n",
      "Requirement already satisfied: pytz>=2011k in /home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages (from pandas->torch-geometric) (2018.4)\n",
      "Requirement already satisfied: python-dateutil>=2.5.0 in /home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages (from pandas->torch-geometric) (2.7.3)\n",
      "Requirement already satisfied: pyparsing in /home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages (from rdflib->torch-geometric) (2.2.0)\n",
      "Requirement already satisfied: isodate in /home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages (from rdflib->torch-geometric) (0.6.0)\n",
      "Requirement already satisfied: MarkupSafe>=0.23 in /home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages (from jinja2->torch-geometric) (1.0)\n",
      "Requirement already satisfied: llvmlite>=0.23.0dev0 in /home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages (from numba->torch-geometric) (0.23.1)\n",
      "Requirement already satisfied: matplotlib>=2.0.0 in /home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages (from ase->torch-geometric) (3.0.0)\n",
      "Requirement already satisfied: decorator>=4.1.0 in /home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages (from networkx->torch-geometric) (4.3.0)\n",
      "Requirement already satisfied: future in /home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages (from torch->torch-geometric) (0.18.2)\n",
      "Requirement already satisfied: urllib3!=1.25.0,!=1.25.1,<1.26,>=1.21.1 in /home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages (from requests->torch-geometric) (1.23)\n",
      "Requirement already satisfied: chardet<4,>=3.0.2 in /home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages (from requests->torch-geometric) (3.0.4)\n",
      "Requirement already satisfied: idna<3,>=2.5 in /home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages (from requests->torch-geometric) (2.6)\n",
      "Requirement already satisfied: certifi>=2017.4.17 in /home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages (from requests->torch-geometric) (2019.11.28)\n",
      "Requirement already satisfied: cycler>=0.10 in /home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages (from matplotlib>=2.0.0->ase->torch-geometric) (0.10.0)\n",
      "Requirement already satisfied: kiwisolver>=1.0.1 in /home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages (from matplotlib>=2.0.0->ase->torch-geometric) (1.0.1)\n",
      "Requirement already satisfied: setuptools in /home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages (from kiwisolver>=1.0.1->matplotlib>=2.0.0->ase->torch-geometric) (50.3.0)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[31mERROR: Error while checking for conflicts. Please file an issue on pip's issue tracker: https://github.com/pypa/pip/issues/new\n",
      "Traceback (most recent call last):\n",
      "  File \"/home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages/pip/_vendor/pkg_resources/__init__.py\", line 3021, in _dep_map\n",
      "    return self.__dep_map\n",
      "  File \"/home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages/pip/_vendor/pkg_resources/__init__.py\", line 2815, in __getattr__\n",
      "    raise AttributeError(attr)\n",
      "AttributeError: _DistInfoDistribution__dep_map\n",
      "\n",
      "During handling of the above exception, another exception occurred:\n",
      "\n",
      "Traceback (most recent call last):\n",
      "  File \"/home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages/pip/_vendor/pkg_resources/__init__.py\", line 3012, in _parsed_pkg_info\n",
      "    return self._pkg_info\n",
      "  File \"/home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages/pip/_vendor/pkg_resources/__init__.py\", line 2815, in __getattr__\n",
      "    raise AttributeError(attr)\n",
      "AttributeError: _pkg_info\n",
      "\n",
      "During handling of the above exception, another exception occurred:\n",
      "\n",
      "Traceback (most recent call last):\n",
      "  File \"/home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages/pip/_internal/commands/install.py\", line 520, in _determine_conflicts\n",
      "    return check_install_conflicts(to_install)\n",
      "  File \"/home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages/pip/_internal/operations/check.py\", line 108, in check_install_conflicts\n",
      "    package_set, _ = create_package_set_from_installed()\n",
      "  File \"/home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages/pip/_internal/operations/check.py\", line 50, in create_package_set_from_installed\n",
      "    package_set[name] = PackageDetails(dist.version, dist.requires())\n",
      "  File \"/home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages/pip/_vendor/pkg_resources/__init__.py\", line 2736, in requires\n",
      "    dm = self._dep_map\n",
      "  File \"/home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages/pip/_vendor/pkg_resources/__init__.py\", line 3023, in _dep_map\n",
      "    self.__dep_map = self._compute_dependencies()\n",
      "  File \"/home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages/pip/_vendor/pkg_resources/__init__.py\", line 3032, in _compute_dependencies\n",
      "    for req in self._parsed_pkg_info.get_all('Requires-Dist') or []:\n",
      "  File \"/home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages/pip/_vendor/pkg_resources/__init__.py\", line 3014, in _parsed_pkg_info\n",
      "    metadata = self.get_metadata(self.PKG_INFO)\n",
      "  File \"/home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages/pip/_vendor/pkg_resources/__init__.py\", line 1420, in get_metadata\n",
      "    value = self._get(path)\n",
      "  File \"/home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages/pip/_vendor/pkg_resources/__init__.py\", line 1616, in _get\n",
      "    with open(path, 'rb') as stream:\n",
      "FileNotFoundError: [Errno 2] No such file or directory: '/home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages/s3transfer-0.3.3.dist-info/METADATA'\u001b[0m\n",
      "\u001b[33mWARNING: You are using pip version 20.2.3; however, version 20.2.4 is available.\n",
      "You should consider upgrading via the '/home/ubuntu/anaconda3/envs/l5kit/bin/python -m pip install --upgrade pip' command.\u001b[0m\n",
      "Requirement already satisfied: neptune-client in /home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages (0.4.124)\n",
      "Requirement already satisfied: click>=7.0 in /home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages (from neptune-client) (7.1.2)\n",
      "Requirement already satisfied: GitPython>=2.0.8 in /home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages (from neptune-client) (3.1.9)\n",
      "Requirement already satisfied: oauthlib>=2.1.0 in /home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages (from neptune-client) (3.1.0)\n",
      "Requirement already satisfied: PyJWT in /home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages (from neptune-client) (1.7.1)\n",
      "Requirement already satisfied: pandas in /home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages (from neptune-client) (0.24.2)\n",
      "Requirement already satisfied: bravado in /home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages (from neptune-client) (11.0.1)\n",
      "Requirement already satisfied: requests>=2.20.0 in /home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages (from neptune-client) (2.24.0)\n",
      "Requirement already satisfied: future>=0.17.1 in /home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages (from neptune-client) (0.18.2)\n",
      "Requirement already satisfied: websocket-client>=0.35.0 in /home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages (from neptune-client) (0.57.0)\n",
      "Requirement already satisfied: requests-oauthlib>=1.0.0 in /home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages (from neptune-client) (1.3.0)\n",
      "Requirement already satisfied: six>=1.12.0 in /home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages (from neptune-client) (1.15.0)\n",
      "Requirement already satisfied: Pillow>=1.1.6 in /home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages (from neptune-client) (6.2.1)\n",
      "Requirement already satisfied: packaging in /home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages (from neptune-client) (20.1)\n",
      "Requirement already satisfied: py3nvml in /home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages (from neptune-client) (0.2.6)\n",
      "Requirement already satisfied: gitdb<5,>=4.0.1 in /home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages (from GitPython>=2.0.8->neptune-client) (4.0.5)\n",
      "Requirement already satisfied: numpy>=1.12.0 in /home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages (from pandas->neptune-client) (1.19.2)\n",
      "Requirement already satisfied: python-dateutil>=2.5.0 in /home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages (from pandas->neptune-client) (2.7.3)\n",
      "Requirement already satisfied: pytz>=2011k in /home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages (from pandas->neptune-client) (2018.4)\n",
      "Requirement already satisfied: bravado-core>=5.16.1 in /home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages (from bravado->neptune-client) (5.17.0)\n",
      "Requirement already satisfied: monotonic in /home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages (from bravado->neptune-client) (1.5)\n",
      "Requirement already satisfied: simplejson in /home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages (from bravado->neptune-client) (3.17.2)\n",
      "Requirement already satisfied: typing-extensions in /home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages (from bravado->neptune-client) (3.7.4.2)\n",
      "Requirement already satisfied: pyyaml in /home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages (from bravado->neptune-client) (5.3.1)\n",
      "Requirement already satisfied: msgpack in /home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages (from bravado->neptune-client) (0.6.0)\n",
      "Requirement already satisfied: urllib3!=1.25.0,!=1.25.1,<1.26,>=1.21.1 in /home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages (from requests>=2.20.0->neptune-client) (1.23)\n",
      "Requirement already satisfied: certifi>=2017.4.17 in /home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages (from requests>=2.20.0->neptune-client) (2019.11.28)\n",
      "Requirement already satisfied: idna<3,>=2.5 in /home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages (from requests>=2.20.0->neptune-client) (2.6)\n",
      "Requirement already satisfied: chardet<4,>=3.0.2 in /home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages (from requests>=2.20.0->neptune-client) (3.0.4)\n",
      "Requirement already satisfied: pyparsing>=2.0.2 in /home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages (from packaging->neptune-client) (2.2.0)\n",
      "Requirement already satisfied: xmltodict in /home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages (from py3nvml->neptune-client) (0.12.0)\n",
      "Requirement already satisfied: smmap<4,>=3.0.1 in /home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages (from gitdb<5,>=4.0.1->GitPython>=2.0.8->neptune-client) (3.0.4)\n",
      "Requirement already satisfied: jsonschema[format]>=2.5.1 in /home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages (from bravado-core>=5.16.1->bravado->neptune-client) (2.6.0)\n",
      "Requirement already satisfied: swagger-spec-validator>=2.0.1 in /home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages (from bravado-core>=5.16.1->bravado->neptune-client) (2.7.3)\n",
      "Requirement already satisfied: jsonref in /home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages (from bravado-core>=5.16.1->bravado->neptune-client) (0.2)\n",
      "Requirement already satisfied: rfc3987 in /home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages (from jsonschema[format]>=2.5.1->bravado-core>=5.16.1->bravado->neptune-client) (1.3.8)\n",
      "Requirement already satisfied: strict-rfc3339 in /home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages (from jsonschema[format]>=2.5.1->bravado-core>=5.16.1->bravado->neptune-client) (0.7)\n",
      "Requirement already satisfied: webcolors in /home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages (from jsonschema[format]>=2.5.1->bravado-core>=5.16.1->bravado->neptune-client) (1.11.1)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[31mERROR: Error while checking for conflicts. Please file an issue on pip's issue tracker: https://github.com/pypa/pip/issues/new\n",
      "Traceback (most recent call last):\n",
      "  File \"/home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages/pip/_vendor/pkg_resources/__init__.py\", line 3021, in _dep_map\n",
      "    return self.__dep_map\n",
      "  File \"/home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages/pip/_vendor/pkg_resources/__init__.py\", line 2815, in __getattr__\n",
      "    raise AttributeError(attr)\n",
      "AttributeError: _DistInfoDistribution__dep_map\n",
      "\n",
      "During handling of the above exception, another exception occurred:\n",
      "\n",
      "Traceback (most recent call last):\n",
      "  File \"/home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages/pip/_vendor/pkg_resources/__init__.py\", line 3012, in _parsed_pkg_info\n",
      "    return self._pkg_info\n",
      "  File \"/home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages/pip/_vendor/pkg_resources/__init__.py\", line 2815, in __getattr__\n",
      "    raise AttributeError(attr)\n",
      "AttributeError: _pkg_info\n",
      "\n",
      "During handling of the above exception, another exception occurred:\n",
      "\n",
      "Traceback (most recent call last):\n",
      "  File \"/home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages/pip/_internal/commands/install.py\", line 520, in _determine_conflicts\n",
      "    return check_install_conflicts(to_install)\n",
      "  File \"/home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages/pip/_internal/operations/check.py\", line 108, in check_install_conflicts\n",
      "    package_set, _ = create_package_set_from_installed()\n",
      "  File \"/home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages/pip/_internal/operations/check.py\", line 50, in create_package_set_from_installed\n",
      "    package_set[name] = PackageDetails(dist.version, dist.requires())\n",
      "  File \"/home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages/pip/_vendor/pkg_resources/__init__.py\", line 2736, in requires\n",
      "    dm = self._dep_map\n",
      "  File \"/home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages/pip/_vendor/pkg_resources/__init__.py\", line 3023, in _dep_map\n",
      "    self.__dep_map = self._compute_dependencies()\n",
      "  File \"/home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages/pip/_vendor/pkg_resources/__init__.py\", line 3032, in _compute_dependencies\n",
      "    for req in self._parsed_pkg_info.get_all('Requires-Dist') or []:\n",
      "  File \"/home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages/pip/_vendor/pkg_resources/__init__.py\", line 3014, in _parsed_pkg_info\n",
      "    metadata = self.get_metadata(self.PKG_INFO)\n",
      "  File \"/home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages/pip/_vendor/pkg_resources/__init__.py\", line 1420, in get_metadata\n",
      "    value = self._get(path)\n",
      "  File \"/home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages/pip/_vendor/pkg_resources/__init__.py\", line 1616, in _get\n",
      "    with open(path, 'rb') as stream:\n",
      "FileNotFoundError: [Errno 2] No such file or directory: '/home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages/s3transfer-0.3.3.dist-info/METADATA'\u001b[0m\n",
      "\u001b[33mWARNING: You are using pip version 20.2.3; however, version 20.2.4 is available.\n",
      "You should consider upgrading via the '/home/ubuntu/anaconda3/envs/l5kit/bin/python -m pip install --upgrade pip' command.\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "# running the installation can take a couple of minutes\n",
    "\n",
    "try:\n",
    "    import zarr\n",
    "except ModuleNotFoundError:\n",
    "    !pip install --use-feature=2020-resolver zarr > /dev/null\n",
    "try:\n",
    "    import pytorch_lightning\n",
    "except ModuleNotFoundError:\n",
    "    !pip install  --use-feature=2020-resolver pytorch-lightning  > /dev/null\n",
    "    \n",
    "try:\n",
    "    import torch_geometric\n",
    "except ModuleNotFoundError:\n",
    "    !pip install torch-geometric\n",
    "\n",
    "!pip install torch-scatter==latest+cu102 -f https://pytorch-geometric.com/whl/torch-1.6.0.html\n",
    "!pip install torch-sparse==latest+cu102 -f https://pytorch-geometric.com/whl/torch-1.6.0.html\n",
    "!pip install torch-cluster==latest+cu102 -f https://pytorch-geometric.com/whl/torch-1.6.0.html\n",
    "!pip install torch-spline-conv==latest+cu102 -f https://pytorch-geometric.com/whl/torch-1.6.0.html\n",
    "!pip install torch-geometric\n",
    "\n",
    "!pip install neptune-client"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import zarr\n",
    "from abc import ABC\n",
    "from pathlib import Path\n",
    "import matplotlib.pyplot as plt\n",
    "from numcodecs import blosc\n",
    "import pandas as pd, numpy as np\n",
    "\n",
    "import bisect\n",
    "import itertools as it\n",
    "from tqdm.notebook import tqdm\n",
    "\n",
    "\n",
    "import torch\n",
    "from torch import nn, optim \n",
    "import torch.nn.functional as F\n",
    "from torch.autograd import Variable\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "\n",
    "from pytorch_lightning import Trainer\n",
    "from pytorch_lightning import LightningModule\n",
    "from pytorch_lightning.callbacks import ModelCheckpoint\n",
    "from pytorch_lightning.loggers import TensorBoardLogger\n",
    "\n",
    "import pickle, copy, re, time, datetime, random, warnings, gc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch import nn\n",
    "from torch.nn import Sequential as Seq, Linear as Lin, ReLU, BatchNorm1d as BN\n",
    "from torch_geometric.nn import SAGEConv\n",
    "from torch_geometric.nn import global_max_pool\n",
    "from torch_geometric.nn import PairNorm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "DATA_ROOT: /data/lyft-motion-prediction-autonomous-vehicles\\TRAIN_ZARR: scenes/train.zarr\n",
      "VALID_ZARR: scenes/validate.zarr\n"
     ]
    }
   ],
   "source": [
    "DATA_ROOT = Path(\"/data/lyft-motion-prediction-autonomous-vehicles\")\n",
    "#DATA_ROOT = Path(\"../input/lyft-motion-prediction-autonomous-vehicles\")\n",
    "TRAIN_ZARR = \"scenes/train.zarr\"\n",
    "VALID_ZARR = \"scenes/validate.zarr\"\n",
    "\n",
    "print(\"DATA_ROOT: {}\\TRAIN_ZARR: {}\\nVALID_ZARR: {}\".format(DATA_ROOT, TRAIN_ZARR, VALID_ZARR))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "DEVICE: cuda\n"
     ]
    }
   ],
   "source": [
    "HBACKWARD = 15\n",
    "HFORWARD = 50\n",
    "NFRAMES = 10\n",
    "FRAME_STRIDE = 15\n",
    "AGENT_FEATURE_DIM = 8\n",
    "MAX_AGENTS = 150\n",
    "\n",
    "DEVICE = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "print(\"DEVICE:\", DEVICE)\n",
    "\n",
    "NUM_WORKERS = 4\n",
    "BATCH_SIZE = 15\n",
    "EPOCHS=15\n",
    "GRADIENT_CLIP_VAL = 1.0\n",
    "LIMIT_VAL_BATCHES = 0.20\n",
    "ROOT = \"GNN\"\n",
    "\n",
    "Path(ROOT).mkdir(exist_ok=True, parents=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "TIME_FORMAT = r\"%Y-%m-%dT%H:%M:%S%Z\"\n",
    "def get_utc():\n",
    "    return datetime.datetime.now(datetime.timezone.utc).strftime(TIME_FORMAT)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "PERCEPTION_LABELS = [\n",
    "    \"PERCEPTION_LABEL_NOT_SET\",\n",
    "    \"PERCEPTION_LABEL_UNKNOWN\",\n",
    "    \"PERCEPTION_LABEL_DONTCARE\",\n",
    "    \"PERCEPTION_LABEL_CAR\",\n",
    "    \"PERCEPTION_LABEL_VAN\",\n",
    "    \"PERCEPTION_LABEL_TRAM\",\n",
    "    \"PERCEPTION_LABEL_BUS\",\n",
    "    \"PERCEPTION_LABEL_TRUCK\",\n",
    "    \"PERCEPTION_LABEL_EMERGENCY_VEHICLE\",\n",
    "    \"PERCEPTION_LABEL_OTHER_VEHICLE\",\n",
    "    \"PERCEPTION_LABEL_BICYCLE\",\n",
    "    \"PERCEPTION_LABEL_MOTORCYCLE\",\n",
    "    \"PERCEPTION_LABEL_CYCLIST\",\n",
    "    \"PERCEPTION_LABEL_MOTORCYCLIST\",\n",
    "    \"PERCEPTION_LABEL_PEDESTRIAN\",\n",
    "    \"PERCEPTION_LABEL_ANIMAL\",\n",
    "    \"AVRESEARCH_LABEL_DONTCARE\",\n",
    "]\n",
    "KEPT_PERCEPTION_LABELS = [\n",
    "    \"PERCEPTION_LABEL_UNKNOWN\",\n",
    "    \"PERCEPTION_LABEL_CAR\",\n",
    "    \"PERCEPTION_LABEL_CYCLIST\",\n",
    "    \"PERCEPTION_LABEL_PEDESTRIAN\",\n",
    "]\n",
    "KEPT_PERCEPTION_LABELS_DICT = {label:PERCEPTION_LABELS.index(label) for label in KEPT_PERCEPTION_LABELS}\n",
    "KEPT_PERCEPTION_KEYS = sorted(KEPT_PERCEPTION_LABELS_DICT.values())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "class LabelEncoder:\n",
    "    def  __init__(self, max_size=500, default_val=-1):\n",
    "        self.max_size = max_size\n",
    "        self.labels = {}\n",
    "        self.default_val = default_val\n",
    "\n",
    "    @property\n",
    "    def nlabels(self):\n",
    "        return len(self.labels)\n",
    "\n",
    "    def reset(self):\n",
    "        self.labels = {}\n",
    "\n",
    "    def partial_fit(self, keys):\n",
    "        nlabels = self.nlabels\n",
    "        available = self.max_size - nlabels\n",
    "\n",
    "        if available < 1:\n",
    "            return\n",
    "\n",
    "        keys = set(keys)\n",
    "        new_keys = list(keys - set(self.labels))\n",
    "\n",
    "        if not len(new_keys):\n",
    "            return\n",
    "        \n",
    "        self.labels.update(dict(zip(new_keys, range(nlabels, nlabels + available) )))\n",
    "    \n",
    "    def fit(self, keys):\n",
    "        self.reset()\n",
    "        self.partial_fit(keys)\n",
    "\n",
    "    def get(self, key):\n",
    "        return self.labels.get(key, self.default_val)\n",
    "    \n",
    "    def transform(self, keys):\n",
    "        return np.array(list(map(self.get, keys)))\n",
    "\n",
    "    def fit_transform(self, keys, partial=True):\n",
    "        self.partial_fit(keys) if partial else self.fit(keys)\n",
    "        return self.transform(keys)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "_cell_guid": "b1076dfc-b9ad-4769-8c92-a6c4dae69d19",
    "_uuid": "8f2839f25d086af736a60e9eeb907d3b93b6e0e5"
   },
   "outputs": [],
   "source": [
    "class CustomLyftDataset(Dataset):\n",
    "    feature_mins = np.array([-17.336, -27.137, 0. , 0., 0. , -3.142, -37.833, -65.583],\n",
    "    dtype=\"float32\")[None,None, None]\n",
    "\n",
    "    feature_maxs = np.array([17.114, 20.787, 42.854, 42.138,  7.079,  3.142, 29.802, 35.722],\n",
    "    dtype=\"float32\")[None,None, None]\n",
    "\n",
    "\n",
    "\n",
    "    def __init__(self, zdataset, scenes=None, nframes=10, frame_stride=15, hbackward=10, \n",
    "                 hforward=50, max_agents=150, agent_feature_dim=8):\n",
    "        \"\"\"\n",
    "        Custom Lyft dataset reader.\n",
    "        \n",
    "        Parmeters:\n",
    "        ----------\n",
    "        zdataset: zarr dataset\n",
    "            The root dataset, containing scenes, frames and agents\n",
    "            \n",
    "        nframes: int\n",
    "            Number of frames per scene\n",
    "            \n",
    "        frame_stride: int\n",
    "            The stride when reading the **nframes** frames from a scene\n",
    "            \n",
    "        hbackward: int\n",
    "            Number of backward frames from  current frame\n",
    "            \n",
    "        hforward: int\n",
    "            Number forward frames from current frame\n",
    "        \n",
    "        max_agents: int \n",
    "            Max number of agents to read for each target frame. Note that,\n",
    "            this also include the backward agents but not the forward ones.\n",
    "        \"\"\"\n",
    "        super().__init__()\n",
    "        self.zdataset = zdataset\n",
    "        self.scenes = scenes if scenes is not None else []\n",
    "        self.nframes = nframes\n",
    "        self.frame_stride = frame_stride\n",
    "        self.hbackward = hbackward\n",
    "        self.hforward = hforward\n",
    "        self.max_agents = max_agents\n",
    "\n",
    "        self.nread_frames = (nframes-1)*frame_stride + hbackward + hforward\n",
    "\n",
    "        self.frame_fields = ['timestamp', 'agent_index_interval']\n",
    "\n",
    "        self.agent_feature_dim = agent_feature_dim\n",
    "\n",
    "        self.filter_scenes()\n",
    "      \n",
    "    def __len__(self):\n",
    "        return len(self.scenes)\n",
    "\n",
    "    def filter_scenes(self):\n",
    "        self.scenes = [scene for scene in self.scenes if self.get_nframes(scene) > self.nread_frames]\n",
    "\n",
    "\n",
    "    def __getitem__(self, index):\n",
    "        return self.read_frames(scene=self.scenes[index])\n",
    "\n",
    "    def get_nframes(self, scene, start=None):\n",
    "        frame_start = scene[\"frame_index_interval\"][0]\n",
    "        frame_end = scene[\"frame_index_interval\"][1]\n",
    "        nframes = (frame_end - frame_start) if start is None else ( frame_end - max(frame_start, start) )\n",
    "        return nframes\n",
    "\n",
    "\n",
    "    def _read_frames(self, scene, start=None):\n",
    "        nframes = self.get_nframes(scene, start=start)\n",
    "        assert nframes >= self.nread_frames\n",
    "\n",
    "        frame_start = scene[\"frame_index_interval\"][0]\n",
    "\n",
    "        start = start or frame_start + np.random.choice(nframes-self.nread_frames)\n",
    "        frames = self.zdataset.frames.get_basic_selection(\n",
    "            selection=slice(start, start+self.nread_frames),\n",
    "            fields=self.frame_fields,\n",
    "            )\n",
    "        return frames\n",
    "    \n",
    "\n",
    "    def parse_frame(self, frame):\n",
    "        return frame\n",
    "\n",
    "    def parse_agent(self, agent):\n",
    "        return agent\n",
    "\n",
    "    def read_frames(self, scene, start=None,  white_tracks=None, encoder=False):\n",
    "        white_tracks = white_tracks or []\n",
    "        frames = self._read_frames(scene=scene, start=start)\n",
    "\n",
    "        agent_start = frames[0][\"agent_index_interval\"][0]\n",
    "        agent_end = frames[-1][\"agent_index_interval\"][1]\n",
    "\n",
    "        agents = self.zdataset.agents[agent_start:agent_end]\n",
    "\n",
    "\n",
    "        X = np.zeros((self.nframes, self.max_agents, self.hbackward, self.agent_feature_dim), dtype=np.float32)\n",
    "        target = np.zeros((self.nframes, self.max_agents, self.hforward, 2),  dtype=np.float32)\n",
    "        target_availability = np.zeros((self.nframes, self.max_agents, self.hforward), dtype=np.uint8)\n",
    "        X_availability = np.zeros((self.nframes, self.max_agents, self.hbackward), dtype=np.uint8)\n",
    "\n",
    "        for f in range(self.nframes):\n",
    "            backward_frame_start = f*self.frame_stride\n",
    "            forward_frame_start = f*self.frame_stride+self.hbackward\n",
    "            backward_frames = frames[backward_frame_start:backward_frame_start+self.hbackward]\n",
    "            forward_frames = frames[forward_frame_start:forward_frame_start+self.hforward]\n",
    "\n",
    "            backward_agent_start = backward_frames[-1][\"agent_index_interval\"][0] - agent_start\n",
    "            backward_agent_end = backward_frames[-1][\"agent_index_interval\"][1] - agent_start\n",
    "\n",
    "            backward_agents = agents[backward_agent_start:backward_agent_end]\n",
    "\n",
    "            le = LabelEncoder(max_size=self.max_agents)\n",
    "            le.fit(white_tracks)\n",
    "            le.partial_fit(backward_agents[\"track_id\"])\n",
    "\n",
    "            for iframe, frame in enumerate(backward_frames):\n",
    "                backward_agent_start = frame[\"agent_index_interval\"][0] - agent_start\n",
    "                backward_agent_end = frame[\"agent_index_interval\"][1] - agent_start\n",
    "\n",
    "                backward_agents = agents[backward_agent_start:backward_agent_end]\n",
    "\n",
    "                track_ids = le.transform(backward_agents[\"track_id\"])\n",
    "                mask = (track_ids != le.default_val)\n",
    "                mask_agents = backward_agents[mask]\n",
    "                mask_ids = track_ids[mask]\n",
    "                X[f, mask_ids, iframe, :2] = mask_agents[\"centroid\"]\n",
    "                X[f, mask_ids, iframe, 2:5] = mask_agents[\"extent\"]\n",
    "                X[f, mask_ids, iframe, 5] = mask_agents[\"yaw\"]\n",
    "                X[f, mask_ids, iframe, 6:8] = mask_agents[\"velocity\"]\n",
    "\n",
    "                X_availability[f, mask_ids, iframe] = 1\n",
    "\n",
    "            \n",
    "            for iframe, frame in enumerate(forward_frames):\n",
    "                forward_agent_start = frame[\"agent_index_interval\"][0] - agent_start\n",
    "                forward_agent_end = frame[\"agent_index_interval\"][1] - agent_start\n",
    "\n",
    "                forward_agents = agents[forward_agent_start:forward_agent_end]\n",
    "\n",
    "                track_ids = le.transform(forward_agents[\"track_id\"])\n",
    "                mask = track_ids != le.default_val\n",
    "\n",
    "                target[f, track_ids[mask], iframe] = forward_agents[mask][\"centroid\"]\n",
    "                target_availability[f, track_ids[mask], iframe] = 1\n",
    "\n",
    "        target -= X[:,:,[-1], :2]\n",
    "        target *= target_availability[:,:,:,None]\n",
    "        X[:,:,:, :2] -= X[:,:,[-1], :2]\n",
    "        X *= X_availability[:,:,:,None]\n",
    "        X -= self.feature_mins\n",
    "        X /= (self.feature_maxs - self.feature_mins)\n",
    "\n",
    "        if encoder:\n",
    "            return X, target, target_availability, le\n",
    "        return X, target, target_availability"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Loss functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "def collate(x):\n",
    "    x = map(np.concatenate, zip(*x))\n",
    "    x = map(torch.from_numpy, x)\n",
    "    return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "def shapefy( xy_pred, xy, xy_av):\n",
    "    NDIM = 3\n",
    "    xy_pred = xy_pred.view(-1, HFORWARD, NDIM, 2)\n",
    "    xy = xy.view(-1, HFORWARD, 2)[:,:,None]\n",
    "    xy_av = xy_av.view(-1, HFORWARD)[:,:,None]\n",
    "    return xy_pred, xy,xy_av\n",
    "\n",
    "def LyftLoss(c, xy_pred, xy, xy_av):\n",
    "    c = c.view(-1,c.shape[-1])\n",
    "    xy_pred, xy, xy_av  = shapefy(xy_pred, xy, xy_av)\n",
    "    print(xy.shape)\n",
    "    print(xy_pred.shape)\n",
    "    \n",
    "    #pos_x, pos_y = xy.detach().cpu().numpy()[:10, : , 0, :].reshape((10, 50, ))\n",
    "    \n",
    "    c = torch.softmax(c, dim=1)\n",
    "    \n",
    "    l = torch.sum(torch.mean(torch.square(xy_pred-xy), dim=3)*xy_av, dim=1)\n",
    "    \n",
    "    # The LogSumExp trick for better numerical stability\n",
    "    # https://en.wikipedia.org/wiki/LogSumExp\n",
    "    m = l.min(dim=1).values\n",
    "    l = torch.exp(m[:, None]-l)\n",
    "    \n",
    "    l = m - torch.log(torch.sum(l*c, dim=1))\n",
    "    denom = xy_av.max(2).values.max(1).values\n",
    "    l = torch.sum(l*denom)/denom.sum()\n",
    "    return 3*l # I found that my loss is usually 3 times smaller than the LB score\n",
    "\n",
    "\n",
    "def MSE(xy_pred, xy, xy_av):\n",
    "    xy_pred, xy, xy_av = shapefy(xy_pred, xy, xy_av)\n",
    "    return 9*torch.mean(torch.sum(torch.mean(torch.square(xy_pred-xy), 3)*xy_av, dim=1))\n",
    "\n",
    "def MAE(xy_pred, xy, xy_av):\n",
    "    xy_pred, xy, xy_av = shapefy(xy_pred, xy, xy_av)\n",
    "    return 9*torch.mean(torch.sum(torch.mean(torch.abs(xy_pred-xy), 3)*xy_av, dim=1))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Pytorch-Lyghtining Modules"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "class BaseNet(LightningModule):   \n",
    "    def __init__(self, batch_size=32, lr=5e-4, weight_decay=1e-8, num_workers=0, \n",
    "                 criterion=LyftLoss, data_root=DATA_ROOT,  epochs=1):\n",
    "        super().__init__()\n",
    "\n",
    "       \n",
    "        self.save_hyperparameters(\n",
    "            dict(\n",
    "                HBACKWARD = HBACKWARD,\n",
    "                HFORWARD = HFORWARD,\n",
    "                NFRAMES = NFRAMES,\n",
    "                FRAME_STRIDE = FRAME_STRIDE,\n",
    "                AGENT_FEATURE_DIM = AGENT_FEATURE_DIM,\n",
    "                MAX_AGENTS = MAX_AGENTS,\n",
    "                TRAIN_ZARR = TRAIN_ZARR,\n",
    "                VALID_ZARR = VALID_ZARR,\n",
    "                batch_size = batch_size,\n",
    "                lr=lr,\n",
    "                weight_decay=weight_decay,\n",
    "                num_workers=num_workers,\n",
    "                criterion=criterion,\n",
    "                epochs=epochs,\n",
    "            )\n",
    "        )\n",
    "        \n",
    "        self._train_data = None\n",
    "        self._collate_fn = None\n",
    "        self._train_loader = None\n",
    "\n",
    "        self.batch_size = batch_size\n",
    "        self.num_workers = num_workers\n",
    "        \n",
    "        \n",
    "        self.lr = lr\n",
    "        self.epochs=epochs\n",
    "        \n",
    "        self.weight_decay = weight_decay\n",
    "        self.criterion = criterion\n",
    "        \n",
    "        self.data_root = data_root\n",
    "    \n",
    "\n",
    "    def train_dataloader(self):\n",
    "        z = zarr.open(self.data_root.joinpath(TRAIN_ZARR).as_posix(), \"r\")\n",
    "        scenes = z.scenes.get_basic_selection(slice(None), fields= [\"frame_index_interval\"])\n",
    "        train_data = CustomLyftDataset(\n",
    "                    z, \n",
    "                    scenes = scenes,\n",
    "                    nframes=NFRAMES,\n",
    "                    frame_stride=FRAME_STRIDE,\n",
    "                    hbackward=HBACKWARD,\n",
    "                    hforward=HFORWARD,\n",
    "                    max_agents=MAX_AGENTS,\n",
    "                    agent_feature_dim=AGENT_FEATURE_DIM,\n",
    "                )\n",
    "        \n",
    "        train_loader = DataLoader(train_data, batch_size = self.batch_size,collate_fn=collate,\n",
    "                                pin_memory=True, num_workers = self.num_workers, shuffle=True)\n",
    "        self._train_data = train_data\n",
    "        self._train_loader = train_loader\n",
    "        \n",
    "        return train_loader\n",
    "\n",
    "    def val_dataloader(self):\n",
    "        z = zarr.open(self.data_root.joinpath(VALID_ZARR).as_posix(), \"r\")\n",
    "        scenes = z.scenes.get_basic_selection(slice(None), fields=[\"frame_index_interval\"])\n",
    "        val_data = CustomLyftDataset(\n",
    "                    z, \n",
    "                    scenes = scenes,\n",
    "                    nframes=NFRAMES,\n",
    "                    frame_stride=FRAME_STRIDE,\n",
    "                    hbackward=HBACKWARD,\n",
    "                    hforward=HFORWARD,\n",
    "                    max_agents=MAX_AGENTS,\n",
    "                    agent_feature_dim=AGENT_FEATURE_DIM,\n",
    "                )\n",
    "        \n",
    "        val_loader = DataLoader(val_data, batch_size = self.batch_size, collate_fn=collate,\n",
    "                                pin_memory=True, num_workers = self.num_workers, shuffle=True)\n",
    "        self._val_data = val_data\n",
    "        self._val_loader = val_loader\n",
    "        return val_loader\n",
    "\n",
    "    def validation_epoch_end(self, outputs):\n",
    "        avg_loss = torch.mean(torch.tensor([x['val_loss'] for x in outputs]))\n",
    "        avg_mse = torch.mean(torch.tensor([x['val_mse'] for x in outputs]))\n",
    "        avg_mae = torch.mean(torch.tensor([x['val_mae'] for x in outputs]))\n",
    "        \n",
    "        tensorboard_logs = {'val_loss': avg_loss, \"val_rmse\": torch.sqrt(avg_mse), \"val_mae\": avg_mae}\n",
    "\n",
    "        torch.cuda.empty_cache()\n",
    "        gc.collect()\n",
    "\n",
    "        return {\n",
    "            'val_loss': avg_loss,\n",
    "            'log': tensorboard_logs,\n",
    "            \"progress_bar\": {\"val_ll\": tensorboard_logs[\"val_loss\"], \"val_rmse\": tensorboard_logs[\"val_rmse\"]}\n",
    "        }\n",
    "\n",
    "    \n",
    "    def configure_optimizers(self):\n",
    "        optimizer =  optim.Adam(self.parameters(), lr= self.lr, betas= (0.9,0.999), \n",
    "                          weight_decay= self.weight_decay, amsgrad=False)\n",
    "        \n",
    "        scheduler = optim.lr_scheduler.CosineAnnealingLR(\n",
    "            optimizer,\n",
    "            T_max=self.epochs,\n",
    "            eta_min=1e-5,\n",
    "        )\n",
    "        return [optimizer], [scheduler]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "_cell_guid": "79c7e3d0-c299-4dcb-8224-4455121ee9b0",
    "_uuid": "d629ff2d2480ee46fbb7e2d37f6b5fab8052498a",
    "collapsed": true
   },
   "source": [
    "# GraphNet"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "class GraphNet(BaseNet):\n",
    "    \"\"\"\n",
    "    Implement GraphNet model\n",
    "    \"\"\"\n",
    "    def __init__(self, **kwargs):\n",
    "        super().__init__(**kwargs)\n",
    "        hidden_size = 256\n",
    "        nb_layers = 7\n",
    "        nb_input_channels = 120\n",
    "        nb_output_channels = 300\n",
    "        self.nb_layers = nb_layers\n",
    "        self.hidden_size = hidden_size\n",
    "        self.conv_in_1 = SAGEConv(nb_input_channels, hidden_size)\n",
    "        self.conv_in_2 = SAGEConv(hidden_size, 2*hidden_size)\n",
    "        self.conv_in_3 = SAGEConv(2*hidden_size, 2*hidden_size)\n",
    "        \n",
    "        self.c_net = nn.Sequential(\n",
    "            nn.Linear(2*hidden_size, hidden_size), \n",
    "            nn.ReLU(),\n",
    "            nn.Linear(hidden_size, 3)\n",
    "        )\n",
    "        \n",
    "        \n",
    "        self.conv_out_1 = SAGEConv(2*hidden_size, 2*hidden_size)\n",
    "        self.conv_out_2 = SAGEConv(2*hidden_size, hidden_size)\n",
    "        self.conv_out_3 = SAGEConv(hidden_size, hidden_size)\n",
    "        self.conv_out_4 = SAGEConv(hidden_size, nb_output_channels)\n",
    "        self.relu = nn.ReLU(inplace=True)\n",
    "        \n",
    "\n",
    "    def forward(self, x):\n",
    "\n",
    "        bsize, npoints, hb, nf = x.shape \n",
    "\n",
    "        # Merge time with features\n",
    "        x = x.reshape(bsize, npoints, hb*nf)\n",
    "        \n",
    "        L1, L2 = [], []\n",
    "        for i in range(MAX_AGENTS):\n",
    "            for j in range(MAX_AGENTS):\n",
    "                L1.append(i)\n",
    "                L2.append(j)\n",
    "                \n",
    "        edge_index = torch.tensor([L1, L2], dtype=torch.long).cuda()\n",
    "        \n",
    "        out = self.relu(self.conv_in_1(x, edge_index))\n",
    "        out = self.relu(self.conv_in_2(out, edge_index))\n",
    "        out = self.relu(self.conv_in_3(out, edge_index))\n",
    "        c = self.c_net(out)\n",
    "\n",
    "        out = self.relu(self.conv_out_1(out, edge_index))\n",
    "        out = self.relu(self.conv_out_2(out, edge_index))\n",
    "        out = self.relu(self.conv_out_3(out, edge_index))\n",
    "        out = self.conv_out_4(out, edge_index)\n",
    "        return c, out\n",
    "\n",
    "    def get_nb_trainable_params(self):\n",
    "        \"\"\"\n",
    "        Return the number of trainable parameters\n",
    "        \"\"\"\n",
    "        model_parameters = filter(lambda p: p.requires_grad, self.parameters())\n",
    "        return sum([np.prod(p.size()) for p in model_parameters])\n",
    "\n",
    "    def describe(self):\n",
    "        s = f\"nb_layers: {self.nb_layers}; hidden size: {self.hidden_size}\"\n",
    "        return s\n",
    "    \n",
    "    def training_step(self, batch, batch_idx):\n",
    "        x, y, y_av = [b.to(DEVICE) for b in batch]\n",
    "        c, preds = self(x)\n",
    "        loss = self.criterion(c,preds,y, y_av)\n",
    "        \n",
    "        with torch.no_grad():\n",
    "            logs = {\n",
    "                'loss': loss,\n",
    "                \"mse\": MSE(preds, y, y_av),\n",
    "                \"mae\": MAE(preds, y, y_av),\n",
    "            }\n",
    "        return {'loss': loss, 'log': logs, \"progress_bar\": {\"rmse\":torch.sqrt(logs[\"mse\"]) }}\n",
    "    \n",
    "    @torch.no_grad()\n",
    "    def validation_step(self, batch, batch_idx):\n",
    "        x, y, y_av =  [b.to(DEVICE) for b in batch]\n",
    "        c,preds = self(x)\n",
    "        loss = self.criterion(c, preds, y, y_av)\n",
    "        \n",
    "        val_logs = {\n",
    "            'val_loss': loss,\n",
    "            \"val_mse\": MSE(preds, y, y_av),\n",
    "            \"val_mae\": MAE(preds, y, y_av),\n",
    "        }\n",
    "        \n",
    "        return val_logs\n",
    "    \n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "learning_rate = 5e-4\n",
    "weight_decay = 5e-7 "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.backends.cudnn.benchmark =  True\n",
    "\n",
    "model = GraphNet(batch_size=BATCH_SIZE, lr= learning_rate, weight_decay=weight_decay, num_workers=NUM_WORKERS)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "GraphNet(\n",
       "  (conv_in_1): SAGEConv(120, 256)\n",
       "  (conv_in_2): SAGEConv(256, 512)\n",
       "  (conv_in_3): SAGEConv(512, 512)\n",
       "  (c_net): Sequential(\n",
       "    (0): Linear(in_features=512, out_features=256, bias=True)\n",
       "    (1): ReLU()\n",
       "    (2): Linear(in_features=256, out_features=3, bias=True)\n",
       "  )\n",
       "  (conv_out_1): SAGEConv(512, 512)\n",
       "  (conv_out_2): SAGEConv(512, 256)\n",
       "  (conv_out_3): SAGEConv(256, 256)\n",
       "  (conv_out_4): SAGEConv(256, 300)\n",
       "  (relu): ReLU(inplace=True)\n",
       ")"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Load trained model\n",
    "model.load_state_dict(torch.load('GraphNet.pt'))\n",
    "model.cuda()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING: There is a new version of neptune-client 0.4.125 (installed: 0.4.124).\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "https://ui.neptune.ai/hvergnes/KaggleGNN/e/KAG2-15\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "NeptuneLogger will work in online mode\n"
     ]
    }
   ],
   "source": [
    "from pytorch_lightning.loggers.neptune import NeptuneLogger\n",
    "\n",
    "neptune_logger = NeptuneLogger(\n",
    "                api_key=\"eyJhcGlfYWRkcmVzcyI6Imh0dHBzOi8vdWkubmVwdHVuZS5haSIsImFwaV91cmwiOiJodHRwczovL3VpLm5lcHR1bmUuYWkiLCJhcGlfa2V5IjoiMDk5YjVmYzYtNTU0My00MzhkLWJiYTAtMGM4ZGVhZmEyMTZiIn0=\",\n",
    "                project_name='hvergnes/KaggleGNN',\n",
    "                params={'epoch_nr': EPOCHS, 'learning_rate': learning_rate, 'batch_size': BATCH_SIZE},  # your hyperparameters, immutable\n",
    "                tags=['GNN'],  # tags\n",
    "                )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "GPU available: True, used: True\n",
      "TPU available: False, using: 0 TPU cores\n",
      "LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]\n",
      "\n",
      "  | Name       | Type       | Params\n",
      "------------------------------------------\n",
      "0 | conv_in_1  | SAGEConv   | 61 K  \n",
      "1 | conv_in_2  | SAGEConv   | 262 K \n",
      "2 | conv_in_3  | SAGEConv   | 524 K \n",
      "3 | c_net      | Sequential | 132 K \n",
      "4 | conv_out_1 | SAGEConv   | 524 K \n",
      "5 | conv_out_2 | SAGEConv   | 262 K \n",
      "6 | conv_out_3 | SAGEConv   | 131 K \n",
      "7 | conv_out_4 | SAGEConv   | 153 K \n",
      "8 | relu       | ReLU       | 0     \n",
      "/home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages/pytorch_lightning/utilities/distributed.py:45: UserWarning: Your val_dataloader has `shuffle=True`, it is best practice to turn this off for validation and test dataloaders.\n",
      "  warnings.warn(*args, **kwargs)\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(HTML(value='Validation sanity check'), FloatProgress(value=1.0, bar_style='info', layout=Layout…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages/pytorch_lightning/utilities/distributed.py:45: UserWarning: The validation_epoch_end should not return anything as of 9.1.to log, use self.log(...) or self.write(...) directly in the LightningModule\n",
      "  warnings.warn(*args, **kwargs)\n",
      "/home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages/pytorch_lightning/utilities/distributed.py:45: UserWarning: The {log:dict keyword} was deprecated in 0.9.1 and will be removed in 1.0.0\n",
      "Please use self.log(...) inside the lightningModule instead.\n",
      "\n",
      "# log on a step or aggregate epoch metric to the logger and/or progress bar\n",
      "# (inside LightningModule)\n",
      "self.log('train_loss', loss, on_step=True, on_epoch=True, prog_bar=True)\n",
      "  warnings.warn(*args, **kwargs)\n",
      "/home/ubuntu/anaconda3/envs/l5kit/lib/python3.6/site-packages/pytorch_lightning/utilities/distributed.py:45: UserWarning: The {progress_bar:dict keyword} was deprecated in 0.9.1 and will be removed in 1.0.0\n",
      "Please use self.log(...) inside the lightningModule instead.\n",
      "\n",
      "# log on a step or aggregate epoch metric to the logger and/or progress bar\n",
      "# (inside LightningModule)\n",
      "self.log('train_loss', loss, on_step=True, on_epoch=True, prog_bar=True)\n",
      "  warnings.warn(*args, **kwargs)\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "0f82145ccf40420a88b8287b0583dd65",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(HTML(value='Training'), FloatProgress(value=1.0, bar_style='info', layout=Layout(flex='2'), max…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "trainer = Trainer(\n",
    "    max_epochs=EPOCHS,\n",
    "    logger=neptune_logger,\n",
    "    gradient_clip_val=GRADIENT_CLIP_VAL,\n",
    "    limit_val_batches=LIMIT_VAL_BATCHES,\n",
    "    gpus=1\n",
    ")\n",
    "\n",
    "trainer.fit(model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Special Mention\n",
    "\n",
    "This work was originaly inspired by the different kernels of KKiller during this competition.\n",
    "We used some of his code as a building block for our model. Please go upvote his kernel here:\n",
    "https://www.kaggle.com/kneroma/training-motion-prediction-with-pointnet"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div style=\"text-align: center\"><strong>Thanks</strong></div>"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Environment (conda_l5kit)",
   "language": "python",
   "name": "conda_l5kit"
  },
  "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.6.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}