| 153 | |
| 154 | # define and parse the arguments |
| 155 | def parse_args(): |
| 156 | parser = argparse.ArgumentParser() |
| 157 | |
| 158 | # processes to choose from. Either diffusion, pdmp, or 'nf' to use a normal normalizing flow. |
| 159 | parser.add_argument("--method", help='generative method to use', default=None, type=str) |
| 160 | |
| 161 | # EXPERIMENT parameters, specific to TRAINING |
| 162 | parser.add_argument("--config", help='config file to use', type=str, required=True) |
| 163 | parser.add_argument("--name", help='name of the experiment. Defines save location: ./models/name/', type=str, required=True) |
| 164 | parser.add_argument('--epochs', help='epochs', default=None, type = int) |
| 165 | parser.add_argument('-r', "--resume", help="resume existing experiment", action='store_true', default=False) |
| 166 | parser.add_argument('--resume_epoch', help='epoch from which to resume', default = None, type=int) |
| 167 | parser.add_argument('--eval', help='evaluation frequency', default=None, type = int) |
| 168 | parser.add_argument('--check', help='checkpoint frequency', default=None, type = int) |
| 169 | parser.add_argument('--n_max_batch', help='max batch per epoch (to speed up testing)', default=None, type = int) |
| 170 | parser.add_argument('--train_reverse_steps', help='number of diffusion steps used for training', default=None, type = int) |
| 171 | |
| 172 | parser.add_argument('--set_seed', help='set random seed', default = None, type=int) |
| 173 | parser.add_argument('--random_seed', help='set random seed to a random number', action = 'store_true', default=None) |
| 174 | |
| 175 | parser.add_argument('--log', help='activate logging to neptune', action='store_true', default=False) |
| 176 | parser.add_argument('--job_id', help='slurm job id', default=None, type = str) |
| 177 | |
| 178 | # EXPERIMENT parameters, specific to EVALUATION |
| 179 | parser.add_argument('--ema_eval', help='evaluate all ema models', action='store_true', default = False) |
| 180 | parser.add_argument('--no_ema_eval', help='dont evaluate ema models', action='store_true', default = False) |
| 181 | parser.add_argument('--generate', help='how many images/datapoints to generate', default = None, type = int) |
| 182 | parser.add_argument('--reverse_steps', help='choose number of reverse_steps', default = None, type = int) |
| 183 | parser.add_argument('--reset_eval', help='reset evaluation metrics', action='store_true', default = False) |
| 184 | |
| 185 | parser.add_argument('--deterministic', help='use deterministic sampling', default = False, action='store_true') |
| 186 | parser.add_argument('--clip', help='use clip denoised (diffusion)', default = False, action='store_true') |
| 187 | |
| 188 | # DATA |
| 189 | parser.add_argument('--dataset', help='choose specific dataset', default = None, type = str) |
| 190 | parser.add_argument('--nsamples', help='choose the size of the dataset (only 2d datasets)', default = None, type = str) |
| 191 | parser.add_argument('--data_std', help='choose data std or scale', default = None, type = float) |
| 192 | |
| 193 | # OPTIMIZER |
| 194 | parser.add_argument('--lr', help='reinitialize learning rate', type=float, default = None) |
| 195 | parser.add_argument('--lr_steps', help='reinitialize learning rate steps', type=int, default = None) |
| 196 | parser.add_argument('--lr_schedule', help='set learning rate schedule', type=str, default = None) |
| 197 | |
| 198 | # MODEL |
| 199 | # only useful for 2d datasets |
| 200 | parser.add_argument('--blocks', help='choose number of blocks in mlp', default = None, type = int) |
| 201 | parser.add_argument('--units', help='choose number of units in mlp', default = None, type = int) |
| 202 | parser.add_argument('--transforms', help='choose number of transforms in neural spline flow', default = None, type = int) |
| 203 | parser.add_argument('--depth', help='choose depth in neural spline flow', default = None, type = int) |
| 204 | parser.add_argument('--width', help='choose width in neural spline flow', default = None, type = int) |
| 205 | parser.add_argument('--t_embedding_type', help='choose time embedding type', default = None, type = str) |
| 206 | parser.add_argument('--t_embedding_size', help='choose time embedding size', default = None, type = int) |
| 207 | |
| 208 | # DIFFUSION |
| 209 | parser.add_argument('--alpha', help='alpha value for diffusion', default=None, type = float) |
| 210 | parser.add_argument('--non_iso', help='use non isotropic noise in the diffusion', action='store_true', default = False) |
| 211 | parser.add_argument('--non_iso_data', help='use non isotropic data', action='store_true', default = False) |
| 212 | parser.add_argument('--median', help='use median of mean. Specify (outer, inner).', nargs ='+', default = None) |