[ALBERT] How to know memory consumption on GPU/TPU of a model?

[ngoanpv]

Hi all, With parameters such as batch size, max sequence length,... How to know the number of megabyte/gigabyte/.. of TPU/GPU need to train the model?
Could you please give me an example? Thank you.

[ddouglas87]

I'm using a 980 (4GB) and in run_glue.sh run_task SST-2 1256 1e-5 20935 100 32 doesn't work for me but using a batch size of 16 instead run_task SST-2 1256 1e-5 20935 100 16 uses 3GB of ram and works. I get around 28 to 30 examples/sec.