{"cells":[{"metadata":{"trusted":true},"cell_type":"code","source":"import sys\nsys.path.insert(0, '/kaggle/input/kb-l5kit/l5kit/')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import os\nfrom pprint import pprint\nfrom typing import Dict\n\nimport numpy as np\nimport torch\nfrom torch import nn, optim\nfrom torch.utils.data import DataLoader, random_split\nfrom torchvision.models.mobilenet import mobilenet_v2\nfrom tqdm.notebook import tqdm\nimport wandb\n\nfrom l5kit.configs import load_config_data\nfrom l5kit.data import LocalDataManager, ChunkedDataset\nfrom l5kit.dataset import AgentDataset, EgoDataset\nfrom l5kit.rasterization import build_rasterizer\nfrom l5kit.evaluation import write_pred_csv","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Prepare Data path and load cfg\nBy setting the L5KIT_DATA_FOLDER variable, we can point the script to the folder where the data lies.\n\nThen, we load our config file with relative paths and other configurations (rasteriser, training params...).","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"# set env variable for data\nos.environ[\"L5KIT_DATA_FOLDER\"] = \"/kaggle/input/lyft-motion-prediction-autonomous-vehicles/\"\ndm = LocalDataManager(None)\n# get config\ncfg = load_config_data(\"/kaggle/input/lyft-config-files/agent_motion_config.yaml\")","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"cfg['model_params']['history_num_frames'] = 10\ncfg['train_params']['max_num_steps'] = 25000\ncfg['train_params']['checkpoint_every_n_steps'] = 10000\ncfg['train_data_loader']['batch_size'] = 12\ncfg['train_data_loader']['num_workers'] = 4\ncfg['train_data_loader']['key'] = 'scenes/train.zarr'\npprint(cfg)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Model\n\nWe will be building MobileNetV2 here and concatenating agent states (velocity, acceleration and heading offsets). We must replace the input and the final layer to address our requirements.","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"class Identity(nn.Module):\n    def __init__(self):\n        super(Identity, self).__init__()\n\n    def forward(self, x):\n        return x\n    \nclass MobilenetV2LSTM(nn.Module):\n    def __init__(self, config: Dict):\n        super(MobilenetV2LSTM, self).__init__()\n        self.cfg = config\n        self.batch_size = self.cfg['train_data_loader']['batch_size']\n        self.hist_frames = self.cfg['model_params']['history_num_frames']\n        self.fc_infeatures = 1280 + (2 * (self.hist_frames + 1)) + (2 * self.hist_frames) + (self.hist_frames + 1)\n        self.num_targets = 2 * self.cfg[\"model_params\"][\"future_num_frames\"]\n#         self.seq_len = 1\n#         self.input_size = 128\n#         self.hidden_size = 128\n        self.cnn = self.build_basecnn()\n        self.fc1 = nn.Sequential(\n            nn.Dropout(p=0.2, inplace=False),\n            nn.Linear(in_features=self.fc_infeatures, out_features=4096),\n            nn.ReLU(inplace=True),\n            nn.Dropout(p=0.2, inplace=False),\n            nn.Linear(in_features=4096, out_features=self.num_targets)\n        )\n#         self.lstm = nn.LSTM(\n#             input_size=self.input_size,\n#             hidden_size=self.hidden_size,\n#             num_layers=1,\n#             batch_first=True\n#         )\n#         self.fc2 = nn.Linear(in_features=128, out_features=100)\n#         self.hidden_cell = (torch.zeros(self.batch_size, 1, self.hidden_size),\n#                             torch.zeros(self.batch_size, 1, self.hidden_size))\n\n    def forward(self, x, vel, accel, yaw):\n        x = self.cnn(x)\n        vel = vel.reshape(self.batch_size, -1)\n        accel = accel.reshape(self.batch_size, -1)\n        yaw = yaw.reshape(self.batch_size, -1)\n        x = torch.cat([x, vel, accel, yaw], dim=1)\n        x = self.fc1(x)\n#         cnn_out = self.fc1(x)\n#         lstm_in = cnn_out.view(self.batch_size, self.seq_len, self.input_size)\n#         lstm_out, self.hidden_cell = self.lstm(lstm_in)\n#         fc_in = lstm_out.view(self.batch_size, lstm_out.shape[2])\n#         x = self.fc2(fc_in)\n\n        return x\n\n    def build_basecnn(self):\n        # change input channels number to match the rasterizer's output\n        mnet = mobilenet_v2(pretrained=True)\n        num_history_channels = (self.cfg[\"model_params\"][\"history_num_frames\"] + 1) * 2\n        num_in_channels = 3 + num_history_channels\n        mnet.features[0][0] = nn.Conv2d(\n            num_in_channels,\n            mnet.features[0][0].out_channels,\n            kernel_size=mnet.features[0][0].kernel_size,\n            stride=mnet.features[0][0].stride,\n            padding=mnet.features[0][0].padding,\n            bias=False,\n        )\n\n        mnet.classifier = Identity()\n        \n        return mnet","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# def build_model(cfg: Dict) -> torch.nn.Module:\n#     # load pre-trained Conv2D model\n#     model = mobilenet_v2(pretrained=True)\n\n#     # change input channels number to match the rasterizer's output\n#     num_history_channels = (cfg[\"model_params\"][\"history_num_frames\"] + 1) * 2\n#     num_in_channels = 3 + num_history_channels\n#     model.features[0][0] = nn.Conv2d(\n#         num_in_channels,\n#         model.features[0][0].out_channels,\n#         kernel_size=model.features[0][0].kernel_size,\n#         stride=model.features[0][0].stride,\n#         padding=model.features[0][0].padding,\n#         bias=False,\n#     )\n#     # change output size to (X, Y) * number of future states\n#     num_targets = 2 * cfg[\"model_params\"][\"future_num_frames\"]\n#     model.classifier[1] = nn.Linear(in_features=model.classifier[1].in_features, out_features=num_targets)\n\n#     return model","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def forward(data, model, device, criterion):\n    im_inputs = data[\"image\"].to(device)\n    vel_inputs = data[\"history_velocities\"].to(device)\n    accel_inputs = data[\"history_accels\"][:, :-1, :].to(device)  # removing last history frame since we don't have accel for it\n    yaw_inputs = data[\"history_yaws\"].to(device)\n\n    target_availabilities = data[\"target_availabilities\"].unsqueeze(-1).to(device)\n    targets = data[\"target_positions\"].to(device)\n    # Forward pass\n    outputs = model(im_inputs, vel_inputs, accel_inputs, yaw_inputs).reshape(targets.shape)\n    loss = criterion(outputs, targets)\n    # not all the output steps are valid, but we can filter them out from the loss using availabilities\n    loss = loss * target_availabilities\n    loss = loss.mean()\n    return loss, outputs","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Load the Train Data\n\nOur data pipeline map a raw `.zarr` folder into a multi-processing instance ready for training by:\n- loading the `zarr` into a `ChunkedDataset` object. This object has a reference to the different arrays into the zarr (e.g. agents and traffic lights);\n- wrapping the `ChunkedDataset` into an `AgentDataset`, which inherits from torch `Dataset` class;\n- passing the `AgentDataset` into a torch `DataLoader`","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"rasterizer = build_rasterizer(cfg, dm)\ntrain_cfg = cfg[\"train_data_loader\"]\ntrain_zarr = ChunkedDataset(dm.require(train_cfg[\"key\"])).open()\ntrain_dataset = AgentDataset(cfg, train_zarr, rasterizer)\n\n# create training splits for transfer learning --- wait for pytorch update on kaggle kernels\n# train1_size = 25000\n# trainrem_size = len(train_dataset) - train1_size\n# train_dataset, _ = random_split(train_dataset, [train1_size, trainrem_size], generator=torch.Generator().manual_seed(42))\ntrain_dataloader = DataLoader(train_dataset, shuffle=train_cfg[\"shuffle\"], batch_size=train_cfg[\"batch_size\"], \n                             num_workers=train_cfg[\"num_workers\"])\nprint(train_dataset)\nprint(len(train_dataset))\nprint(len(train_dataloader))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# ==== INIT MODEL\ndevice = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\nmodel = MobilenetV2LSTM(cfg).to(device)\n# model = build_model(cfg).to(device)\noptimizer = optim.Adam(model.parameters(), lr=1e-3)\ncriterion = nn.MSELoss(reduction=\"none\")","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Training","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"def train(model, device, train_loader, criterion, optimizer, max_steps, checkpoint_steps):\n    # Switch model to training mode. This is necessary for layers like dropout, batchnorm etc \n    # which behave differently in training and evaluation mode\n    model.train()\n    losses_train = []\n    step = 1\n    \n    # We loop over the data iterator, and feed the inputs to the network and adjust the weights.\n    train_pbar = tqdm(enumerate(train_loader), desc=\"Training steps\", leave=True, total=max_steps)\n    for batch_idx, data in train_pbar:\n        # Forward pass to calculate loss\n        loss, _ = forward(data, model, device, criterion)\n\n        # Reset the gradients to 0 for all learnable weight parameters\n        optimizer.zero_grad()\n\n        # Backward pass: compute the gradients of the loss w.r.t. the model's parameters\n        loss.backward()\n\n        # Update the model weights\n        optimizer.step()\n\n        # Get average loss of iterations so far\n        losses_train.append(loss.item())\n        avg_train_loss = np.mean(losses_train)\n\n        # wandb logging - metrics to track\n        wandb.log({'Training Loss': loss.item(), 'Avg Training Loss': avg_train_loss})\n\n        train_pbar.set_description(f\" Avg train loss: {avg_train_loss}\")\n        \n        if step % checkpoint_steps == 0 and step != max_steps:\n            torch.save(model.state_dict(), f'l5run3_iter{step}_mobilenetv2.pth')\n        \n        if step >= max_steps:\n            return avg_train_loss\n        step += 1\n\n    return avg_train_loss","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Main training and validation loop\n\n# Initialize wandb\nwandb.init(project='lyft-competition', name='mobilenetv2', anonymous='allow')\n\n# WandB – Config is a variable that holds and saves hyperparameters and inputs\nwandb_cfg = wandb.config                                         # initialize wandb config\nwandb_cfg.batch_size = cfg['train_data_loader']['batch_size']    # input batch size for training (default: 64)\nwandb_cfg.test_batch_size = cfg['val_data_loader']['batch_size'] # input batch size for testing (default: 1000)\nwandb_cfg.steps = cfg['train_params']['max_num_steps']\nwandb_cfg.epochs = 1                                             # number of epochs to train (default: 10)\nwandb_cfg.lr = 1e-3                                              # learning rate (default: 0.01)\nwandb_cfg.seed = 42                                              # random seed (default: 42)\nwandb_cfg.log_interval = 0                                       # how many batches to wait before logging training status\n\n# Set random seeds and deterministic pytorch for reproducibility\n# random.seed(wandb_cfg.seed)       # python random seed\n# torch.manual_seed(wandb_cfg.seed) # pytorch random seed\n# np.random.seed(wandb_cfg.seed)    # numpy random seed\n# torch.backends.cudnn.deterministic = True\n\n# WandB – wandb.watch() automatically fetches all layer dimensions, gradients, model parameters and logs them automatically to your dashboard.\n# Using log=\"all\" log histograms of parameter values in addition to gradients\n# wandb.watch(model)\n\n# pbar = tqdm(range(1, wandb_cfg.epochs+1), desc=\"Epochs\")\n# for epoch in pbar:\n#     train_loss = train(model, device, train_dataloader, criterion, optimizer, epoch)\n#     val_loss = val(model, device, val_dataloader, criterion, epoch)\n#     print(f\"Epoch: {epoch}, Train loss: {train_loss}, Val loss: {val_loss}\")\n\ncheckpoint_steps =  cfg['train_params']['checkpoint_every_n_steps']   \ntrain_loss = train(model, device, train_dataloader, criterion, optimizer, wandb_cfg.steps, checkpoint_steps)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Save model\ntorch.save(model.state_dict(), 'l5run3_mobilenetv2.pth')","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}