in hyperparams.py [0:0]
def add_arguments(parser):
parser.add_argument('--out_dir', type=str, default=DEFAULT_OUT_DIR)
parser.add_argument('--desc', type=str, default='test')
parser.add_argument('--print_params', action="store_true")
parser.add_argument('--hparam_sets', '--hps', type=str, default='')
# dataset params
parser.add_argument('--dataset', type=str, default="cifar10")
parser.add_argument('--auxiliary_dataset', type=str, default=None)
parser.add_argument('--auxiliary_dataset_fraction', type=float, default=0.5)
parser.add_argument('--auxiliary_dataset_subset_size', type=int, default=None)
parser.add_argument('--auxiliary_dataset_seed', type=int, default=42)
# Training params
parser.add_argument('--n_batch', type=int, default=128)
parser.add_argument('--max_grad_norm', type=float, default=1.0)
# Transformer architectural parameters
parser.add_argument('--n_embd', type=int, default=512)
parser.add_argument('--n_ctx', type=int, default=256)
parser.add_argument('--n_head', type=int, default=8)
parser.add_argument('--n_layer', type=int, default=6)
parser.add_argument('--dropout_broadcast_dims', type=str, default=None)
parser.add_argument('--embd_pdrop', type=float, default=0.1)
parser.add_argument('--resid_pdrop', type=float, default=0.1)
parser.add_argument('--mlp_multiple', type=float, default=4.0)
parser.add_argument('--qk_ratio', type=float, default=1.0)
parser.add_argument('--attention_layers', type=str, default='a')
parser.add_argument('--local_attn_ctx', type=int, default=64)
parser.add_argument('--pos_embd_std', type=float, default=0.007)
parser.add_argument('--w_embd_std', type=float, default=0.013)
parser.add_argument('--mlp_w1', type=float, default=0.125)
parser.add_argument('--mlp_w2', type=float, default=0.125)
parser.add_argument('--qk_w', type=float, default=0.125)
parser.add_argument('--v_w', type=float, default=0.125)
parser.add_argument('--post_w', type=float, default=0.125)
parser.add_argument('--logits_w', type=float, default=0.125)
parser.add_argument('--preconv_w', type=float, default=0.125)
# rand augment params
# https://arxiv.org/pdf/1909.13719.pdf
parser.add_argument('--rand_augment', action="store_true")
parser.add_argument('--rand_augment_conditioning', action="store_true")
parser.add_argument('--rand_augment_rate', type=float, default=0.95)
parser.add_argument('--rand_augment_n', type=int, default=1) # Number of sequential perturbations -- range [1, 3]
parser.add_argument('--rand_augment_m', type=int, default=2) # Magnitude of pertubations -- range [2, 30]
# Distr Aug Params
parser.add_argument('--aug', action='store_true')
parser.add_argument('--permute_embeddings', dest='permute_embeddings', action="store_true")
parser.add_argument('--no_permute_embeddings', dest='permute_embeddings', action="store_false")
parser.set_defaults(permute_embeddings=True)
parser.add_argument('--use_imagenet_fraction', type=float, default=1.0)
parser.add_argument('--unaugmented_data_rate', type=float, default=None)
parser.add_argument('--use_rotation', action="store_true")
parser.add_argument('--use_dataset_conditioning', action="store_true")
parser.add_argument('--no_dataset_conditioning', action="store_false", dest="use_dataset_conditioning")
parser.add_argument('--use_color', action="store_true")
parser.add_argument('--use_transposition', action="store_true")
parser.add_argument('--use_randomly_determined_order', action="store_true")
parser.add_argument('--randomly_determined_order_num_perms', type=int, default=3)
parser.add_argument('--randomly_determined_order_seed', type=int, default=42)
parser.add_argument('--randomly_determined_order_use_lookahead', action="store_true")
parser.add_argument('--use_reverse', action="store_true")
parser.add_argument('--use_linf_pgd', action="store_true")
parser.add_argument('--use_jigsaw', action="store_true")
parser.add_argument('--jigsaw_grid_size', type=int, default=2)
parser.add_argument('--use_unconditional_augmentation', action='store_true')
parser.add_argument('--datapoints', type=int, default=None)
parser.add_argument('--test_size', type=int, default=None)
# Training params
parser.add_argument('--seed', type=int, default=42)
parser.add_argument('--aug_seed', type=int, default=314)
parser.add_argument('--optimizer', type=str, default='bs_adam')
parser.add_argument('--activation', type=str, default='quick_gelu')
parser.add_argument('--beta2', type=float, default=0.999)
parser.add_argument('--l2_loss', type=float, default=0.0)
parser.add_argument('--recompute', action="store_true", dest="recompute")
parser.add_argument('--no_recompute', action="store_false", dest="recompute")
parser.add_argument('--float16', action="store_true")
parser.add_argument('--no_float16', action="store_false", dest='float16')
parser.add_argument('--blocksparse_op', action="store_true")
parser.add_argument('--no_blocksparse_op', action="store_false", dest="blocksparse_op")
parser.add_argument('--blocksize', type=int, default=64)
parser.add_argument('--fp16_allreduce', action="store_true")
parser.add_argument('--no_fp16_allreduce', action="store_false", dest='fp16_allreduce')
parser.add_argument('--merge_layer_allreduce', default=0, type=int)
parser.add_argument('--fp32_gains_biases', action="store_true")
parser.add_argument('--fp16_loss_scale', type=float, default=2.0**16)
parser.add_argument('--min_loss_scale', type=float, default=2.0**10)
parser.add_argument('--fp16_loss_freq', type=int, default=1000)
parser.add_argument('--fp16_mean_var', action='store_true')
parser.add_argument('--no_fp16_mean_var', action='store_false',
dest='fp16_mean_var')
parser.add_argument('--dynamic_loss_scaling', action='store_true')
parser.add_argument('--no_dynamic_loss_scaling', action='store_false',
dest='dynamic_loss_scaling')
parser.add_argument('--lr', type=float, default=0.0005)
parser.add_argument('--lr_offset', type=int, default=0)
parser.add_argument('--decay_lr_linearly', action="store_true")
parser.add_argument('--no_vocab_rounding', action="store_true")
parser.add_argument('--disable_ema_vars', action="store_true")
parser.add_argument('--total_epochs', type=int, default=100)
parser.add_argument('--exit_after_n_epochs', type=int, default=None)
parser.add_argument('--warmup_iters', type=int, default=5000)
parser.add_argument('--weights_beta', type=float, default=0.999)
parser.add_argument('--iters_per_log', type=int, default=500)
parser.add_argument('--aug_eval', type=str, default=None)
parser.add_argument('--aug_eval_n_examples', type=int, default=None)
parser.add_argument('--eval_after_n_examples', type=int, default=None)
parser.add_argument('--epochs_per_save', type=int, default=1)
parser.add_argument('--epochs_per_backup', type=int, default=1)
parser.add_argument('--epochs_per_eval', type=int, default=1)
# eval stuff
parser.add_argument('--skip_initial_evals', action="store_true")
parser.add_argument('--eval_and_exit', action="store_true")
parser.add_argument('--no_skip_initial_evals', action="store_false",
dest='skip_initial_evals')
parser.add_argument('--eval_test', action="store_true")
parser.add_argument('--eval_start_idx', type=int, default=0)
parser.add_argument('--eval_n_examples', type=int, default=100000)
# Generating unconditional samples
parser.add_argument('--sample_batch', type=int, default=4)
parser.add_argument('--samples_to_generate', type=int, default=4)
parser.add_argument('--sample_grid_dim', type=int, default=4)
parser.add_argument('--sample_and_exit', action="store_true")
parser.add_argument('--sample_during_eval', action="store_true")
parser.add_argument('--sample_f16', action="store_true")
parser.add_argument('--temperature', type=float, default=1.0)
parser.add_argument('--no_sample_during_eval', action="store_false", dest='sample_during_eval')
# Restoring jobs
parser.add_argument('--restore_path', type=str, default='')
return parser