mahmoudnafifi / HistoGAN

Reference code for the paper HistoGAN: Controlling Colors of GAN-Generated and Real Images via Color Histograms (CVPR 2021).
MIT License
267 stars 27 forks source link

Huge GPU memory usage of the Color Histogram Loss #8

Closed FunkyKoki closed 2 years ago

FunkyKoki commented 2 years ago

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 commented 2 years ago

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 commented 2 years ago

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.