| |
| import argparse |
| import os |
| import yaml |
| import json |
| import random |
| import numpy as np |
| import torch |
|
|
| def mkdir(dirpath): |
| os.makedirs(dirpath, exist_ok=True) |
|
|
| def get_config(): |
| parser = argparse.ArgumentParser() |
|
|
| parser.add_argument('--config', type=str, default='config.json', |
| help="Path to base config file (json or yaml)") |
| parser.add_argument('--tokenizer_dir', type=str) |
| parser.add_argument('--data_dir', type=str) |
| parser.add_argument('--train_glob', type=str) |
| parser.add_argument('--valid_glob', type=str) |
| parser.add_argument('--output_dir', type=str) |
|
|
| |
| parser.add_argument('--datapoint_length', type=int) |
| parser.add_argument('--training_type', type=str, choices=["strict", "strict_small"]) |
| parser.add_argument('--n_epochs', type=int) |
| parser.add_argument('--batch_size', type=int) |
| parser.add_argument('--learning_rate', type=float) |
| parser.add_argument('--weight_decay', type=float) |
| parser.add_argument('--num_training_steps', type=int) |
| parser.add_argument('--num_warmup_steps', type=int) |
| parser.add_argument('--gradient_clip_norm', type=float) |
|
|
| |
| parser.add_argument('--seed', type=int) |
| parser.add_argument('--base_folder', type=str) |
| parser.add_argument('--experiment_name', type=str) |
| parser.add_argument('--use_wandb', action='store_true') |
| parser.add_argument('--wandb_project_name', type=str) |
| parser.add_argument('--wandb_experiment_name', type=str) |
|
|
|
|
|
|
| args = parser.parse_args() |
| config = construct_config(args) |
| return config |
|
|
| def setup_experiment(cfg): |
| |
| if cfg.get("seed", -1) == -1: |
| cfg["seed"] = random.randint(0, 10**9) |
| random.seed(cfg["seed"]) |
| np.random.seed(cfg["seed"]) |
| torch.manual_seed(cfg["seed"]) |
| torch.cuda.manual_seed_all(cfg["seed"]) |
| torch.backends.cudnn.deterministic = True |
| torch.backends.cudnn.benchmark = False |
| print(f"[utils] Using seed {cfg['seed']}") |
|
|
| |
| cfg["expdir"] = os.path.join(cfg["base_folder"], cfg["experiment_name"]) |
| cfg["checkpoint_dir"] = os.path.join(cfg["expdir"], 'checkpoints') |
| cfg["logdir"] = os.path.join(cfg["expdir"], 'logging') |
| mkdir(cfg["expdir"]); mkdir(cfg["checkpoint_dir"]); mkdir(cfg["logdir"]) |
|
|
| |
| with open(os.path.join(cfg["logdir"], "exp_cfg.yaml"), 'w') as cfg_file: |
| yaml.safe_dump(cfg, cfg_file, sort_keys=False) |
|
|
| def setup_wandb(cfg): |
| try: |
| import wandb |
| except ImportError: |
| raise RuntimeError("use_wandb is true but wandb is not installed") |
| wandb.init( |
| project=cfg["wandb_project_name"], |
| name=cfg["wandb_experiment_name"] |
| ) |
|
|
| def load_file_any(filepath): |
| ext = os.path.splitext(filepath)[1].lower() |
| with open(filepath, 'r') as f: |
| if ext in ['.yaml', '.yml']: |
| return yaml.safe_load(f) |
| else: |
| return json.load(f) |
|
|
| def construct_config(args): |
| base_cfg = load_file_any(args.config) |
| |
| for k, v in vars(args).items(): |
| if k == "config": |
| continue |
| if v is not None: |
| base_cfg[k] = v |
| return base_cfg |
