theEricMa / OTAvatar

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

Pretrained model batchsize and gpus #29

Closed szh-bash closed 8 months ago

szh-bash commented 8 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 8 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.