code: import numpy as np
  import torch
  from torch import nn, optim
  from torch.utils.data import DataLoader
  from torchvision.models.resnet import resnet50, resnet34
  from torch import Tensor
  from typing import Dict
  
  from l5kit.configs import load_config_data
  from l5kit.data import LocalDataManager, ChunkedDataset
  from l5kit.dataset import AgentDataset, EgoDataset
  from l5kit.rasterization import build_rasterizer
  from l5kit.evaluation import write_pred_csv, compute_metrics_csv, read_gt_csv, create_chopped_dataset
  from l5kit.evaluation.chop_dataset import MIN_FUTURE_STEPS
  from l5kit.evaluation.metrics import neg_multi_log_likelihood, time_displace
  from l5kit.geometry import transform_points
  from l5kit.visualization import PREDICTED_POINTS_COLOR, TARGET_POINTS_COLOR, draw_trajectory
  from pathlib import Path
  import pandas as pd
  import os
  import random
  import time
  import gc, psutil
  
  cfg = {
      'format_version': 4,
      'model_params': {
          'model_architecture': "resnet34",
          '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,
          'model_name': "model_resnet34",
          'lr': 1e-3,
          'train': True,
          'predict': True
      },
  
      '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': 16,
          'shuffle': True,
          'num_workers': 0
      },
      'test_data_loader': {
          'key': 'scenes/test.zarr',
          'batch_size': 16,
          'shuffle': False,
          'num_workers': 0,
      },
      'train_params': {
          'steps': 120,
          'update_steps': 50,
          'checkpoint_steps': 100,
          'precision': True
      }
  }
  
  class LyftCNNModel(nn.Module):
  
      def __init__(self, cfg: Dict, num_modes=3):
          super().__init__()
  
          architecture = cfg["model_params"]["model_architecture"]
          backbone = eval(architecture)(pretrained=True, progress=True)
          self.backbone = backbone
  
          num_history_channels = (cfg["model_params"]["history_num_frames"] + 1) * 2
          num_in_channels = 3 + num_history_channels
  
          self.backbone.conv1 = nn.Conv2d(
              num_in_channels,
              self.backbone.conv1.out_channels,
              kernel_size=self.backbone.conv1.kernel_size,
              stride=self.backbone.conv1.stride,
              padding=self.backbone.conv1.padding,
              bias=False,
          )
  
          if architecture == "resnet50":
              backbone_out_features = 2048
          else:
              backbone_out_features = 512
  
          # X, Y coords for the future positions (output shape: batch_sizex50x2)
          self.future_len = cfg["model_params"]["future_num_frames"]
          num_targets = 2 * self.future_len
  
          # You can add more layers here.
          self.head = nn.Sequential(
              # nn.Dropout(0.2),
              nn.Linear(in_features=backbone_out_features, out_features=4096),
          )
  
          self.num_preds = num_targets * num_modes
          self.num_modes = num_modes
  
          self.logit = nn.Linear(4096, out_features=self.num_preds + num_modes)
  
      def forward(self, x):
          x = self.backbone.conv1(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)
  
          # pred (batch_size)x(modes)x(time)x(2D coords)
          # confidences (batch_size)x(modes)
          bs, _ = x.shape
          pred, confidences = torch.split(x, self.num_preds, dim=1)
          pred = pred.view(bs, self.num_modes, self.future_len, 2)
          assert confidences.shape == (bs, self.num_modes)
          confidences = torch.softmax(confidences, dim=1)
          return pred, confidences
  
  
apis: torch.nn.Linear torch.split torch.flatten torch.nn.Conv2d torch.softmax