add some new configs & whatnot

Author: vedantroy

config_model/learned_variance_256.yaml

@@ -0,0 +1,13 @@
+unet:
+  in_channels: 3
+  # 6 channels b/c need to output learned variance interpolation vector
+  out_channels: 6
+  model_channels: 256
+  channel_mult: [1, 2, 3, 4]
+  layer_attn: [false, false, true, true]
+  res_blocks: 3
+  attention_heads: 4
+diffusion:
+  schedule: cosine
+  steps: 1000
+  learn_sigma: true

config_run/1xA100-80GB-learned.yaml

@@ -6,6 +6,7 @@ warmup_batches: 500
 # b/c I suspect it will increase convergence
 # significantly
 lr: 1.00*1e-4
+cosine_factor: 1.2
 
 batch_size: 302
 micro_batches: 1

config_run/1xA100-80GB.bash

@@ -17,6 +17,7 @@ class RunConfig(hp.Hparams):
 
     warmup_batches: int = hp.required("Warmup batches")
     lr: str = hp.required("Learning rate")
+    cosine_factor: float = hp.required("# the factor to use for the cosine scheduler")
 
     batch_size: int = hp.required("Batch size")
     micro_batches: int = hp.required("Micro-batches")
@@ -27,6 +28,7 @@ class RunConfig(hp.Hparams):
     evals: int = hp.required("# Evals")
 
 
+
 def main(
     model_config_file: Path = typer.Option(...),
     run_config_file: Path = typer.Option(...),
@@ -62,7 +64,7 @@ def get_num(expr, typ):
             batch_rate=c.batch_rate,
             target_time=target_time,
             warmup=c.warmup_batches,
-            cosine_factor=1.4,
+            cosine_factor=c.cosine_factor,
         )
 
         train_dl = dataloader(dir_train, c.batch_size, workers=8)