Scaling Out Motion Prediction for Autonomous Vehicles with L5Kit, Ray, and W&B

import torch # 🔦
from torch import nn
from torchvision.models.resnet import resnet50

# 🔦
def build_model(config: Dict) -> torch.nn.Module:
    # load pre-trained Conv2D model
    model = resnet50(pretrained=True)

    # change input channels number to match the rasterizer's output
    num_history_channels = (config["model_params"]["history_num_frames"] + 1) * 2
    num_in_channels = 3 + num_history_channels
    model.conv1 = nn.Conv2d(
        num_in_channels,
        model.conv1.out_channels,
        kernel_size=model.conv1.kernel_size,
        stride=model.conv1.stride,
        padding=model.conv1.padding,
        bias=False,
    )

    # change output size to (X, Y) * number of future states
    num_targets = 2 * config["model_params"]["future_num_frames"]
    model.fc = nn.Linear(in_features=2048, out_features=num_targets)

    return model

import torch # 🔦
import multiprocessing # 🔎
from ray import train
from ray.air import session, Checkpoint # 💾
from ray.train.torch import TorchTrainer # 🔦💨⚡️
from ray.air.config import ScalingConfig

# 🔦+🔎 Get all your useful hardware information
USE_GPU = torch.cuda.is_available()
NUM_GPUS = torch.cuda.device_count()
NUM_CPUS = multiprocessing.cpu_count()

# 🧮 Do some quick math to efficiently spread the workload to each GPU
if USE_GPU:
    num_actors = NUM_GPUS
    num_data_workers = NUM_CPUS // num_actors
else:
    num_data_workers = 4 if NUM_CPUS>=4 else NUM_CPUS
    ideal_num_actors = NUM_CPUS // num_data_workers
    num_actors = ideal_num_actors if ideal_num_actors else 1

# 🔦 Define the details of a training run
def train_model(config): 
	train_dataloader = load_data(config) # 🫵🏽 You implement this
	# ⚡️ Prepares Data for DDP
	train_dataloader = train.torch.prepare_data_loader(train_dataloader)

	model, criterion, optimizer = build_model(config) # 🔦+🫵🏽 and this...
	# ⚡️ Prepares Model for DDP
    	model = train.torch.prepare_model(model)

	steps_before_checkpointing = config.get("steps_before_checkpointing", 100)
	max_epochs = config.get("max_epochs", 100)
	
	for epoch in range(max_epochs):
		model.train()
		torch.set_grad_enabled(True)
        	metrics, model_outputs = train_model_epoch(train_dataloader,
								model, criterion,
								optimizer) # 🫵🏽 AND this!
        
		# 🪄🐝 Report your training metrics to your logger of choice
		# 💾 This includes the model checkpoints serialized by Ray
      		if (step%steps_before_checkpointing==0) or (step==max_epochs-1):
            		session.report(
                		metrics=metrics,
                		checkpoint=Checkpoint.from_dict(dict(epoch=epoch, model=model)))
        	else:
            		session.report(
                		metrics=metrics
            		)

# 🔦💨⚡️ Ray deals with the rest of the hardware management for you SUPAFAST
trainer = TorchTrainer(
    train_loop_per_worker=train_model,
    scaling_config=ScalingConfig(num_workers=num_actors, use_gpu=USE_GPU),
)


from ray import tune
from ray.tune.tuner import Tuner # 📻♻️
from ray.air.callbacks.wandb import WandbLoggerCallback # 🪄🐝
from ray.tune.stopper import ExperimentPlateauStopper # 🛑
from ray.tune.search.optuna import OptunaSearch # 🕵🏽

project_name = "your_cool_project_name"
n_search_attempts = 25

# 🕵🏽 Define your hyperparameter optimization search strategy
# You can define for which config values we want to search and how
config["max_epochs"] = tune.quniform(1000, 5000, 250)
optuna_search = OptunaSearch()

# 📻♻️ Define details of the ray.tune job by providing:
tuner = Tuner(
        trainer, # 🔦💨⚡️ The trainer from above...
        tune_config=tune.TuneConfig(
            metric="your_important_metric",
            mode="min",
            search_alg=optuna_search, # 🕵🏽 How we want to search.
            num_samples=n_search_attempts,
        ),
        param_space={
            "train_loop_config": config # 🔦💨⚡️ ...with the config we want to pass to the trainer!
        },
        run_config=RunConfig(
	    # 🛑 How to stop the experiment based on a provided metric.
            stop=ExperimentPlateauStopper("another_important_metric"),
	    # 🪄🐝 Your ML System of Record (Tracked Experiments & Checkpoint Artifacts).
            callbacks=[WandbLoggerCallback(project=f"{project_name}-trials", 
						save_checkpoints=True),]))
# 🏃🏽‍♂️ Run the tune job with
# - Distributed Training
# - Hyperparameter Search
# - Experiment Stopping
# - MLOps
analysis = tuner.fit()