PVD/utils/file_utils.py

85 lines
2.1 KiB
Python
Raw Normal View History

2021-10-19 20:54:46 +00:00
import os
import random
import sys
from shutil import copyfile
import datetime
import torch
import logging
logger = logging.getLogger()
import numpy as np
def set_global_gpu_env(opt):
os.environ['CUDA_DEVICE_ORDER'] = 'PCI_BUS_ID'
os.environ['CUDA_VISIBLE_DEVICES'] = str(opt.gpu)
torch.cuda.set_device(opt.gpu)
def copy_source(file, output_dir):
copyfile(file, os.path.join(output_dir, os.path.basename(file)))
def setup_logging(output_dir):
log_format = logging.Formatter("%(asctime)s : %(message)s")
logger = logging.getLogger()
logger.handlers = []
output_file = os.path.join(output_dir, 'output.log')
file_handler = logging.FileHandler(output_file)
file_handler.setFormatter(log_format)
logger.addHandler(file_handler)
console_handler = logging.StreamHandler(sys.stdout)
console_handler.setFormatter(log_format)
err_handler = logging.StreamHandler(sys.stderr)
err_handler.setFormatter(log_format)
logger.addHandler(err_handler)
logger.setLevel(logging.INFO)
return logger
def get_output_dir(prefix, exp_id):
t = datetime.datetime.now().strftime('%Y-%m-%d-%H-%M-%S')
output_dir = os.path.join(prefix, 'output/' + exp_id, t)
if not os.path.exists(output_dir):
os.makedirs(output_dir)
return output_dir
def set_seed(opt):
if opt.manualSeed is None:
opt.manualSeed = random.randint(1, 10000)
print("Random Seed: ", opt.manualSeed)
random.seed(opt.manualSeed)
torch.manual_seed(opt.manualSeed)
np.random.seed(opt.manualSeed)
if opt.gpu is not None and torch.cuda.is_available():
torch.cuda.manual_seed_all(opt.manualSeed)
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = True
def setup_output_subdirs(output_dir, *subfolders):
output_subdirs = output_dir
try:
os.makedirs(output_subdirs)
except OSError:
pass
subfolder_list = []
for sf in subfolders:
curr_subf = os.path.join(output_subdirs, sf)
try:
os.makedirs(curr_subf)
except OSError:
pass
subfolder_list.append(curr_subf)
return subfolder_list