{"cells":[{"metadata":{},"cell_type":"markdown","source":"# Motion Prediction with PointNet"},{"metadata":{},"cell_type":"markdown","source":"# About"},{"metadata":{},"cell_type":"markdown","source":"In this work, I will be using the [pointnet architecture](https://github.com/charlesq34/pointnet) to predict autonomous vehicle motions. This work is completely different from the main public kernels and exposes another way of dealing with motion prediction.\n\nOne of the edges of the pointnet architecture is that it works directly on the agents set and, therefore, no rastarization is required. Hence, we can can get a decent accuracy  while training only on the small train set and just for 5 to 7 hours. Hence, this model could be run even if one doesn't have huge GPU capacities. But, if you do have enough GPU, you can increase the model size and train on the full train set."},{"metadata":{},"cell_type":"markdown","source":"### This kernel is for training only, the one for inference could [be found here](https://www.kaggle.com/kneroma/inference-motion-prediction-with-pointnet)."},{"metadata":{},"cell_type":"markdown","source":"The PointNet model uses a combination of feed forward (conv1d) models along with somme transformation matrices and symmetric functions (max pooling) to deal with unordered set of points (agents). For more reading, please visit the [PointNet project page](https://stanford.edu/~rqi/pointnet/) or [the original paper](https://arxiv.org/pdf/1612.00593.pdf). My code is inspired from [this github repos](https://github.com/charlesq34/pointnet)."},{"metadata":{},"cell_type":"markdown","source":"![](https://stanford.edu/~rqi/pointnet/images/pointnet.jpg)"},{"metadata":{"trusted":true},"cell_type":"code","source":"","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"> PointNet can be applied for classification and image segmentation. Since we will be needing an output for each point (agent), only the segmentation version will be suitable.\n\n\n![PointNet architecture](https://github.com/charlesq34/pointnet/blob/master/doc/teaser.png?raw=true)"},{"metadata":{"trusted":true},"cell_type":"code","source":"","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Training summary"},{"metadata":{},"cell_type":"markdown","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 : the **great pytorch-lightning** trainer\n* checkpoints : 5 best models based on val_loss"},{"metadata":{"trusted":true},"cell_type":"code","source":"","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"try:\n    import zarr\nexcept ModuleNotFoundError:\n    !pip install --use-feature=2020-resolver zarr > /dev/null\ntry:\n    import pytorch_lightning\nexcept ModuleNotFoundError:\n    !pip install  --use-feature=2020-resolver pytorch-lightning  > /dev/null","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import zarr\nfrom abc import ABC\nfrom pathlib import Path\nfrom numcodecs import blosc\nimport pandas as pd, numpy as np\n\nimport bisect\nimport itertools as it\nfrom tqdm.notebook import tqdm\n\n\nimport torch\nfrom torch import nn, optim \nimport torch.nn.functional as F\nfrom torch.autograd import Variable\nfrom torch.utils.data import Dataset, DataLoader\n\nfrom pytorch_lightning import Trainer\nfrom pytorch_lightning import LightningModule\nfrom pytorch_lightning.callbacks import ModelCheckpoint\nfrom pytorch_lightning.loggers import TensorBoardLogger\n\nimport pickle, copy, re,time, datetime, random, warnings, gc","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# blosc.set_nthreads(6)  \n# blosc.use_threads = True","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"DATA_ROOT = Path(\"../input/lyft-motion-prediction-autonomous-vehicles\")\nTRAIN_ZARR = \"scenes/train.zarr\"\nVALID_ZARR = \"scenes/validate.zarr\"\n\nprint(\"DATA_ROOT: {}\\TRAIN_ZARR: {}\\nVALID_ZARR: {}\".format(DATA_ROOT, TRAIN_ZARR, VALID_ZARR))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"HBACKWARD = 15\nHFORWARD = 50\nNFRAMES = 10\nFRAME_STRIDE = 15\nAGENT_FEATURE_DIM = 8\nMAX_AGENTS = 150","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"DEVICE = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\nprint(\"DEVICE:\", DEVICE)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"NUM_WORKERS = 2\nBATCH_SIZE = 150\nEPOCHS=24\nGRADIENT_CLIP_VAL = 1.0\nLIMIT_VAL_BATCHES = 0.20\nROOT = \"pointnet\"\n\nPath(ROOT).mkdir(exist_ok=True, parents=True)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"TIME_FORMAT = r\"%Y-%m-%dT%H:%M:%S%Z\"\ndef get_utc():\n    return datetime.datetime.now(datetime.timezone.utc).strftime(TIME_FORMAT)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Dataset"},{"metadata":{"trusted":true},"cell_type":"code","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]\nKEPT_PERCEPTION_LABELS = [\n    \"PERCEPTION_LABEL_UNKNOWN\",\n    \"PERCEPTION_LABEL_CAR\",\n    \"PERCEPTION_LABEL_CYCLIST\",\n    \"PERCEPTION_LABEL_PEDESTRIAN\",\n]\nKEPT_PERCEPTION_LABELS_DICT = {label:PERCEPTION_LABELS.index(label) for label in KEPT_PERCEPTION_LABELS}\nKEPT_PERCEPTION_KEYS = sorted(KEPT_PERCEPTION_LABELS_DICT.values())","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","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)","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","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","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Loss functions"},{"metadata":{"trusted":true},"cell_type":"code","source":"def collate(x):\n    x = map(np.concatenate, zip(*x))\n    x = map(torch.from_numpy, x)\n    return x","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","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\ndef 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    \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\ndef 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\ndef 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","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"markdown","source":"# Pytorch-Lyghtining Modules"},{"metadata":{"trusted":true},"cell_type":"code","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]","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","collapsed":true,"_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":false},"cell_type":"markdown","source":"# PointNet"},{"metadata":{"trusted":true},"cell_type":"code","source":"class STNkd(nn.Module):\n    def __init__(self,  k=64):\n        super(STNkd, self).__init__()\n        \n        self.conv = nn.Sequential(\n            nn.Conv1d(k, 256, kernel_size=1), nn.ReLU(),\n            nn.Conv1d(256, 256, kernel_size=1), nn.ReLU(),\n            nn.Conv1d(256, 512, kernel_size=1), nn.ReLU(),\n        )\n        \n        self.fc = nn.Sequential(\n            nn.Linear(512, k*k),nn.ReLU(),\n        )\n        self.k = k\n\n    def forward(self, x):\n        batchsize = x.size()[0]\n        x = self.conv(x)\n        x = torch.max(x, 2)[0]\n        x = self.fc(x)\n\n        iden = Variable(torch.from_numpy(np.eye(self.k).flatten().astype(np.float32))).view(1,\n                                                                            self.k*self.k).repeat(batchsize,1)\n        if x.is_cuda:\n            iden = iden.cuda()\n        x = x + iden\n        x = x.view(-1, self.k, self.k)\n        return x","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class PointNetfeat(nn.Module):\n    def __init__(self, global_feat = False, feature_transform = False, stn1_dim = 120,\n                 stn2_dim = 64):\n        super(PointNetfeat, self).__init__()\n        \n        self.global_feat = global_feat\n        self.feature_transform = feature_transform\n        self.stn1_dim = stn1_dim\n        self.stn2_dim = stn2_dim\n        \n        self.stn = STNkd(k=stn1_dim)\n        \n        self.conv1 = nn.Sequential(\n            nn.Conv1d(stn1_dim, 256, kernel_size=1), nn.ReLU(),\n        )\n        \n        self.conv2 = nn.Sequential(\n            nn.Conv1d(256, 256, kernel_size=1), nn.ReLU(),\n            nn.Conv1d(256, 1024, kernel_size=1), nn.ReLU(),\n            nn.Conv1d(1024, 2048, kernel_size=1), nn.ReLU(),\n        )\n        \n        if self.feature_transform:\n            self.fstn = STNkd(k=stn2_dim)\n\n    def forward(self, x):\n        n_pts = x.size()[2]\n        trans = self.stn(x)\n        x = x.transpose(2, 1)\n        x = torch.bmm(x, trans)\n        x = x.transpose(2, 1)\n        \n        x = self.conv1(x)\n\n        if self.feature_transform:\n            trans_feat = self.fstn(x)\n            x = x.transpose(2,1)\n            x = torch.bmm(x, trans_feat)\n            x = x.transpose(2,1)\n        else:\n            trans_feat = None\n\n        pointfeat = x\n        \n        x = self.conv2(x)\n        x = torch.max(x, 2)[0]\n        if self.global_feat:\n            return x, trans, trans_feat\n        else:\n            x = x[:,:,None].repeat(1, 1, n_pts)\n            return torch.cat([x, pointfeat], 1), trans, trans_feat","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class LyftNet(BaseNet):   \n    def __init__(self, **kwargs):\n        super().__init__(**kwargs)\n        \n        self.pnet = PointNetfeat()\n\n        self.fc0 = nn.Sequential(\n            nn.Linear(2048+256, 1024), nn.ReLU(),\n        )\n\n        self.fc = nn.Sequential(\n            nn.Linear(1024, 300),\n        )\n\n        self.c_net = nn.Sequential(\n            nn.Linear(1024, 3),\n        )\n        \n    \n    def forward(self, x):\n        bsize, npoints, hb, nf = x.shape \n        \n        # Push points to the last  dim\n        x = x.transpose(1, 3)\n\n        # Merge time with features\n        x = x.reshape(bsize, hb*nf, npoints)\n\n        x, trans, trans_feat = self.pnet(x)\n\n        # Push featuresxtime to the last dim\n        x = x.transpose(1,2)\n\n        x = self.fc0(x)\n\n        c = self.c_net(x)\n        x = self.fc(x)\n\n        return c,x\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","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def get_last_checkpoint(root):\n    res = None\n    mtime = -1\n    for model_name in Path(root).glob(\"lyfnet*.ckpt\"):\n        e = model_name.stat().st_ctime\n        if e > mtime:\n            mtime=e\n            res = model_name\n    return res\n\ndef get_last_version(root):\n\n    last_version = 0\n    for model_name in Path(root).glob(\"version_*\"):\n        version = int(model_name.as_posix().split(\"_\")[-1])\n        if version > last_version:\n            last_version = version\n    return last_version","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Training"},{"metadata":{"trusted":true},"cell_type":"code","source":"torch.backends.cudnn.benchmark =  True\n\nlast_checkpoint = get_last_checkpoint(ROOT)\n\nif last_checkpoint is not None:\n    print(f'\\n***** RESUMING FROM CHECKPOINT `{last_checkpoint.as_posix()}`***********\\n')\n    model = LyftNet.load_from_checkpoint(Path(last_checkpoint).as_posix(), \n    map_location=DEVICE, num_workers = NUM_WORKERS, batch_size = BATCH_SIZE)\nelse:\n    print('\\n***** NEW MODEL ***********\\n')\n\n    model = LyftNet(batch_size=BATCH_SIZE, \n                lr= 1e-3, weight_decay=5e-7, num_workers=NUM_WORKERS)\n\ncheckpoint_callback = ModelCheckpoint(\n    filepath=ROOT,\n    save_top_k=5,\n    verbose=0,\n    monitor='val_loss',\n    mode='min',\n    prefix='lyfnet_',\n)\n\n\nlogger = TensorBoardLogger(\n    save_dir=ROOT,\n    version=get_last_version(ROOT),\n    name='lightning_logs'\n)\n\nprint(model)\ntrainer = Trainer(\n    max_epochs=EPOCHS,\n    gradient_clip_val=GRADIENT_CLIP_VAL,\n    logger=logger,\n    checkpoint_callback=checkpoint_callback,\n    limit_val_batches=LIMIT_VAL_BATCHES,\n    gpus=1\n)\n\ntrainer.fit(model)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### This kernel is for training only, the one for inference could [be found here](https://www.kaggle.com/kneroma/inference-motion-prediction-with-pointnet)."},{"metadata":{"trusted":true},"cell_type":"code","source":"","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"<div style=\"text-align: center\"><strong>Thanks</strong></div>"},{"metadata":{"trusted":true},"cell_type":"code","source":"","execution_count":null,"outputs":[]}],"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat":4,"nbformat_minor":4}