google-research / pix2seq

Pix2Seq codebase: multi-tasks with generative modeling (autoregressive and diffusion)
Apache License 2.0
857 stars 71 forks source link

Training from scratch on local MSCOCO data on A100 #6

Closed dipendra009 closed 2 years ago

dipendra009 commented 2 years ago

Hi, I can not find any resource on how to train this model on local available MSCOCO dataset using GPU. Could you please provide some training guide for that? I know the code is for running on TPUs, what will be the changes required to train on GPUs?

Regards, Dipendra

chentingpc commented 2 years ago

The code works for both TPUs and GPUs in pretty similar ways. The instructions given in the README is for GPU (in finetuning setup). For training from scratch on COCO, please refer to this issue #3. Hope this helps.

dipendra009 commented 2 years ago

Thank you for your reply. Could you please also shed some light on how to train on a custom dataset available locally when training from scratch? I have MSCOCO dataset available on my local disk and I tried running the training by setting config.data_dir argument to the /datasets/coco - which contains the entire file but it failed.

chentingpc commented 2 years ago

If you only train on COCO dataset, the tensorflow_datasets module will download COCO dataset automatically, and you only need to additionally download the annotation file as instructed in README. If you have other datasets that are not in tensorflow_datasets, you will need to write a wrapper class using tf.Data which implements extract method as shown here.

bilalorhanlr commented 2 years ago

Thank you for your reply. Could you please also shed some light on how to train on a custom dataset available locally when training from scratch? I have MSCOCO dataset available on my local disk and I tried running the training by setting config.data_dir argument to the /datasets/coco - which contains the entire file but it failed.

Hello there, did you train on your custom dataset ?