| import os | |
| import yaml | |
| import random | |
| import numpy | |
| import torch | |
| def save_config(cfg, exp_dir): | |
| if not os.path.exists(exp_dir): | |
| os.makedirs(exp_dir, exist_ok=True) | |
| else: | |
| print( | |
| f"save model to {exp_dir}, old configs, checkpoints, and evaluation results will be overwritten." | |
| ) | |
| with open(os.path.join(exp_dir, "config.yaml"), "w") as f: | |
| yaml.dump(vars(cfg), f) | |
| def deterministic(seed): | |
| torch.manual_seed(seed) | |
| torch.cuda.manual_seed(seed) | |
| numpy.random.seed(seed) | |
| random.seed(seed) | |
| torch.backends.cudnn.deterministic = True | |
| torch.backends.cudnn.benchmark = False | |
| def set_cuda_visible_device(cfg, outmost=True): | |
| if cfg.device == "cpu": | |
| return | |
| if outmost: | |
| cuda_num = cfg.device.split(":")[-1] | |
| os.environ["CUDA_VISIBLE_DEVICES"] = cuda_num | |
| for key, value in cfg.__dict__.items(): | |
| if key == "device": | |
| cfg.device = "cuda:0" | |
| if hasattr(value, "__dict__"): | |
| set_cuda_visible_device(value, outmost=False) | |