Add new kwarg to load model on CPU or GPU

[odulcy-mindee]

At the moment, I don't think there is a way to load models on CPU or GPU exclusively. One way is to set os.environ['CUDA_VISIBLE_DEVICES']. It would be nice have a parameter device to control explicitly this behavior.

Example:

from doctr.io import DocumentFile
from doctr.models import ocr_predictor

model = ocr_predictor(pretrained=True, device='cpu')

doc = DocumentFile.from_images('/path/to/img.jpg')

result = model(doc)

What do you think ?

[frgfm]

Hi @odulcy-mindee 👋

Yes it's indeed something that has been bothering me as well!
Perhaps we should make this a method so that it can be called at different moments. But thinking a bit more about this, I've had the problem with Tensorflow:

• on PyTorch backend, the models are all created on CPU by default, and the predictor does has the proper interface (similar to torch.nn.Module) to move the model to another device
• on TensorFlow though, the backend doesn't ask and move the model to GPU if it's available. So perhaps a quick win would be to change that

What do you think? :)

[odulcy-mindee]

@frgfm So,

• First step is to replicate the behavior from PyTorch for Tensorflow models, i.e. loading everything on CPU
• Second step is to create and propagate a kwarg to chose on which device we want to load the model (?)

For the latter, I think we need a function to standardize name for devices for Torch and Tensorflow before forwarding kwargs to internal methods.

[frgfm]

I agree with your first point.

For the second, if you mean passing kwargs down the call method of all modules, I disagree. In my opinion, considering moving from one device to the other require both time and resources (RAM), that should be done as a constructor or independent method without increasing latency for inference-related method. Here is a broad idea of what I have in mind:

# Method 1
predictor = ocr_predictor()
predictor.set_device("cuda:0")
result = predictor(input)
# Method 2
predictor = ocr_predictor(..., device="cuda:0")
result = predictor(input)

And I prefer the first since you might want a public method to switch the model to CPU or another GPU. Regarding standardizing the device name, we could easily borrow the ones from nvidia:

• cpu (default)
• cuda:<GPU_ID>

and we'd throw an error when the device is not accessible :)
How does it sound?

[odulcy-mindee]

@frgfm I thought to pass the kwarg at model initialization, i.e. in __init__ and not in __call__. I'm not sure if I fully understood your point 😅.

Considering your example, I think we can implement both methods. Personaly, I prefer Method 2 because your PC or GPU may lack of memory so it could make impossible to load a model on GPU if you need to instantiate it first on CPU then move it on GPU.