in src/run.py [0:0]
def parse_arguments():
parser = argparse.ArgumentParser()
# data and I/O
parser.add_argument("--data_path", type=str, default="/root/downloads/imagenet")
parser.add_argument("--ckpt_path", type=str, default="/root/downloads/model.ckpt-1000000")
parser.add_argument("--color_cluster_path", type=str, default="/root/downloads/kmeans_centers.npy")
parser.add_argument("--save_dir", type=str, default="/root/save/")
# model
parser.add_argument("--n_embd", type=int, default=512)
parser.add_argument("--n_head", type=int, default=8)
parser.add_argument("--n_layer", type=int, default=24)
parser.add_argument("--n_px", type=int, default=32, help="image height or width in pixels")
parser.add_argument("--n_vocab", type=int, default=512, help="possible values for each pixel")
parser.add_argument("--bert", action="store_true", help="use the bert objective (defaut: autoregressive)")
parser.add_argument("--bert_mask_prob", type=float, default=0.15)
parser.add_argument("--clf", action="store_true", help="add a learnable classification head")
# parallelism
parser.add_argument("--n_sub_batch", type=int, default=8, help="per-gpu batch size")
parser.add_argument("--n_gpu", type=int, default=8, help="number of gpus to distribute training across")
# mode
parser.add_argument("--eval", action="store_true", help="evaluates the model, requires a checkpoint and dataset")
parser.add_argument("--sample", action="store_true", help="samples from the model, requires a checkpoint and clusters")
# reproducibility
parser.add_argument("--seed", type=int, default=42, help="seed for random, np, tf")
args = parser.parse_args()
print("input args:\n", json.dumps(vars(args), indent=4, separators=(",", ":")))
return args