search

ZhengPeng7 / BiRefNet

[CAAI AIR'24] Bilateral Reference for High-Resolution Dichotomous Image Segmentation
https://www.birefnet.top
MIT License
1.1k stars 84 forks source link

Train with my custom large dataset #98

Open yaju1234 opened 4 days ago

yaju1234 commented 4 days ago

I have a large dataset of 300K images with different objects,car,human, animals, etc. Should I need to train from scratch or I can do transfer learning from your general pre-trained model?

ZhengPeng7 commented 4 days ago

300K is huge enough to set up the training from scratch. I recommend you do that to achieve better performance in your cases. But you can also first do an inference on my model to take a look at how the performance is, to provide a comparison with the results you obtain on your custom data.

yaju1234 commented 3 days ago

Thank you for your valuable response. If I were to train the model from scratch, how many epochs would I need, and how long would the training take if I use an A100 80 GPU?

ZhengPeng7 commented 3 days ago

Hi, you can refer to the training log of BiRefNet on DIS5K with 8 A100-80G GPUs, where the training set has 3,000 images.

Since your training set is so large and you have only one GPU, I suggest you save the checkpoints from the first epoch to check the performance of each of them. You can keep it training and copy the checkpoint of epoch-3 to another place to conduct inference and evaluation, which should already have some decent results.