Cuda out of memory when train with TRADES On resnet50.

[yfreedomliTHU]

Hi, @yaodongyu , I'm very interested in your work at ICML'19, and I attempt to use it in the competition. I tried to train resnet50 with trades_loss but there was an Error alert:CUDA out of memory. I wonder if trades_loss needs more CUDA memory.

I trained the model resnet50 on NVIDIA 1080ti with cross entropy loss, and the batch size can be set to 128.However, when I trained with trades_loss, it raise an error"CUDA out of memory" with batch_size 16.
I'm not sure whether there's a problem with my code, or trades_loss needs more CUDA memory.

thank you!

[yaodongyu]

Thanks for your interests in our paper.

In trades_loss() ，

TRADES/trades.py

Line 71 in a2183f3

loss = loss_natural + beta * loss_robust

the loss includes two terms, loss_natural(i.e., cross-entropy loss) and loss_robust. This means that trades_loss() needs more memory compared with natural training.

Btw, as you have already changed the batch size from 128 to 16, you should be able to train with trades_loss(). Maybe you can check your code and figure it out.