from typing import Callable
import os
from typing import Optional, Tuple, Union
import warnings
from mmcv import Config
import torch
import wandb
from risk_biased.predictors.biased_predictor import (
LitTrajectoryPredictor,
LitTrajectoryPredictorParams,
)
from risk_biased.utils.config_argparse import config_argparse
from risk_biased.utils.cost import TTCCostParams
from risk_biased.utils.torch_utils import load_weights
from risk_biased.scene_dataset.loaders import SceneDataLoaders
from risk_biased.scene_dataset.scene import load_create_dataset
from risk_biased.utils.waymo_dataloader import WaymoDataloaders
def get_predictor(
config: Config, unnormalizer: Callable[[torch.Tensor, torch.Tensor], torch.Tensor]
):
params = LitTrajectoryPredictorParams.from_config(config)
model_class = LitTrajectoryPredictor
ttc_params = TTCCostParams.from_config(config)
return model_class(params=params, unnormalizer=unnormalizer, cost_params=ttc_params)
def load_from_wandb_id(
log_id: str,
log_path: str,
entity: str,
project: str,
config: Optional[Config] = None,
load_last=False,
) -> Tuple[Union[LitTrajectoryPredictor, LitTrajectoryPredictor], Config]:
"""
Load a model using a wandb id code.
Args:
log_id: the wandb id code
log_path: the wandb log directory path
config: An optional configuration argument, use these settings if not None, use the settings from the log directory otherwise
load_last: An optional argumument, set to True to load the last checkpoint instead of the best one
Returns:
Predictor model and config file either loaded from the checkpoint or the one passed as argument.
"""
list_matching = list(filter(lambda path: log_id in path, os.listdir(log_path)))
if len(list_matching) == 1:
list_ckpt = list(
filter(
lambda path: "epoch" in path and ".ckpt" in path,
os.listdir(os.path.join(log_path, list_matching[0], "files")),
)
)
if not load_last and len(list_ckpt) == 1:
print(f"Loading best model: {list_ckpt[0]}.")
checkpoint_path = os.path.join(
log_path, list_matching[0], "files", list_ckpt[0]
)
else:
print(f"Loading last checkpoint.")
checkpoint_path = os.path.join(
log_path, list_matching[0], "files", "last.ckpt"
)
config_path = os.path.join(
log_path, list_matching[0], "files", "learning_config.py"
)
if config is None:
config = config_argparse(config_path)
distant_model_type = None
else:
distant_config = config_argparse(config_path)
distant_model_type = distant_config.model_type
config["load_from"] = log_id
if config.model_type == "interaction_biased":
dataloaders = WaymoDataloaders(config)
else:
[data_train, data_val, data_test] = load_create_dataset(config)
dataloaders = SceneDataLoaders(
state_dim=config.state_dim,
num_steps=config.num_steps,
num_steps_future=config.num_steps_future,
batch_size=config.batch_size,
data_train=data_train,
data_val=data_val,
data_test=data_test,
num_workers=config.num_workers,
)
try:
if len(config.gpus):
map_location = "cpu"
else:
map_location = "gpu"
model = load_weights(
get_predictor(config, dataloaders.unnormalize_trajectory),
torch.load(checkpoint_path, map_location=map_location),
strict=True,
)
except RuntimeError:
raise RuntimeError(
f"The source model is of type {distant_model_type}."
+ " It cannot be used to load the weights of the interaction biased model."
)
return model, dataloaders, config
else:
print("Trying to download logs from WandB...")
api = wandb.Api()
run = api.run(entity + "/" + project + "/" + log_id)
if run is not None:
checkpoint_path = os.path.join(
log_path, "downloaded_run-" + log_id, "files"
)
os.makedirs(checkpoint_path)
for file in run.files():
if file.name.endswith("ckpt") or file.name.endswith("config.py"):
file.download(checkpoint_path)
return load_from_wandb_id(
log_id, log_path, entity, project, config, load_last
)
else:
raise RuntimeError(
f"Error while loading checkpoint: Found {len(list_matching)} occurences of the given id {log_id} in the logs at {log_path}."
)
def load_from_config(cfg: Config):
"""
This function loads the predictor model and the data depending on which one is selected in the config.
If a "load_from" field is not empty, then tries to load the pre-trained model from the checkpoint.
The matching config file is loaded
Args:
cfg : Configuration that defines the model to be loaded
Returns:
loaded model and a new version of the config that is compatible with the checkpoint model that it could be loaded from
"""
log_path = os.path.join(cfg.log_path, "wandb")
ignored_keys = [
"project",
"dataset_parameters",
"load_from",
"force_config",
"load_last",
]
if "load_from" in cfg.keys() and cfg.load_from != "" and cfg.load_from:
if "load_last" in cfg.keys():
load_last = cfg["load_last"]
else:
load_last = False
if cfg.force_config:
warnings.warn(
f"Using local configuration but loading from run {cfg.load_from}. Will fail if local configuration is not compatible."
)
predictor, dataloaders, config = load_from_wandb_id(
log_id=cfg.load_from,
log_path=log_path,
entity=cfg.entity,
project=cfg.project,
config=cfg,
load_last=load_last,
)
else:
predictor, dataloaders, config = load_from_wandb_id(
log_id=cfg.load_from,
log_path=log_path,
entity=cfg.entity,
project=cfg.project,
load_last=load_last,
)
difference = False
warning_message = ""
for key, item in cfg.items():
try:
if config[key] != item:
if not difference:
warning_message += "When loading the model, the configuration was changed to match the configuration of the pre-trained model to be loaded.\n"
difference = True
if key not in ignored_keys:
warning_message += f" The value of '{key}' is now '{config[key]}' instead of '{item}'."
except KeyError:
if not difference:
warning_message += "When loading the model, the configuration was changed to match the configuration of the pre-trained model to be loaded."
difference = True
warning_message += f" The parameter '{key}' with value '{item}' does not exist for the model you are loading from, it is added."
config[key] = item
if warning_message != "":
warnings.warn(warning_message)
return predictor, dataloaders, config
else:
if cfg.model_type == "interaction_biased":
dataloaders = WaymoDataloaders(cfg)
else:
[data_train, data_val, data_test] = load_create_dataset(cfg)
dataloaders = SceneDataLoaders(
state_dim=cfg.state_dim,
num_steps=cfg.num_steps,
num_steps_future=cfg.num_steps_future,
batch_size=cfg.batch_size,
data_train=data_train,
data_val=data_val,
data_test=data_test,
num_workers=cfg.num_workers,
)
predictor = get_predictor(cfg, dataloaders.unnormalize_trajectory)
return predictor, dataloaders, cfg