Huge GPU memory usage of the Color Histogram Loss

[FunkyKoki]

Hi, thanks for your great work, it makes sense very much.

When I was going to apply the Color Histogram Loss in my work, I found that this loss occupies a lot of CUDA memory, is this normal?

For example, the model I used to train can be trained with batchsize = 16 and GPU_num = 4, but when I use this loss, the model can only be trained with batchsize = 8 and GPU_num = 8.

😭

[mahmoudnafifi]

Hi, thanks for your question. There are a couple of options to solve this. First, the RGBuvHistBlock class has some useful options that can reduce memory use.

You can reduce the size of images before computing the histogram loss using insz. Also, you can reduce the number of bins in the histogram feature using h. Another option is green_only which computes a single channel histogram feature instead of the 3-channel one. Please note that for now, the green_only option is available only on the
Colab example.

I am going to add it to the code in the next version. Also, please note that if you are going to reduce the number of bins, you may need to change the histogram boundary (from hist_boundary) to avoid compressing important bins (which are often located around the center of histogram features), so probably you could use h=32 with hist_boundary=[-2, 2]. Again, this hist_boundary option is only available on the Colab example for now and will be added to the next code version.

Finally, I just added other versions of the histogram class to compute histogram loss for other color spaces, like rg-chroma and Lab (for Lab, the input is supposed to be already in CIE LAB space before computing the loss).

I hope this helps.

[mahmoudnafifi]

The code of the other histogram classes is now provided in ./histogram_classes. Please check the readme for more information on how to switch from the RGB-uv histogram to rg-chroma or Lab histograms.

@FunkyKoki I am going to close this issue. Feel free to reopen it if you are still facing an issue with the GPU memory.