Need tuning suggestions

Hi all, I’m trying to train a model with ~1K and ~5K images. So far model is collapsing, or this is what I understand from output.

I tried first with 1K , following parameters and 150K epochs: stylegan2_pytorch --aug_prob 0.25 --attn_layers [1,2]

And got:

Then I tried with ~5K, following parameters and 150K epochs: stylegan2_pytorch --aug_prob 0.25 --attn_layers [1,2]

And got:

What do you suggest to do to improve results? Cleaning dataset? There are a lot of parameters to try and by reading other threads cl_reg and trunc_psi are suggested but by using them, the images start to get worse. Right now I’m trying to reduce batch_size to 3 and increase network_capacity to 32

These are the parameters by default I think I could tune:

data = './data' 
results_dir = './results'
models_dir = './models'
name = 'default'
new = False
load_from = -1
image_size = 128
network_capacity = 16
transparent = False
batch_size = 5 
gradient_accumulate_every = 6
num_train_steps = 150000
learning_rate = 2e-4
lr_mlp = 0.1
ttur_mult = 1.5
num_workers =  None
save_every = 1000
generate = False
generate_interpolation = False
save_frames = False
num_image_tiles = 8
trunc_psi = 0.75
fp16 = False
cl_reg = False
fq_layers = []
fq_dict_size = 256
attn_layers = []
no_const = False
aug_prob = 0. 
dataset_aug_prob = 0. 

[lucidrains]

Please try https://github.com/lucidrains/unet-stylegan2 ! And yes, up the network capacity and make sure the batch size x gradient accumulate is at least 32

[lucidrains]

yup, keep attention and aug prob! dataset aug prob augments the images immediately when they are fetched from the file system, and not when it is fed to the discriminator. the former is not differentiable while the latter is