From c2a402f7f7a61e21cb1699a11ba6a3b081b415e8 Mon Sep 17 00:00:00 2001 From: Andrej Karpathy Date: Wed, 11 Jan 2023 01:00:22 +0000 Subject: [PATCH] guess the config from globals() and log all of it with wandb --- train.py | 14 ++++++-------- 1 file changed, 6 insertions(+), 8 deletions(-) diff --git a/train.py b/train.py index 581635d..37bd3fb 100644 --- a/train.py +++ b/train.py @@ -48,7 +48,8 @@ dropout = 0.0 # for pretraining 0 is good, for finetuning try 0.1+ learning_rate = 6e-4 # max learning rate max_iters = 600000 # total number of training iterations weight_decay = 1e-2 -betas = (0.9, 0.95) +beta1 = 0.9 +beta2 = 0.95 # learning rate decay settings decay_lr = True # whether to decay the learning rate warmup_iters = 2000 # how many steps to warm up for @@ -61,7 +62,9 @@ device = 'cuda' # examples: 'cpu', 'cuda', 'cuda:0', 'cuda:1', etc. dtype = 'bfloat16' # 'float32' or 'bfloat16' compile = True # use PyTorch 2.0 to compile the model to be faster # ----------------------------------------------------------------------------- +config_keys = [k for k,v in globals().items() if not k.startswith('_') and isinstance(v, (int, float, bool, str))] exec(open('configurator.py').read()) # overrides from command line or config file +config = {k: globals()[k] for k in config_keys} # will be useful for logging # ----------------------------------------------------------------------------- # various inits, derived attributes, I/O setup @@ -142,7 +145,7 @@ if block_size < model.config.block_size: model.to(device) # optimizer -optimizer = model.configure_optimizers(weight_decay, learning_rate, betas) +optimizer = model.configure_optimizers(weight_decay, learning_rate, (beta1, beta2)) if init_from == 'resume': optimizer.load_state_dict(checkpoint['optimizer']) @@ -188,12 +191,7 @@ def get_lr(iter): # logging if wandb_log and gpu_id == 0: import wandb - wandb.init(project=wandb_project, name=wandb_run_name) - wandb.config = { - "batch_size": batch_size, - "block_size": block_size, - "learning_rate": learning_rate, # TODO log everything else too - } + wandb.init(project=wandb_project, name=wandb_run_name, config=config) # training loop t0 = time.time()