Required GPU Memory for Training PixArt on Pokemon Dataset with 512 Image Size

[mnslarcher]

Hi,

I'm working on training PixArt from scratch and following the train.ipynb notebook. I've made only some small changes to adapt to the new code (like renaming files from 0.npy to images_0.npy etc.).

I'm on Paperspace Gradient using a P6000 GPU (24 GB) but running into OOM issues. Is this expected? What's the recommended memory to train on the Pokemon dataset with 512 image size?

Thanks,
Mario

[lawrence-cj]

You can try to change the line here to fp32_attention = True according to your config file.

[mnslarcher]

This was not enough, but setting

fp32_attention = False
train_batch_size = 4 # 32
eval_batch_size = 4
gradient_accumulation_steps = 8

it remains around 80% GPU usage.

Thanks!