<jupyter_start><jupyter_text><jupyter_code>import os 

os.makedirs('/content/lyft-motion-prediction-autonomous-vehicles/', exist_ok=True)
os.chdir('/content/lyft-motion-prediction-autonomous-vehicles/')

%cd '/content/lyft-motion-prediction-autonomous-vehicles/'
!pwd
os.makedirs('/content/lyft-motion-prediction-autonomous-vehicles/scenes/', exist_ok=True)
os.chdir('/content/lyft-motion-prediction-autonomous-vehicles/scenes/')

!wget --header="Host: storage.googleapis.com" --header="User-Agent: Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/84.0.4147.125 Safari/537.36" --header="Accept: text/html,application/xhtml+xml,application/xml;q=0.9,image/webp,image/apng,*/*;q=0.8,application/signed-exchange;v=b3;q=0.9" --header="Accept-Language: en-US,en;q=0.9" --header="Referer: https://www.kaggle.com/" --header="Cookie: _ga=GA1.3.1723755215.1594569279" --header="Connection: keep-alive" "https://storage.googleapis.com/kaggle-competitions-data/kaggle-v2/19990/1472735/compressed/scenes.zip?GoogleAccessId=<EMAIL>&Expires=1598865619&Signature=K6DsZz1qgin5fv8TMaTVg124%2FXRHRDgCOMeB4TimJIl%2F5hvXQo8T9OnVwqJ5LeeoXWvMTmyF7JFWK1EOlPx3mGKHT7Ofilnf7nSeHxIDxRO6QxjtDfMtJ3DQeTNeyIlAHJVqErjSEozQYCT6FM9sH8c%2B8yPx2hTl9qtiB92ISGUjmIKGsx%2BUIfQC6OfDxlyQa%2BO1edlpy6MFCHZkSJCYFsbbUYaoZYLbPKramldw%2Bn5TgmsCd%2F95r7PWEfjkZ2Niv26aROi6eSM3tHcygmZT3%2BOJPmXj%2FSJkGDRV3DdTsq8Ho%2F4z7CwFBa5U4Kihmp%2BG060mz5V%2BZoqj0LR9SEzbww%3D%3D&response-content-disposition=attachment%3B+filename%3Dscenes.zip" -c -O 'scenes.zip'
!unzip ./scenes.zip
%cd '/content/lyft-motion-prediction-autonomous-vehicles/scenes/'
!pwd
!ls
%rm -r validate.zarr/
%rm -r sample.zarr/
%rm scenes.zip
%cd ..
!pwd
!ls
os.makedirs('/content/lyft-motion-prediction-autonomous-vehicles/aerial_map/', exist_ok=True)
os.chdir('/content/lyft-motion-prediction-autonomous-vehicles/aerial_map/')

!wget --header="Host: storage.googleapis.com" --header="User-Agent: Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/84.0.4147.125 Safari/537.36" --header="Accept: text/html,application/xhtml+xml,application/xml;q=0.9,image/webp,image/apng,*/*;q=0.8,application/signed-exchange;v=b3;q=0.9" --header="Accept-Language: en-US,en;q=0.9" --header="Referer: https://www.kaggle.com/" --header="Cookie: _ga=GA1.3.1723755215.1594569279" --header="Connection: keep-alive" "https://storage.googleapis.com/kaggle-competitions-data/kaggle-v2/19990/1472735/compressed/aerial_map.zip?GoogleAccessId=<EMAIL>&Expires=1598867413&Signature=YBvVVAxzxbJhAIRNIfNklh2Qo%2By0kR8JggS33fKtpx%2BtO2LClR5e6OmRZ2g2UGsbmpMPAiuoZKwiL6vZqSh%2BJvnbbdTXrB%2BJphbDvWP6u3oc9udleGsHZKvJ4HLZvfdGF7%2B6Kpy8bVKwfyb4PrNrfSc4Q41%2BoAoGvCLFLG2onatz3QhgORhzlWfMqrbm1MY3KG4%2BZB5%2B3L0cJs2yKSfMclc89TkPTmORBMkSoNugfOGryEeru9qMyYRqL3kuUUliULfHrAnQB6q3F7R9u%2BNSuVEIViElfXlcdGVl79DC%2FNHawimzZ%2B%2BQcXIiI1fXC8w%2BhIiVnGZTYFqHd8hjSKQQfQ%3D%3D&response-content-disposition=attachment%3B+filename%3Daerial_map.zip" -c -O 'aerial_map.zip'
!unzip ./aerial_map.zip
%rm -r aerial_map.zip
%cd ..
!pwd
!ls
os.makedirs('/content/lyft-motion-prediction-autonomous-vehicles/semantic_map/', exist_ok=True)
os.chdir('/content/lyft-motion-prediction-autonomous-vehicles/semantic_map/')

!wget --header="Host: storage.googleapis.com" --header="User-Agent: Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/84.0.4147.125 Safari/537.36" --header="Accept: text/html,application/xhtml+xml,application/xml;q=0.9,image/webp,image/apng,*/*;q=0.8,application/signed-exchange;v=b3;q=0.9" --header="Accept-Language: en-US,en;q=0.9" --header="Referer: https://www.kaggle.com/" --header="Cookie: _ga=GA1.3.1723755215.1594569279" --header="Connection: keep-alive" "https://storage.googleapis.com/kaggle-competitions-data/kaggle-v2/19990/1472735/upload/semantic_map.zip?GoogleAccessId=<EMAIL>&Expires=1598867848&Signature=Hz%2F6RINYOAsBvD7sZF72nQnm5obsWKtzFrrK4%2FQePMEkQki9R5aSobiVGj8TijXWMJIRek%2F6o28k1cp4ld6TArQ9v%2F78hwCVjRN49YD3qQ3DBmNCUjSiYmp3FH5SeClQfjGQiydyPiaA3qaf3%2Bt3dHsvyUYwxmlmpPnF07vLu%2FcgVGhMQS3PS%2BJ1Ml0dhX7rP7i23zxsmQnc7Szr6mmolyd%2FoPRdBB9hXrHa5vc0axWIg3P44sqamtuVr9V0wF%2FHT9Zrsidc6sKOLnL9tGS%2BWYHS%2FWU58yVSgPUr%2FdOLCi2t7RUI8adP4YqQxa%2Bfle6zloLoeJCIibhroGTWyT2niQ%3D%3D&response-content-disposition=attachment%3B+filename%3Dsemantic_map.zip" -c -O 'semantic_map.zip'
!unzip ./semantic_map.zip
%rm -r semantic_map.zip
%cd ..
!pwd
!ls
## this script transports l5kit and dependencies
!pip -q install pymap3d==2.1.0 
!pip -q install protobuf==3.12.2 
!pip -q install transforms3d 
!pip -q install zarr 
!pip -q install ptable

!pip -q install --no-dependencies l5kit
# import packages
from google.colab import files
import numpy as np
import torch
import gc

from torch import nn, optim
from torch.utils.data import DataLoader
from torchvision.models.resnet import resnet18, resnet34, resnet50
from torchvision.models.densenet import densenet121
from tqdm import tqdm
from typing import Dict

from l5kit.data import LocalDataManager, ChunkedDataset
from l5kit.dataset import AgentDataset, EgoDataset
from l5kit.rasterization import build_rasterizer
INPUT_DIR = '/content/lyft-motion-prediction-autonomous-vehicles/'
cfg = {
    'format_version': 4,
    'model_params': {
        'model_architecture': 'densenet121',
        'history_num_frames': 10,
        'history_step_size': 1,
        'history_delta_time': 0.1,
        'future_num_frames': 50,
        'future_step_size': 1,
        'future_delta_time': 0.1
    },

    'raster_params': {
        'raster_size': [224, 224],
        'pixel_size': [0.5, 0.5],
        'ego_center': [0.25, 0.5],
        'map_type': 'py_semantic',
        'satellite_map_key': 'aerial_map/aerial_map.png',
        'semantic_map_key': 'semantic_map/semantic_map.pb',
        'dataset_meta_key': 'meta.json',
        'filter_agents_threshold': 0.5
    },
    
    'train_data_loader': {
        'key': 'scenes/train.zarr',
        'batch_size': 8,
        'shuffle': True,
        'num_workers': 4
    },
    
    'train_params': {
        'max_num_steps': 30000,
        'checkpoint_every_n_steps': 5000,
    }
}
# set env variable for data
os.environ["L5KIT_DATA_FOLDER"] = INPUT_DIR
dm = LocalDataManager(None)

# get config
print(cfg)
# ===== INIT DATASET
train_cfg = cfg["train_data_loader"]

# Rasterizer
rasterizer = build_rasterizer(cfg, dm)

# Train dataset/dataloader
train_zarr = ChunkedDataset(dm.require(train_cfg["key"])).open()
train_dataset = AgentDataset(cfg, train_zarr, rasterizer)
train_dataloader = DataLoader(train_dataset,
                              shuffle=train_cfg["shuffle"],
                              batch_size=train_cfg["batch_size"],
                              num_workers=train_cfg["num_workers"])

print(train_dataset)
gc.collect()
class LyftModel(nn.Module):
    
    def __init__(self, cfg: Dict):
        super().__init__()
        
        self.backbone = densenet121(pretrained=True, progress=True)
        
        num_history_channels = (cfg["model_params"]["history_num_frames"] + 1) * 2
        num_in_channels = 3 + num_history_channels

        self.backbone.features.conv0 = nn.Conv2d(
            num_in_channels,
            self.backbone.features.conv0.out_channels,
            kernel_size=self.backbone.features.conv0.kernel_size,
            stride=self.backbone.features.conv0.stride,
            padding=self.backbone.features.conv0.padding,
            bias=False,
        )
        
        # This is 512 for resnet18 and resnet34;
        # And it is 2048 for the other resnets
        # 1024 for densenet121.
        backbone_out_features = 1024

        # X, Y coords for the future positions (output shape: Bx50x2)
        num_targets = 2 * cfg["model_params"]["future_num_frames"]

        # You can add more layers here.
        self.head = nn.Sequential(
            # nn.Dropout(0.2),
            nn.Linear(in_features=backbone_out_features, out_features=2048),
        )

        self.logit = nn.Linear(2048, out_features=num_targets)
        
    def forward(self, x):
        x = self.backbone.features(x)
        # x = self.backbone.bn1(x)
        # x = self.backbone.relu(x)
        # x = self.backbone.maxpool(x)

        # x = self.backbone.layer1(x)
        # x = self.backbone.layer2(x)
        # x = self.backbone.layer3(x)
        # x = self.backbone.layer4(x)

        # x = self.backbone.avgpool(x)
        # x = torch.flatten(x, 1)
        
        x = self.head(x)
        x = self.logit(x)
        
        return x
# ==== INIT MODEL
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

model = LyftModel(cfg)
model.to(device)
optimizer = optim.Adam(model.parameters(), lr=1e-3)

# Later we have to filter the invalid steps.
criterion = nn.MSELoss(reduction="none")
device
# ==== TRAIN LOOP
tr_it = iter(train_dataloader)

progress_bar = tqdm(range(cfg["train_params"]["max_num_steps"]))
losses_train = []

for itr in progress_bar:

    try:
        data = next(tr_it)
    except StopIteration:
        tr_it = iter(train_dataloader)
        data = next(tr_it)

    model.train()
    torch.set_grad_enabled(True)
    
    # Forward pass
    inputs = data["image"].to(device)
    target_availabilities = data["target_availabilities"].unsqueeze(-1).to(device)
    targets = data["target_positions"].to(device)
    
    outputs = model(inputs).reshape(targets.shape)
    loss = criterion(outputs, targets)

    # not all the output steps are valid, but we can filter them out from the loss using availabilities
    loss = loss * target_availabilities
    loss = loss.mean()

    # Backward pass
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    losses_train.append(loss.item())

    if (itr+1) % cfg['train_params']['checkpoint_every_n_steps'] == 0:
        torch.save({
            'model_state_dict': model.state_dict()}, 
            f'resnet50_model_state_{itr}.pth')        
    progress_bar.set_description(f"loss: {loss.item()} loss(avg): {np.mean(losses_train[-100:])}")
torch.save(model.state_dict(), f'resnet50_model_state_last.pth')
from google.colab import files
files.download(f'model_state_last.pth') 
densenet121(pretrained=False)
resnet34(pretrained=False)
model<jupyter_output><empty_output>