update available models

run.py

@@ -1,8 +1,7 @@
 import os
-from clize import run
 from glob import glob
 from subprocess import call
-
+import json
 def base():
     return {
         "slurm":{
@@ -34,7 +33,7 @@ def base():
             "save_ckpt_every": 1,
             "masked_mean": "",
             "resume": "",
-        }
+        },
     }
 def ddgan_cc12m_v2():
     cfg =  base()
@@ -69,7 +68,6 @@ def ddgan_cc12m_v9():
     cfg['model']['batch_size'] = 1
     return cfg
 
-
 def ddgan_cc12m_v11():
     cfg = base()
     cfg['model']['text_encoder'] = "google/t5-v1_1-large"    
@@ -77,22 +75,78 @@ def ddgan_cc12m_v11():
     cfg['model']['cross_attention'] = ""
     return cfg
 
-models = [
-    ddgan_cc12m_v2,
-    ddgan_cc12m_v6,
-    ddgan_cc12m_v7,
-    ddgan_cc12m_v8,
-    ddgan_cc12m_v9,
-    ddgan_cc12m_v11,
+def ddgan_cc12m_v12():
+    cfg = ddgan_cc12m_v11()
+    cfg['model']['text_encoder'] = "google/t5-v1_1-xl"    
+    cfg['model']['preprocessing'] = 'random_resized_crop_v1'
+    return cfg
+
+def ddgan_cc12m_v13():
+    cfg = ddgan_cc12m_v12()
+    cfg['model']['discr_type'] = "large_cond_attn"
+    return cfg
+
+def ddgan_cc12m_v14():
+    cfg = ddgan_cc12m_v12()
+    cfg['model']['num_channels_dae'] = 192
+    return cfg
+
+
+def ddgan_cifar10_cond17():
+    cfg = base()
+    cfg['model']['image_size'] = 32    
+    cfg['model']['classifier_free_guidance_proba'] = 0.2
+    cfg['model']['ch_mult'] = "1 2 2 2"
+    cfg['model']['cross_attention'] = ""
+    cfg['model']['dataset'] = "cifar10"
+    cfg['model']['n_mlp'] = 4
+    return cfg
 
+def ddgan_cifar10_cond18():
+    cfg = ddgan_cifar10_cond17()
+    cfg['model']['text_encoder'] = "google/t5-v1_1-xl"    
+    return cfg
+
+def ddgan_laion_aesthetic_v1():
+    cfg = ddgan_cc12m_v11()
+    cfg['model']['dataset_root'] = '"/p/scratch/ccstdl/cherti1/LAION-aesthetic/output/{00000..05038}.tar"'
+    return cfg
+
+def ddgan_laion_aesthetic_v2():
+    cfg = ddgan_laion_aesthetic_v1()
+    cfg['model']['discr_type'] = "large_cond_attn"
+    return cfg
+
+def ddgan_laion_aesthetic_v3():
+    cfg = ddgan_laion_aesthetic_v1()
+    cfg['model']['text_encoder'] = "google/t5-v1_1-xl" 
+    return cfg
+
+
+models = [
+    ddgan_cifar10_cond17, # cifar10, cross attn for discr
+    ddgan_cifar10_cond18, # cifar10, xl encoder
+    ddgan_cc12m_v2, # baseline (no large text encoder, no classifier guidance)
+    ddgan_cc12m_v6, # like v2 but using large T5 text encoder
+    ddgan_cc12m_v7, # like v2 but with classifier guidance
+    ddgan_cc12m_v8, # like v6 but classifier guidance
+    ddgan_cc12m_v9, # ~1B model but 64x64 resolution
+    ddgan_cc12m_v11, # large text encoder + cross attention + classifier free guidance
+    ddgan_cc12m_v12, # T5-XL + cross attention + classifier free guidance + random_resized_crop_v1
+    ddgan_cc12m_v13, # T5-XL + cross attention + classifier free guidance + random_resized_crop_v1 + cond attn
+    ddgan_cc12m_v14, # T5-XL + cross attention + classifier free guidance + random_resized_crop_v1 + 300M model
+    ddgan_laion_aesthetic_v1, # like ddgan_cc12m_v11 but fine-tuned on laion aesthetic
+    ddgan_laion_aesthetic_v2, # like ddgan_laion_aesthetic_v1 but trained from scratch with the new cross attn discr
+    ddgan_laion_aesthetic_v3, # like ddgan_laion_aesthetic_v1 but trained from scratch with T5-XL
 ]
+
 def get_model(model_name):
     for model in models:
         if model.__name__ == model_name:
             return model()
 
 
-def test(model_name, *, cond_text="", batch_size:int=None, epoch:int=None, guidance_scale:float=0, fid=False, real_img_dir=""):
+def test(model_name, *, cond_text="", batch_size:int=None, epoch:int=None, guidance_scale:float=0, fid=False, real_img_dir="", q=0.0, seed=0, nb_images_for_fid=0):
 
     cfg = get_model(model_name)
     model = cfg['model']
@@ -104,6 +158,7 @@ def test(model_name, *, cond_text="", batch_size:int=None, epoch:int=None, guida
     args = {}
     args['exp'] = model_name
     args['image_size'] = model['image_size']
+    args['seed'] = seed
     args['num_channels'] = model['num_channels']
     args['dataset'] = model['dataset']
     args['num_channels_dae'] = model['num_channels_dae']
@@ -116,12 +171,35 @@ def test(model_name, *, cond_text="", batch_size:int=None, epoch:int=None, guida
     args['text_encoder'] = model.get("text_encoder")
     args['cross_attention'] = model.get("cross_attention")
     args['guidance_scale'] = guidance_scale
+    args['masked_mean'] = model.get("masked_mean")
+    args['dynamic_thresholding_quantile'] = q
+    args['n_mlp'] = model.get("n_mlp")
 
     if fid:
         args['compute_fid'] = ''
         args['real_img_dir'] = real_img_dir 
-    cmd = "python test_ddgan.py " + " ".join(f"--{k} {v}" for k, v in args.items() if v is not None)
+        args['nb_images_for_fid'] = nb_images_for_fid
+    cmd = "python -u test_ddgan.py " + " ".join(f"--{k} {v}" for k, v in args.items() if v is not None)
     print(cmd)
     call(cmd, shell=True)
 
-run([test])
+def eval_results(model_name):
+    import pandas as pd
+    rows = []
+    cfg = get_model(model_name)
+    model = cfg['model']
+    paths = glob('./saved_info/dd_gan/{}/{}/fid*.json'.format(model["dataset"], model_name))
+    for path in paths:
+        with open(path, "r") as fd:
+            data = json.load(fd)
+        row = {}
+        row['fid'] = data['fid']
+        row['epoch'] = data['epoch_id']
+        rows.append(row)
+    out = './saved_info/dd_gan/{}/{}/fid.csv'.format(model["dataset"], model_name)
+    df = pd.DataFrame(rows)
+    df.to_csv(out, index=False)
+
+if __name__ == "__main__":
+    from clize import run
+    run([test, eval_results])