ss_vae_conv3d_16l8_fp16.json

{
    "models": {
        "encoder": {
            "name": "SparseStructureEncoder",
            "args": {
                "in_channels": 1,
                "latent_channels": 8,
                "num_res_blocks": 2,
                "num_res_blocks_middle": 2,
                "channels": [32, 128, 512],
                "use_fp16": true
            }
        },
        "decoder": {
            "name": "SparseStructureDecoder",
            "args": {
                "out_channels": 1,
                "latent_channels": 8,
                "num_res_blocks": 2,
                "num_res_blocks_middle": 2,
                "channels": [512, 128, 32],
                "use_fp16": true
            }
        }
    },
    "dataset": {
        "name": "SparseStructure",
        "args": {
            "resolution": 64,
            "min_aesthetic_score": 4.5
        }
    },
    "trainer": {
        "name": "SparseStructureVaeTrainer",
        "args": {
            "max_steps": 1000000,
            "batch_size_per_gpu": 4,
            "batch_split": 1,
            "optimizer": {
                "name": "AdamW",
                "args": {
                    "lr": 1e-4,
                    "weight_decay": 0.0
                }
            },
            "ema_rate": [
                0.9999
            ],
            "fp16_mode": "inflat_all",
            "fp16_scale_growth": 0.001,
            "grad_clip": {
                "name": "AdaptiveGradClipper",
                "args": {
                    "max_norm": 1.0,
                    "clip_percentile": 95
                }
            },
            "i_log": 500,
            "i_sample": 10000,
            "i_save": 10000,
            "loss_type": "dice",
            "lambda_kl": 0.001
        }
    }
}