Upload hps.py

hps.py

@@ -0,0 +1,300 @@
+import argparse
+
+HPARAMS_REGISTRY = {}
+
+
+class Hparams:
+    def update(self, dict):
+        for k, v in dict.items():
+            setattr(self, k, v)
+            
+brset = Hparams()
+brset.lr = 1e-3
+brset.bs = 16
+brset.wd = 0.01
+brset.z_dim = 16
+brset.input_res = 384 #192
+brset.pad = 9
+brset.hflip = 0.5
+
+brset.input_channels = 3
+# the first number is never used, it is just a placeholder to know the expected dimension of the output
+# b is the number of convolutional blocks, so for example 32b3d2 means 3 convolutional blocks
+# d is used to create a downsampling layer (represented as projection layer, and a 2D average pooling layer), so 32b3d2 means that we will add a 2D average pooling layer block with a stride and and kernel size of 2, at the end of the 3 convolutional blocks
+# The widths are the number of channels of each convolutional block
+#brset.enc_arch = "384b1d4,96b3d2,48b7d2,24b11d2,12b7d2,6b3d6,1b2" # Also for 384 but requires more memory 
+#brset.dec_arch = "1b2,6b4,12b8,24b12,48b8,96b4,384b2" # Also for 384 but requires more memory
+brset.enc_arch = "384b1d4,96b3d4,24b11d2,12b7d2,6b3d6,1b2" # for 384
+brset.dec_arch = "1b2,6b4,12b8,24b12,96b4,384b2" # for 384
+brset.widths = [32, 64, 128, 160, 192, 512] # for 384
+#brset.enc_arch = "192b1d2,96b3d2,48b7d2,24b11d2,12b7d2,6b3d6,1b2" # for 192
+#brset.dec_arch = "1b2,6b4,12b8,24b12,48b8,96b4,192b2" # for 192
+#brset.widths = [32, 64, 96, 128, 160, 192, 512] # for 192
+brset.bias_max_res = 64 # Used for the max resolution of the bias parameter
+brset.bottleneck = 4 # Used for the number of channels of the bottleneck layer in the block = width/bottleneck
+brset.parents_x = ['patient_age', 'patient_sex', 'DR_ICDR']
+brset.context_norm = "[-1,1]"
+brset.context_dim = 7 # Number of context variables. In our case it is 7 because we have age (1 - Continuous), sex (1 - Binary) and DR_ICDR (5 one-hot encoded)
+brset.n_classes = 5
+brset.concat_pa = True
+HPARAMS_REGISTRY["brset"] = brset
+
+
+morphomnist = Hparams()
+morphomnist.lr = 1e-3
+morphomnist.bs = 32
+morphomnist.wd = 0.01
+morphomnist.z_dim = 16
+morphomnist.input_res = 32
+morphomnist.pad = 4
+morphomnist.enc_arch = "32b3d2,16b3d2,8b3d2,4b3d4,1b4"
+morphomnist.dec_arch = "1b4,4b4,8b4,16b4,32b4"
+morphomnist.widths = [16, 32, 64, 128, 256]
+morphomnist.parents_x = ["thickness", "intensity", "digit"]
+morphomnist.concat_pa = True
+morphomnist.context_norm = "[-1,1]"
+morphomnist.context_dim = 12
+HPARAMS_REGISTRY["morphomnist"] = morphomnist
+
+
+cmnist = Hparams()
+cmnist.lr = 1e-3
+cmnist.bs = 32
+cmnist.wd = 0.01
+cmnist.z_dim = 16
+cmnist.input_res = 32
+cmnist.input_channels = 3
+cmnist.pad = 4
+cmnist.enc_arch = "32b3d2,16b3d2,8b3d2,4b3d4,1b4"
+cmnist.dec_arch = "1b4,4b4,8b4,16b4,32b4"
+cmnist.widths = [16, 32, 64, 128, 256]
+cmnist.parents_x = ["digit", "colour"]
+cmnist.context_dim = 20
+HPARAMS_REGISTRY["cmnist"] = cmnist
+
+
+ukbb64 = Hparams()
+ukbb64.lr = 1e-3
+ukbb64.bs = 32
+ukbb64.wd = 0.1
+ukbb64.z_dim = 16
+ukbb64.input_res = 64
+ukbb64.pad = 3
+ukbb64.enc_arch = "64b3d2,32b31d2,16b15d2,8b7d2,4b3d4,1b2"
+ukbb64.dec_arch = "1b2,4b4,8b8,16b16,32b32,64b4"
+ukbb64.widths = [32, 64, 128, 256, 512, 1024]
+HPARAMS_REGISTRY["ukbb64"] = ukbb64
+
+
+ukbb192 = Hparams()
+ukbb192.update(ukbb64.__dict__)
+ukbb192.input_res = 384
+ukbb192.pad = 9
+ukbb192.enc_arch = "384b2d2,192b2d2,96b3d2,48b7d2,24b11d2,12b7d2,6b3d6,1b2"
+ukbb192.dec_arch = "1b2,6b4,12b8,24b12,48b8,96b4,192b2,384b2"
+ukbb192.widths = [32, 64, 96, 128, 160, 192, 512, 1024]
+HPARAMS_REGISTRY["ukbb192"] = ukbb192
+
+
+mimic192 = Hparams()
+mimic192.lr = 1e-3
+mimic192.bs = 16
+mimic192.wd = 0.1
+mimic192.z_dim = 16
+mimic192.input_res = 192
+mimic192.pad = 9
+mimic192.enc_arch = "192b1d2,96b3d2,48b7d2,24b11d2,12b7d2,6b3d6,1b2"
+mimic192.dec_arch = "1b2,6b4,12b8,24b12,48b8,96b4,192b2"
+mimic192.widths = [32, 64, 96, 128, 160, 192, 512]
+HPARAMS_REGISTRY["mimic192"] = mimic192
+
+mimic384 = Hparams()
+mimic384.lr = 1e-3
+mimic384.bs = 16
+mimic384.wd = 0.1
+mimic384.z_dim = 16
+mimic384.input_res = 384
+mimic384.pad = 9
+mimic384.enc_arch = "384b1d2,192b1d2,96b3d2,48b7d2,24b11d2,12b7d2,6b3d6,1b2"
+mimic384.dec_arch = "1b2,6b4,12b8,24b12,48b8,96b4,192b2,384b2"
+mimic384.widths = [32, 64, 96, 128, 160, 192, 512,1024]
+HPARAMS_REGISTRY["mimic384"] = mimic384
+
+def setup_hparams(parser: argparse.ArgumentParser) -> Hparams:
+    hparams = Hparams()
+    args = parser.parse_known_args()[0]
+    valid_args = set(args.__dict__.keys())
+    hparams_dict = HPARAMS_REGISTRY[args.hps].__dict__
+    for k in hparams_dict.keys():
+        if k not in valid_args:
+            raise ValueError(f"{k} not in default args")
+    parser.set_defaults(**hparams_dict)
+    hparams.update(parser.parse_known_args()[0].__dict__)
+    return hparams
+
+
+def add_arguments(parser: argparse.ArgumentParser):
+    parser.add_argument("--exp_name", help="Experiment name.", type=str, default="")
+    parser.add_argument(
+        "--data_dir", help="Data directory to load form.", type=str, default=""
+    )
+    parser.add_argument("--hps", help="hyperparam set.", type=str, default="ukbb64")
+    parser.add_argument(
+        "--resume", help="Path to load checkpoint.", type=str, default=""
+    )
+    parser.add_argument("--seed", help="Set random seed.", type=int, default=7)
+    parser.add_argument(
+        "--deterministic",
+        help="Toggle cudNN determinism.",
+        action="store_true",
+        default=False,
+    )
+    # training
+    parser.add_argument("--epochs", help="Training epochs.", type=int, default=5000)
+    parser.add_argument("--bs", help="Batch size.", type=int, default=32)
+    parser.add_argument("--lr", help="Learning rate.", type=float, default=1e-3)
+    parser.add_argument(
+        "--lr_warmup_steps", help="lr warmup steps.", type=int, default=100
+    )
+    parser.add_argument("--wd", help="Weight decay penalty.", type=float, default=0.01)
+    parser.add_argument(
+        "--betas",
+        help="Adam beta parameters.",
+        nargs="+",
+        type=float,
+        default=[0.9, 0.9],
+    )
+    parser.add_argument(
+        "--ema_rate", help="Exp. moving avg. model rate.", type=float, default=0.999
+    )
+    parser.add_argument(
+        "--input_res", help="Input image crop resolution.", type=int, default=64
+    )
+    parser.add_argument(
+        "--input_channels", help="Input image num channels.", type=int, default=1
+    )
+    parser.add_argument("--pad", help="Input padding.", type=int, default=3)
+    parser.add_argument(
+        "--hflip", help="Horizontal flip prob.", type=float, default=0.5
+    )
+    parser.add_argument(
+        "--grad_clip", help="Gradient clipping value.", type=float, default=350
+    )
+    parser.add_argument(
+        "--grad_skip", help="Skip update grad norm threshold.", type=float, default=500
+    )
+    parser.add_argument(
+        "--accu_steps", help="Gradient accumulation steps.", type=int, default=1
+    )
+    parser.add_argument(
+        "--beta", help="Max KL beta penalty weight.", type=float, default=1.0
+    )
+    parser.add_argument(
+        "--beta_warmup_steps", help="KL beta penalty warmup steps.", type=int, default=0
+    )
+    parser.add_argument(
+        "--kl_free_bits", help="KL min free bits constraint.", type=float, default=0.0
+    )
+    parser.add_argument(
+        "--viz_freq", help="Steps per visualisation.", type=int, default=10000
+    )
+    parser.add_argument(
+        "--eval_freq", help="Train epochs per validation.", type=int, default=5
+    )
+    parser.add_argument(
+        "--n_classes", help="Number of classes for DR ICDR.", type=int, default=10
+    )
+    
+    # model
+    parser.add_argument(
+        "--vae",
+        help="VAE model: simple/hierarchical.",
+        type=str,
+        default="hierarchical",
+    )
+    parser.add_argument(
+        "--enc_arch",
+        help="Encoder architecture config.",
+        type=str,
+        default="64b1d2,32b1d2,16b1d2,8b1d8,1b2",
+    )
+    parser.add_argument(
+        "--dec_arch",
+        help="Decoder architecture config.",
+        type=str,
+        default="1b2,8b2,16b2,32b2,64b2",
+    )
+    parser.add_argument(
+        "--cond_prior",
+        help="Use a conditional prior.",
+        action="store_true",
+        default=False,
+    )
+    parser.add_argument(
+        "--widths",
+        help="Number of channels.",
+        nargs="+",
+        type=int,
+        default=[16, 32, 48, 64, 128],
+    )
+    parser.add_argument(
+        "--bottleneck", help="Bottleneck width factor.", type=int, default=4
+    )
+    parser.add_argument(
+        "--z_dim", help="Numver of latent channel dims.", type=int, default=16
+    )
+    parser.add_argument(
+        "--z_max_res",
+        help="Max resolution of stochastic z layers.",
+        type=int,
+        default=192,
+    )
+    parser.add_argument(
+        "--bias_max_res",
+        help="Learned bias param max resolution.",
+        type=int,
+        default=64,
+    )
+    parser.add_argument(
+        "--x_like",
+        help="x likelihood: {fixed/shared/diag}_{gauss/dgauss}.",
+        type=str,
+        default="diag_dgauss",
+    )
+    parser.add_argument(
+        "--std_init",
+        help="Initial std for x scale. 0 is random.",
+        type=float,
+        default=0.0,
+    )
+    parser.add_argument(
+        "--parents_x",
+        help="Parents of x to condition on.",
+        nargs="+",
+        default=["mri_seq", "brain_volume", "ventricle_volume", "sex"],
+    )
+    parser.add_argument(
+        "--concat_pa",
+        help="Whether to concatenate parents_x.",
+        action="store_true",
+        default=False,
+    )
+    parser.add_argument(
+        "--context_dim",
+        help="Num context variables conditioned on.",
+        type=int,
+        default=4,
+    )
+    parser.add_argument(
+        "--context_norm",
+        help='Conditioning normalisation {"[-1,1]"/"[0,1]"/log_standard}.',
+        type=str,
+        default="log_standard",
+    )
+    parser.add_argument(
+        "--q_correction",
+        help="Use posterior correction.",
+        action="store_true",
+        default=False,
+    )
+    return parser