theEricMa / OTAvatar

This is the official repository for OTAvatar: One-shot Talking Face Avatar with Controllable Tri-plane Rendering [CVPR2023].
312 stars 39 forks source link

Pretrained model batchsize and gpus #29

Closed szh-bash closed 10 months ago

szh-bash commented 10 months ago
          Hi, may I know how much GPU is used for the training? Mine is 4 A100s (80GB mem), so the batchsize is 8(per GPU) * 4 (GPU num) = 32, therefore the 2000 iters will spend more than 1 epoch. If you cannot support batchsize=8 per GPU, please try more GPUs. Larger batchsize leads to more stable training.  

Originally posted by @theEricMa in https://github.com/theEricMa/OTAvatar/issues/10#issuecomment-1537972669

I trained with 4(per GPU) 6(GPU num), 1500 iters spent exactlly 1 epoch. image The pretrained model named epoch_00005_iteration_000002000, maybe you trained this model with more than 8(per GPU) 8(GPU num)?

theEricMa commented 10 months ago

Hello, thanks for the feedback. The training configuration has been optimized for typical GPUs with 24-32GB memory. However, we used 80GB A100 GPUs which allowed much larger batch sizes. Feel free to modify the configs as needed to fit your hardware.