Multiple GPU Training Support

[udasinnayan]

In the config file trainer part, if I set auto_select_gpus: true and devices: 2, I got an error like
AttributeError: Can't pickle local object 'FeatureExtractor.get_features.<locals>.hook', it comes from this line trainer.fit(model=model, datamodule=datamodule) in anomalib/tools/train.py script but for single GPU it works without an error. Any comments about this error / is it possible to use multiple GPUs for training and testing?

[ashwinvaidya17]

Hi @udasinnayan that's a known issue. This might probably be solved when me move towards torchfx but the only option is to use single GPU for now.

[samet-akcay]

@ashwinvaidya17, I'm wondering if this is the case even when we use timm's feature extractor

[ashwinvaidya17]

I tested this on the current main branch and it still does not work with timm's feature extractor. I am closing as it will require a significant change.

[ashwinvaidya17]

Multi-gpu training would be a nice feature to have however the horizon to implement it is going to be longer so I've converted this into a feature request.

[ashwinvaidya17]

Upon further investigation on the padim model I found that if 1. the metrics are not moved to cpu in metrics configuration callback and base anomaly module, and 2. the embeddings are left on the gpu in on_validation_start, then the model trains on multiple gpus. A change which fixes this depends on solving existing OOM issues in a few model + dataset combinations.

[zyy0877]

Have you resolved your issue with multi GPU training

[samet-akcay]

Closing this as it is currently being addressed here #1449