idealo / imagededup

😎 Finding duplicate images made easy!
https://idealo.github.io/imagededup/
Apache License 2.0
5.17k stars 458 forks source link

Add a device parameter to perform encoding using different torch devices such as mps #189

Open CarloNicolini opened 1 year ago

CarloNicolini commented 1 year ago

Working on a MacBook M1Pro, I often enjoy the concrete speedup that the Metal Performance Shader (mps) device has over the cpu.

I typically see increases around 20x at inference time with standard convolutional neural networks on Pytorch 1.12, which is as easy as a .to("mps") for model and tensors.

I am wondering if you think it's possible to add it to the parameters of the CNN encoder for duplicate removal.

tanujjain commented 1 year ago

The current convention to select between cpu and gpu is through the use of the environment variable CUDA_VISIBLE_DEVICES . Instead of explicitly passing a device selection parameter, I would like to stay with this convention.

However, it maybe possible to switch to mps by default if pytorch detects an M1 chip on the system. Need to see what kind of edge cases this might lead to and the fact that one would lose the ability to use cpu if mps is present on the system. I can pick this up somewhere next month, would be useful to add this functionality somehow.

tanujjain commented 1 year ago

Tried adding mps but it seems some pytorch operations are not yet implemented for mps backend. As a result, CNN usage breaks. Will keep the issue open and keep checking if the corresponding operators are supported.