save_yaml(out_dir / 'flags.yaml', flags_dict)
save_yaml(out_dir / 'cfg.yaml', cfg)
debug = flags.debug
# set env variable for data
os.environ["L5KIT_DATA_FOLDER"] = flags.l5kit_data_folder
dm = LocalDataManager(None)
# ===== INIT DATASET
if local_rank == 0:
print("INIT DATASET...")
train_cfg = cfg["train_data_loader"]
valid_cfg = cfg["valid_data_loader"]
# Rasterizer
rasterizer = build_custom_rasterizer(cfg, dm)
rasterizer_eval = build_custom_rasterizer(cfg, dm, eval=True)
print("rasterizer", rasterizer)
# Train dataset/dataloader
if flags.pred_mode in ["multi_agent", "rnn_head_multi", "yaw"]:
transform = pred_mode_to_transform[flags.pred_mode]
else:
transform = ImageAugmentation(flags).transform
transform_validation = transform
if not flags.augmentation_in_validation:
transform_validation = pred_mode_to_transform[flags.pred_mode]
collate_fn = pred_mode_to_collate_fn[flags.pred_mode]
train_path = "scenes/sample.zarr" if debug else train_cfg["key"]
'key': 'scenes/test.zarr',
'batch_size': 32,
'shuffle': False,
'num_workers': 4
}
test_cfg = cfg.get("test_data_loader", default_test_cfg)
# from copy import deepcopy
# cfg2 = deepcopy(cfg)
# cfg2["model_params"]["history_num_frames"] = 50
cfg["model_params"]["history_num_frames"] = 100
cfg["raster_params"][
"map_type"] = "stub_debug" # For faster calculation...
# Rasterizer
rasterizer = build_custom_rasterizer(cfg, dm)
valid_cfg = cfg["valid_data_loader"]
# valid_path = "scenes/sample.zarr" if debug else valid_cfg["key"]
valid_path = valid_cfg["key"]
valid_agents_mask = None
if flags.validation_chopped:
num_frames_to_chop = 100
th_agent_prob = cfg["raster_params"]["filter_agents_threshold"]
min_frame_future = 1
num_frames_to_copy = num_frames_to_chop
valid_agents_mask = load_mask_chopped(dm.require(valid_path),
th_agent_prob,
num_frames_to_copy,
min_frame_future)
print("valid_path", valid_path, "valid_agents_mask",
# Not use flags.l5kit_data_folder, but use fixed test data.
l5kit_data_folder = "../../input/lyft-motion-prediction-autonomous-vehicles"
os.environ["L5KIT_DATA_FOLDER"] = l5kit_data_folder
dm = LocalDataManager(None)
print("Load dataset...")
default_test_cfg = {
'key': 'scenes/test.zarr',
'batch_size': 32,
'shuffle': False,
'num_workers': 4
}
test_cfg = cfg.get("test_data_loader", default_test_cfg)
# Rasterizer
rasterizer = build_custom_rasterizer(cfg, dm)
tuned_cfg = deepcopy(cfg)
assert tuned_cfg["raster_params"]["map_type"] == "py_semantic"
# tuned_cfg["raster_params"]["map_type"] = "tuned_box+semantic_debug"
tuned_cfg["raster_params"]["map_type"] = "tuned_box+tuned_semantic"
tuned_rasterizer = build_custom_rasterizer(tuned_cfg, dm)
print("tuned_rasterizer", tuned_rasterizer)
valid_cfg = cfg["valid_data_loader"]
# valid_path = "scenes/sample.zarr" if debug else valid_cfg["key"]
valid_path = valid_cfg["key"]
valid_agents_mask = None
if flags.validation_chopped:
num_frames_to_chop = 100
th_agent_prob = cfg["raster_params"]["filter_agents_threshold"]