training_config.json

{
  "name": "multitrace_unet_cifar10_15m",
  "architecture": "ConditionalMultiTraceUNet",
  "objective": "multi_trace_imf",
  "seed": 42,
  "dataset": "cifar10",
  "data_root": "/mnt/imf-training/data/cifar10",
  "image_size": 32,
  "num_workers": 2,
  "lr": 0.0001,
  "min_lr_ratio": 0.1,
  "betas": [
    0.9,
    0.99
  ],
  "weight_decay": 0.0,
  "grad_clip": 1.0,
  "mixed_precision": "bf16",
  "channels_last": true,
  "allow_tf32": true,
  "fused_optimizer": true,
  "preload_gpu": true,
  "log_step": 50,
  "sample_step": 1000,
  "checkpoint_step": 1000,
  "keep_last_checkpoints": 3,
  "sample_nrow": 10,
  "parameter_target": 15000000,
  "parameter_tolerance": 0.01,
  "model": {
    "input_size": 32,
    "in_channels": 3,
    "num_classes": 10,
    "base_channels": 69,
    "channel_mult": [
      1,
      2,
      2,
      4
    ],
    "num_res_blocks": 1,
    "cond_dim": 312,
    "fourier_dim": 64,
    "attn_resolutions": [
      8
    ]
  },
  "multi_trace_imf": {
    "lambda_diag": 1.0,
    "max_nfe": 5,
    "jvp_precision": "bf16",
    "channels_last": true
  },
  "monitoring": {
    "enabled": true,
    "tensorboard": true,
    "save_best": true,
    "validate_raw": true,
    "plot_every": 500,
    "validation_every": 1000,
    "validation_batches": 5,
    "sample_n_per_class": 2
  },
  "evaluation": {
    "num_generated": 10000,
    "batch_size": 128,
    "nfes": [
      1,
      2,
      3,
      4,
      5
    ],
    "seed": 2026
  },
  "batch_size": 1024,
  "gradient_accumulation_steps": 1,
  "n_steps": 50000,
  "warmup_steps": 2500,
  "ema_decay": 0.9998000100000001,
  "source_manifest": {
    "repo_url": "https://github.com/namdeptraivcd/IMF.git",
    "git_commit": "59d5164a2b20753b846568264ebb10c3e49beace",
    "files": {
      "train.py": "34f72a2bbba5341aeda9cd3fe021e18ed6bf1bbc73b823187a97ed83e03cc3aa",
      "monitoring.py": "5315ee70585789a684162fbd4fa3a8f2479d5186d282d22a3fc314432b3100a9",
      "hub.py": "ca23b68ef728ea2ae514918a09bb678d80df303fc97949593f2c6607cdc0fa03",
      "sample.py": "90f1089547475b5ab1d8f5f8d160a33c639826eda95c6ca79ee2d3f527b24be1",
      "models/dit.py": "843080c2702d596dd3b16e35b78501b168ca008d013d4394ee9c64c086361819",
      "models/__init__.py": "1d9d5ec9cc26f0e76e570735a42e420314fdccc0a4aa113ddac642fd4d147c36",
      "trace_imf.py": "43e2e8105541257376f648f3ee82a92cce8402530be59372f8d7ed0516501263",
      "data.py": "a346a102fac83a40823501b68cf8055ac385ab3837cb6ef62b9212bdc7b9dffc",
      "configs/__init__.py": "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855",
      "configs/cifar10_22m.py": "1599fb09bab2ac556f46f434f82d22d61aa22901140679b4bfef62aa65245588",
      "requirements-notebook.txt": "858b7587f8c9b4778b2f13955ce8a00ff1f9a8f1cbfabb0ce94c2e97c55cc5fc",
      "models/unet.py": "57baed6e4ef61ca23aa0ddf62ec0c30b85e6d27f2504416269bb6d1245ee22df",
      "multi_trace_imf.py": "a807809500be68da9d33f13fd528d987e61471161ad869f97c7411bb225fe5b8",
      "objectives.py": "691d152ee1f0970fa9c279c596d09070c5a5c6618ac465bb462fac0188796f7d",
      "ema.py": "1bc020387ec0020724d0a5c066180100f08204102fd79e804c3624203ebc4579",
      "evaluation.py": "6f2d19263ea1b8c13c9fdc4678062adae99ea0773e0f03c8001f81a598f03238",
      "configs/multitrace_cifar10_22m.py": "e0a8fc408170b618d88d19ccf71993af5f469e59c36c31dd38d920b57ee727bb",
      "configs/multitrace_cifar10_15m.py": "96863dc436dce0accb6a2722b4398aa1a8a6ebf4e96386511767fe9b3c2db69d"
    }
  }
}