Pytorch_lightning Trainer Predict execution time

1

Hi,

I have a model and dataloader with 1 image.

The main function is updating the image by looping then predicting the image.

import time
# load model and weight
model = ...
if __name__ == "__main__":
    while(True):
        # load image
        start_time = time.time()
        dataloader = process_image_to_dataloader
        trainer.predict(model=model, dataloaders=[dataloader], return_predictions=True)
        print("--- %s seconds ---" % (time.time() - start_time))

As my experiment, the execution time is around 12s for 1 image prediction.
Another experiment that predicts 10 images, the execution time is also around 12s.
→ That means the number of images is not counting for execution time. (dataloader is not take long time ~0s)

Is there any efficient way to predict 1 image in loop that minimizes the execution time?

Thank you!

2

Are the predictions being done on CPU or GPU? If they are being done on GPU, you may want to add torch.cuda.synchronize() before you start and stop timing to ensure that the timing is done in a valid fashion.

In general, there are many ways to improve the performance of inference, you may want to check out torch.jit.script as well as torch.compile:
https://pytorch.org/docs/stable/generated/torch.jit.script.html

https://pytorch.org/tutorials/intermediate/torch_compile_tutorial.html

(the latter of which is new PyTorch 2.0 feature)

3

I predict by GPU. I added torch.cuda.synchronize() but the execution time was still the same.
I am using torch==1.12.1+cu116, so torch.compile might not be supported.
Let’s me try to use torch.jit.script, it look quite difficult to me.~