Need help with running inference faster in GPU

[Vishnu280412]

Hi @felixdittrich92 👋,

So I was trying to create a fastapi endpoint of DocTR for the project and I am using a finetuned model. I have also installed cuda enabled torch and I wanted to make sure if what I am doing to run the inference is correct, because it took around 10 seconds to process 5 images.

The model loading part:

1. Initially I was doing something like this:

class DocTR:
    def __init__(self):
        custom_vocab = CUSTOM_VOCAB

        # Detection model initialisation
        det_model = db_resnet50(pretrained=False, pretrained_backbone=False)
        det_state_dict = torch.load(
            DOCTR_DETECT_MODEL_PATH, map_location=COMPUTE_DEVICE, weights_only=True)
        det_model.load_state_dict(det_state_dict)
        det_model = det_model.to(COMPUTE_DEVICE)

        # Recognition model initialisation
        reco_model = crnn_vgg16_bn(
            pretrained=False, pretrained_backbone=False, vocab=custom_vocab)
        recog_state_dict = torch.load(
            DOCTR_RECOG_MODEL_PATH, map_location=COMPUTE_DEVICE, weights_only=True)
        reco_model.load_state_dict(recog_state_dict)
        reco_model = reco_model.to(COMPUTE_DEVICE)

        self.model = ocr_predictor(
            det_arch=det_model, reco_arch=reco_model, pretrained=False).to(COMPUTE_DEVICE)

2. After that I got around documentations and discussions and somewhere found another way for model initialization and did this:

class DocTR:
    def __init__(self):
        custom_vocab = CUSTOM_VOCAB

        # Detection model initialisation
        det_model = db_resnet50(pretrained=False, pretrained_backbone=False)
        det_model.from_pretrained(
            DOCTR_DETECT_MODEL_PATH, map_location=COMPUTE_DEVICE, weights_only=True)

        # Recognition model initialisation
        reco_model = crnn_vgg16_bn(
            pretrained=False, pretrained_backbone=False, vocab=custom_vocab)
        reco_model.from_pretrained(
            DOCTR_RECOG_MODEL_PATH, map_location=COMPUTE_DEVICE, weights_only=True)

        self.model = ocr_predictor(
            det_arch=det_model, reco_arch=reco_model, pretrained=False).to(COMPUTE_DEVICE)

The COMPUTE_DEVICE in both the case is set to cuda and I am running it on a decive with NVIDIA graphic card which has cuda.

Tell me if what I am doing is correct or any changes has to be made? Or is the time taken is expected to be this long?

[felixdittrich92]

Hi @Vishnu280412 👋

This class is initialized one time in a fastapi lifespan so before the API is up ?

Otherwise it will be re-initialized every time.

Additional you can try to run with half precision or try if the compiled models are faster (https://mindee.github.io/doctr/using_doctr/using_model_export.html#compiling-your-models-pytorch-only)

        self.model = ocr_predictor(
            det_arch=det_model, reco_arch=reco_model, pretrained=False).to(COMPUTE_DEVICE).half()

[Vishnu280412]

Hello @felixdittrich92 👋,
Thanks for the response, and yes it is being initialised only once throughout the fastapi lifespan and I am using that object to run the model. So I wanted to know if I am loading the model on GPU correctly. So if doing this is enough:

reco_model = crnn_vgg16_bn(
            pretrained=False, pretrained_backbone=False, vocab=custom_vocab)
reco_model.from_pretrained(
            DOCTR_RECOG_MODEL_PATH, map_location=COMPUTE_DEVICE, weights_only=True)
self.model = ocr_predictor(
            det_arch=det_model, reco_arch=reco_model, pretrained=False).to(COMPUTE_DEVICE).half()

or I'll also have to all .to(COMPUTE_DEVICE to both det_model and reco_model?

I'll try to run with half precision and also try compiling the models and using it now.

Thank you.

[felixdittrich92]

It's correct ocr_predictor is enough 👍