danbider / lightning-pose

Accelerated pose estimation and tracking using semi-supervised convolutional networks.
MIT License
223 stars 32 forks source link

how to speed up prediction #174

Closed Wulin-Tan closed 1 week ago

Wulin-Tan commented 2 weeks ago

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 1 week ago

@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 1 week ago

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?

  2. 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 1 week ago

@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 1 week ago

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

themattinthehatt commented 1 week ago

closed by 786644e4c3d187b95397b51dd01cbe098c87664f