paninski-lab/lightning-pose

how to speed up prediction

Closed this issue · 5 comments

Wulin-Tan commented

Hi, LP team:
before prediction, LP needs to load the model and the dataset.

model = load_model_from_checkpoint(cfg=cfg, ckpt_file=ckpt_file, eval=True)

AND

dataset = get_dataset(cfg=cfg, data_dir=data_dir, imgaug_transform=imgaug_transform)

these 2 steps of loading are the most time consuming. (about 2~3 min for each in my toy example, the dataset step would be longer if I have hundreds of lableled images).

And I found that I can pickle the dataset into my disk and load this pickle file much faster. (about 2 min down to 14ms).

However, the model seems to be a dali pipeline that I cannot pickle.

Any suggestion about this?

themattinthehatt commented

@Wulin-Tan there are several ways to speed this up; for example you do not necessarily need to call both load_model_from_checkpoint and get_dataset. Are you using the predict_new_vids.py script? I see that's what we're doing there, I can speed that up. If not using that script, can you share with me the code snippet you're using for inference?

Another question: do you want to compute the various unsupervised losses on your video data during inference? You can also speed this step up if you don't (though those losses are useful for outlier detection, so I would recommend keeping them).

Wulin-Tan commented

Hi, @themattinthehatt

  1. Yes, I am using predict_new_vids.py. model is for the parameter 'model', and dataset is for the parameter 'data_module'.
model = load_model_from_checkpoint(cfg=cfg, ckpt_file=ckpt_file, eval=True)
......
data_module = get_data_module(cfg=cfg, dataset=dataset, video_dir=video_dir)

and then

export_predictions_and_labeled_video(
                video_file=video_file,
                cfg=cfg,
                ckpt_file=ckpt_file,
                prediction_csv_file=prediction_csv_file,
                labeled_mp4_file=labeled_mp4_file,
                trainer=trainer,
                model=model,
                data_module=data_module,
                save_heatmaps=cfg.eval.get("predict_vids_after_training_save_heatmaps", False),
            )

in this code, it seems that ckpt and dataset are both necessary. Any improvement we can have?

  1. which step or chunk of code to get losses? finally I get 2 more csv files other the predicted coordinates, like 'video_1_pca_singleview_error.csv' and 'video_1_temporal_norm.csv', are they the losses you mentioned?
themattinthehatt commented

@Wulin-Tan I have also updated predict_new_vids.py, which you will have with version 1.4.0. I rearranged the order of the code so that the data module only gets created once, which is what is taking so long (in the previous version it was unneccesarily created twice).

Note: I removed some of the configuration options, and force the output to be saved in the model folder under the subfolder video_preds!

Yes, video_1_pca_singleview_error.csv and video_1_temporal_norm.csv are the unsupervised losses, which you can use to flag outlier frames (kind of like the confidence; large values of the pca error or temporal norm are potentially problematic predictions). Like I said before I would recommend computing these values (as you see it happens automatically), but if you really want to speed things up you can not compute these, which means you don't need to create the data module. If you're interested in that let me know and I can show you how to do it (you'll need to write your own custom prediction script, but it's not much work).

Wulin-Tan commented

HI, @themattinthehatt
I tried the updated predict_new_vids.py, it works.

  • 👍1