compute_features.py

# # Copyright (c) Meta Platforms, Inc. and affiliates.
# # LICENSE file in the root directory of this source tree.
# import random
# import os
# from datetime import datetime
# import torch
# import pytorch_lightning as pl
# from dataset.egoexo4d.dataloader import filter_narration, clean_narration_text
# from lib.clip_model import ClipPLModel
# from lib.data_modules import EgoExo4dDataModule, UnsupEgoExo4dDataModule
# from lib.evaluation import evaluate
# from argparse import ArgumentParser
# import yaml

# def train(configs):

#     random.seed(1234)

#     # Load Model Parameters
#     model_hparams = configs.get("model_hparams", {})
#     model_name = model_hparams.get("model_name")
#     model_suffix = model_hparams.get("model_suffix", "")
#     imu_encoder_name = model_hparams.get("imu_encoder_name")
#     audio_encoder_name = model_hparams.get("audio_encoder_name")
#     video_encoder_name = model_hparams.get("video_encoder_name")
#     window_sec = model_hparams.get("window_sec")
#     target_fps = model_hparams.get("target_fps")
#     datasetname = model_hparams.get("datasetname", "ego4d")
#     final_embedding_size = model_hparams.get("final_embedding_size", 512)

#     # Params for the trainer
#     train_hparams = configs.get("train_hparams", {})
#     source_modality = train_hparams.get("source_modality")
#     target_modalities = train_hparams.get("target_modalities")
#     limit_train_batches = train_hparams.get("limit_train_batches")
#     batch_size = train_hparams.get("batch_size")
#     max_epochs = train_hparams.get("max_epochs")
#     gpus = train_hparams.get("gpus")
#     num_workers_for_dm = train_hparams.get("num_workers_for_dm")
#     test_only = train_hparams.get("test_only")
#     trainer_strategy = train_hparams.get("trainer_strategy")
#     freeze_modalities = train_hparams.get("freeze_modalities")
#     path_load_pretrained_imu_encoder = train_hparams.get(
#         "path_load_pretrained_imu_encoder"
#     )

#     # Paths, etc.
#     path_root_save_dir = f"./saved/{model_name}"
#     if not os.path.exists(path_root_save_dir):
#         os.makedirs(path_root_save_dir)
#     target_modalities.sort()
#     list_modalities = [source_modality] + target_modalities
#     source_modality_initial = source_modality[0]
#     target_modality_initials = "".join([m[0] for m in target_modalities])
#     if source_modality == "imu":
#         source_encoder_name = imu_encoder_name
#     if source_modality == "audio":
#         source_encoder_name = audio_encoder_name
#     model_identifier = (
#         f"{model_name}_s_{source_modality_initial}_t_{target_modality_initials}"
#         + f"_se_{source_encoder_name}_w_{window_sec}"
#     )
#     if model_suffix != "":
#         model_identifier += "_" + model_suffix
#     else:
#         model_identifier += "_%d" % (int(datetime.now().timestamp() % 10000))
#     path_save_checkpoint = f"{path_root_save_dir}/{model_identifier}_best.ckpt"
#     path_save_src_encoder = f"{path_root_save_dir}/{model_identifier}_src_encoder.pt"
#     result_path = f"./results/{model_identifier}"
#     configs["path_save_checkpoint"] = path_save_checkpoint

#     # Initialize the data module
#     dataset_params = {
#         "window_sec": window_sec,
#         "target_fps": target_fps,
#         "list_modalities": list_modalities,
#         "clean_narration_func": clean_narration_text,
#         "filter_narration_func": filter_narration,
#         "imu_sampling_rate": imu_sampling_rate,
#     }

#     if "text" in list_modalities:
#         datamodule = EgoExo4dDataModule(
#             batch_size=batch_size,
#             num_workers=num_workers_for_dm,
#             pin_memory=True,
#             drop_last=True,
#             dataset_params=dataset_params,
#         )
#     else:
#         datamodule = UnsupEgoExo4dDataModule(
#             batch_size=batch_size,
#             num_workers=num_workers_for_dm,
#             pin_memory=True,
#             drop_last=True,
#             dataset_params=dataset_params,
#         )

#     # Initialize encoder models
#     text_encoder, video_encoder, imu_encoder = None, None, None
#     modality_to_encoder = {}

#     if "text" in list_modalities:
#         # For now we only use a CLIP-based text model
#         text_encoder = ClipPLModel(freeze=True)
#         modality_to_encoder["text"] = text_encoder

#     if "video" in list_modalities:
#         # For now we only use a CLIP-based image model as a video encoder
#         video_encoder = (
#             ClipPLModel(freeze=True) if text_encoder is None else text_encoder
#         )
#         video_encoder.video_encoder_name = video_encoder_name

#         modality_to_encoder["video"] = video_encoder

#     for modality in list_modalities:
#         if modality in freeze_modalities:
#             modality_to_encoder[modality].eval()
#             print("Freezing modality: ", modality)
#             modality_to_encoder[modality].freeze()

#     # Initialize the training module for contrastive training
#     model = MultimodalContrastiveLearningModule(
#         modality_to_encoder=modality_to_encoder,
#         source_modality=source_modality,
#         target_modalities=target_modalities,
#     )

#     # Checkpoint settings
#     checkpoint_callback = pl.callbacks.ModelCheckpoint(
#         monitor="val_loss",
#         dirpath=path_root_save_dir,
#         filename=f"{model_identifier}" + "-{epoch:02d}-{val_loss:.2f}",
#         save_top_k=3,
#         mode="min",
#     )

#     # Initialize Trainer
#     trainer = pl.Trainer(
#         max_epochs=max_epochs,
#         gpus=gpus,
#         strategy=trainer_strategy,
#         limit_train_batches=limit_train_batches,
#         enable_checkpointing=True,
#         callbacks=[checkpoint_callback],
#     )

#     if not test_only:
#         # Start training
#         print("Start training: [%s] ..." % path_save_checkpoint)
#         trainer.fit(model, datamodule=datamodule)

#         # Save the checkpoint & encoder to a temp folder
#         torch.distributed.barrier()
#         print("Best checkpoint:", checkpoint_callback.best_model_path)
#         model.load_from_checkpoint(
#             checkpoint_callback.best_model_path,
#             modality_to_encoder=modality_to_encoder,
#             source_modality=source_modality,
#             target_modalities=target_modalities,
#         )
#         src_encoder = None
#         if source_modality == "imu":
#             src_encoder = model.imu_encoder
#         elif source_modality == "audio":
#             src_encoder = model.audio_encoder
#         elif source_modality == "video":
#             src_encoder = model.video_encoder
#         torch.save(src_encoder.state_dict(), path_save_src_encoder)
#     else:
#         print("Skipping training ...")

#     # Test the performance
#     print("Start evaluating ...")
#     metrics = evaluate(
#         datamodule.get_dataset(
#             "test",
#             window_sample_rate=1.0,
#             video_uid_sample_rate=0.25,
#             max_n_windows_per_video=2,
#         ),
#         datamodule.collate_fn,
#         model,
#         source_modality,
#         target_modalities,
#         result_path,
#         configs,
#     )
#     print(metrics)
#     return metrics

# if __name__ == "__main__":


#     parser = ArgumentParser()

#     # Main parameters are defined in a YAML file
#     parser.add_argument(
#         "--path_configs", default="./configs/train_contrastive/default.yaml"
#     )

#     # Override-params for a quick resource allocation adjustment or for debugging purposes
#     # If it is *not* None, the values in args override the values in the YAML file.
#     parser.add_argument("--gpus", default=None)
#     parser.add_argument("--num_workers_for_dm", default=None)
#     parser.add_argument("--max_epochs", default=None)
#     parser.add_argument("--test_only", default=None)
#     parser.add_argument("--path_load_pretrained_imu_encoder", default=None)
#     args = parser.parse_args()

#     # Load the YAML file
#     with open(args.path_configs) as f:
#         configs = yaml.load(f, Loader=yaml.FullLoader)

#     # Override the configs with args, if requested
#     if args.gpus is not None:
#         configs["train_hparams"]["gpus"] = int(args.gpus)
#     if args.num_workers_for_dm is not None:
#         configs["train_hparams"]["num_workers_for_dm"] = int(args.num_workers_for_dm)
#     if args.max_epochs is not None:
#         configs["train_hparams"]["max_epochs"] = int(args.max_epochs)
#     if args.test_only is not None:
#         configs["train_hparams"]["test_only"] = eval(args.test_only)
#     if args.path_load_pretrained_imu_encoder is not None:
#         configs["train_hparams"][
#             "path_load_pretrained_imu_encoder"
#         ] = args.path_load_pretrained_imu_encoder

#     print(configs)
#     train(configs)