nyx-ai / stylegan2-flax-tpu

🖼 Training StyleGAN2 on TPUs in JAX
https://nyx-ai.github.io/stylegan2-flax-tpu
130 stars 11 forks source link

training on grayscale images #6

Open GeorvityLabs opened 2 years ago

GeorvityLabs commented 2 years ago

i have a dataset of grayscale images and wanted to train on it. do you have any suggestions on how to get this working with grayscale input images.

currently, when trying to create a tfrecord , i'm getting the following error : Screenshot from 2022-09-29 14-36-31

GeorvityLabs commented 2 years ago

Traceback (most recent call last): File "dataset_utils/images_to_tfrecords.py", line 144, in <module> images_to_tfrecords(args.image_dir, args.data_dir, args.has_labels) File "dataset_utils/images_to_tfrecords.py", line 115, in images_to_tfrecords channels = img.shape[2] IndexError: tuple index out of range

GeorvityLabs commented 2 years ago

@MasterScrat @mar-muel it would be great if you guys could make a training notebook as well , keeping the above issue in mind , that way users could train this model using kaggle TPU or colab

mar-muel commented 2 years ago

Hi @GeorvityLabs - Indeed the code is not in a very mature state.

The error you experience is because your images don't have 3 dimensions. You can either reshape your greyscale images to have a third dimension or just hardcode channels to 1 in the script. The training code should in theory support grey scale image training, however this was not yet tested.

GeorvityLabs commented 2 years ago

Hi @GeorvityLabs - Indeed the code is not in a very mature state.

The error you experience is because your images don't have 3 dimensions. You can either reshape your greyscale images to have a third dimension or just hardcode channels to 1 in the script. The training code should in theory support grey scale image training, however this was not yet tested.

I made a simple training notebook , so as to run on colab tpu, but there seems to be some issue with regards to connecting with the tpu. I first changed version on requirements.txt to match the needed version ,then i removed jax from the requirements.txt file. but none of them seem to work.

I've linked my notebook here , could you have a look and check if you are able to train on images at your end : https://colab.research.google.com/drive/1tzGI7hzZvXGWRh-Q9yz_me_WQugHzVZn?usp=sharing

@mar-muel

mar-muel commented 2 years ago

Hi @GeorvityLabs - we've just made a new release, which comes with an (experimental) Colab training notebook. However, Colab TPUs are an older generation TPUs compared to TPU VMs and training is fairly slow unfortunately :(

GeorvityLabs commented 2 years ago

Hi @GeorvityLabs - we've just made a new release, which comes with an (experimental) Colab training notebook. However, Colab TPUs are an older generation TPUs compared to TPU VMs and training is fairly slow unfortunately :(

Thanks for the update. I will test it out. For starters, how many input/training image samples would you recommend to use and for how many epochs would you recommend to train for..

mar-muel commented 2 years ago

Really depends on the subjects and resolution, but I would argue you need 5k examples for a good model at the very least. Regarding training time we noticed improvements even after several days of training (however marginal). From scratch training at 512 resolution took on the order of 2-3 days at least. You may also want to check out our W&B report.

GeorvityLabs commented 2 years ago

Really depends on the subjects and resolution, but I would argue you need 5k examples for a good model at the very least. Regarding training time we noticed improvements even after several days of training (however marginal). From scratch training at 512 resolution took on the order of 2-3 days at least. You may also want to check out our W&B report.

Thanks for the reply. Say I have 50 images only, what would you recommend in that case. Is it possible to train any generator with low sample size?

mar-muel commented 2 years ago

That's definitely too small. You may want to look for pre-trained models or models which cover your mode at least partially. Also check out LAION/img2dataset to get to a decent sized dataset.

GeorvityLabs commented 2 years ago

That's definitely too small. You may want to look for pre-trained models or models which cover your mode at least partially. Also check out LAION/img2dataset to get to a decent sized dataset.

Got it. Will have a look at that. For now, I will run a test on a 5k image dataset like you mentioned.