| import torch |
| import numpy as np |
| import random |
|
|
| import os |
| import logging |
| from omegaconf import OmegaConf |
|
|
| import sys |
| import fsspec |
| import pdb |
|
|
|
|
| def makedirs(dirname): |
| os.makedirs(dirname, exist_ok=True) |
|
|
|
|
| def get_logger(logpath, package_files=[], displaying=True, saving=True, debug=False): |
| logger = logging.getLogger() |
| if debug: |
| level = logging.DEBUG |
| else: |
| level = logging.INFO |
|
|
| if (logger.hasHandlers()): |
| logger.handlers.clear() |
|
|
| logger.setLevel(level) |
| formatter = logging.Formatter('%(asctime)s - %(message)s') |
| if saving: |
| info_file_handler = logging.FileHandler(logpath, mode="a") |
| info_file_handler.setLevel(level) |
| info_file_handler.setFormatter(formatter) |
| logger.addHandler(info_file_handler) |
| if displaying: |
| console_handler = logging.StreamHandler() |
| console_handler.setLevel(level) |
| console_handler.setFormatter(formatter) |
| logger.addHandler(console_handler) |
|
|
| for f in package_files: |
| logger.info(f) |
| with open(f, "r") as package_f: |
| logger.info(package_f.read()) |
|
|
| return logger |
|
|
|
|
| def restore_checkpoint(ckpt_dir, state, device): |
| if not os.path.exists(ckpt_dir): |
| makedirs(os.path.dirname(ckpt_dir)) |
| logging.warning(f"No checkpoint found at {ckpt_dir}. Returned the same state as input") |
| return state |
| else: |
| loaded_state = torch.load(ckpt_dir, map_location=device) |
| try: |
| state['optimizer'].load_state_dict(loaded_state['optimizer']) |
| except: |
| pass |
| state['model'].module.load_state_dict(loaded_state['model'], strict=False) |
| state['ema'].load_state_dict(loaded_state['ema']) |
| state['step'] = loaded_state['step'] |
| state['config'] = loaded_state['config'] |
| return state |
|
|
|
|
| def save_checkpoint(ckpt_dir, state): |
| saved_state = { |
| 'optimizer': state['optimizer'].state_dict(), |
| 'model': state['model'].module.state_dict(), |
| 'ema': state['ema'].state_dict(), |
| 'step': state['step'], |
| 'config': state['config'] |
| } |
| torch.save(saved_state, ckpt_dir) |
|
|
|
|
| def set_seed(seed): |
| |
| random.seed(seed) |
| np.random.seed(seed) |
| torch.manual_seed(seed) |
| if torch.cuda.is_available(): |
| torch.cuda.manual_seed(seed) |
| torch.cuda.manual_seed_all(seed) |
| |
|
|
| def fsspec_exists(filename): |
| """Check if a file exists using fsspec.""" |
| fs, _ = fsspec.core.url_to_fs(filename) |
| return fs.exists(filename) |
|
|
|
|
| def fsspec_listdir(dirname): |
| """Listdir in manner compatible with fsspec.""" |
| fs, _ = fsspec.core.url_to_fs(dirname) |
| return fs.ls(dirname) |
|
|
|
|
| def fsspec_mkdirs(dirname, exist_ok=True): |
| """Mkdirs in manner compatible with fsspec.""" |
| fs, _ = fsspec.core.url_to_fs(dirname) |
| fs.makedirs(dirname, exist_ok=exist_ok) |
