{"cells":[{"metadata":{},"cell_type":"markdown","source":"# Motion Prediction with PointVoxelNet"},{"metadata":{},"cell_type":"markdown","source":"# About"},{"metadata":{},"cell_type":"markdown","source":"For more reading, please visit the [PointVoxelNet repo page](https://github.com/mit-han-lab/pvcnn) or [the original paper](https://arxiv.org/abs/1907.03739)."},{"metadata":{},"cell_type":"markdown","source":"![](https://hanlab.mit.edu/projects/pvcnn/figures/overview.png)\n\nPointVoxelNet work on 3D, I set the coord(x, y, z) with the z dim is zero. Because there is an issue on library (../input/point-voxel-cnn/modules/voxelization.py): coords.mean(2, keepdim=True)\nI still under contruction. So, some suggestions are very wellcome.\n\nThis source follow on this notebook: https://www.kaggle.com/kneroma/training-motion-prediction-with-pointnet\n\nThanks.\n"},{"metadata":{},"cell_type":"markdown","source":"# Library"},{"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 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\n        coords = X\n        # here, under contructing \n        coords = np.zeros((self.nframes, 3),  dtype=np.float32)\n        for i in range(self.nframes):\n            coords[i, :] = [X[i, 1, 1, 0], X[i, 1, 1, 1], 0]\n\n\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, coords\n        return X, target, target_availability, coords","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":"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=False)\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\n        else:\n            x = x[:,:,None].repeat(1, 1, n_pts)\n            return torch.cat([x, pointfeat], 1)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import sys\nsys.path.append('../input/point-voxel-cnn')\n\nimport modules.functional as F\nfrom modules.voxelization import Voxelization\nfrom modules.shared_mlp import SharedMLP\nfrom modules.se import SE3d\nfrom models.utils import create_pointnet_components\n\n\nclass PVConv(nn.Module):\n    def __init__(self, in_channels, out_channels, kernel_size, resolution, training=True, with_se=False, normalize=True, eps=0):\n        super().__init__()\n        self.in_channels = in_channels\n        self.out_channels = out_channels\n        self.kernel_size = kernel_size\n        self.resolution = resolution\n        self.training = training\n\n        self.voxelization = Voxelization(resolution, normalize=normalize, eps=eps)\n        voxel_layers = [\n            nn.Conv1d(in_channels, out_channels, kernel_size, stride=1, padding=kernel_size // 2),\n            nn.BatchNorm1d(out_channels, eps=1e-4),\n            nn.LeakyReLU(0.1, True),\n            nn.Conv1d(out_channels, out_channels, kernel_size, stride=1, padding=kernel_size // 2),\n            nn.BatchNorm1d(out_channels, eps=1e-4),\n            nn.LeakyReLU(0.1, True),\n         ]\n        if with_se:\n            voxel_layers.append(SE3d(out_channels))\n        self.voxel_layers = nn.Sequential(*voxel_layers)\n        self.point_features = SharedMLP(in_channels, out_channels)\n\n    def forward(self, features, coords):\n        print(coords.shape)\n        voxel_features, voxel_coords = self.voxelization(features, coords)\n        voxel_features = self.voxel_layers(voxel_features)\n        voxel_features = F.trilinear_devoxelize(voxel_features, voxel_coords, self.resolution, self.training)\n        fused_features = voxel_features + self.point_features(features)\n        return fused_features, coords\n\nclass PVConvfeat(nn.Module):\n    # out_channels, num_blocks, voxel_resolution\n    blocks = ((256, 1, 32), (1024, 1, 16), (2048, 1, None))\n\n    def __init__(self, global_feat = False):\n        super(PVConvfeat, self).__init__()\n        \n        self.global_feat = global_feat\n        \n        self.pvconv1 = PVConv(in_channels=120, out_channels=256, kernel_size=1, with_se=True, normalize=False, resolution=32)\n        \n        layers, channels_point, concat_channels_point = create_pointnet_components(\n            blocks=self.blocks, in_channels=256, with_se=True, normalize=False,\n            width_multiplier=1, voxel_resolution_multiplier=1\n        )\n        self.pvconv2 = nn.ModuleList(layers)\n\n    def forward(self, x, coord):\n        n_pts = x.size()[2]\n        x = self.pvconv1(x, coord)\n        pointfeat = x\n        x = self.pvconv2(x)\n\n        x = torch.max(x, 2)[0]\n\n        if self.global_feat:\n            return x\n        else:\n            x = x[:,:,None].repeat(1, 1, n_pts)\n            return torch.cat([x, pointfeat], 1)","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        self.pnet = PVConvfeat()\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, coord):\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 = self.pnet(x, coord)\n        # x = 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, k = [b.to(DEVICE) for b in batch]\n        c, preds = self(x, k)\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, k =  [b.to(DEVICE) for b in batch]\n        c, preds = self(x, k)\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    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, 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":{"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}