diff --git a/configs/train_decoder_config.example.json b/configs/train_decoder_config.example.json new file mode 100644 index 0000000..dd387ed --- /dev/null +++ b/configs/train_decoder_config.example.json @@ -0,0 +1,99 @@ +{ + "unets": [ + { + "dim": 128, + "image_embed_dim": 768, + "cond_dim": 64, + "channels": 3, + "dim_mults": [1, 2, 4, 8], + "attn_dim_head": 32, + "attn_heads": 16 + } + ], + "decoder": { + "image_sizes": [64], + "channels": 3, + "timesteps": 1000, + "loss_type": "l2", + "beta_schedule": "cosine", + "learned_variance": true + }, + "data": { + "webdataset_base_url": "pipe:s3cmd get s3://bucket/path/{}.tar -", + "embeddings_url": "s3://bucket/embeddings/path/", + "num_workers": 4, + "batch_size": 64, + "start_shard": 0, + "end_shard": 9999999, + "shard_width": 6, + "index_width": 4, + "splits": { + "train": 0.75, + "val": 0.15, + "test": 0.1 + }, + "shuffle_train": true, + "resample_train": false, + "preprocessing": { + "RandomResizedCrop": { + "size": [128, 128], + "scale": [0.75, 1.0], + "ratio": [1.0, 1.0] + }, + "ToTensor": true + } + }, + "train": { + "epochs": 20, + "lr": 1e-4, + "wd": 0.01, + "max_grad_norm": 0.5, + "save_every_n_samples": 100000, + "n_sample_images": 6, + "device": "cuda:0", + "epoch_samples": null, + "validation_samples": null, + "use_ema": true, + "ema_beta": 0.99, + "amp": false, + "save_all": false, + "save_latest": true, + "save_best": true, + "unet_training_mask": [true] + }, + "evaluate": { + "n_evaluation_samples": 1000, + "FID": { + "feature": 64 + }, + "IS": { + "feature": 64, + "splits": 10 + }, + "KID": { + "feature": 64, + "subset_size": 10 + }, + "LPIPS": { + "net_type": "vgg", + "reduction": "mean" + } + }, + "tracker": { + "tracker_type": "console", + "data_path": "./models", + + "wandb_entity": "", + "wandb_project": "", + + "verbose": false + }, + "load": { + "source": null, + + "run_path": "", + "file_path": "", + + "resume": false + } +}