Spaces:
Sleeping
Sleeping
import os, shutil | |
import torch | |
# from tensorboardX import SummaryWriter | |
from .options import * | |
import torch.distributed as dist | |
import time | |
""" ==================== Save ======================== """ | |
def make_path(): | |
return "{}_{}_bs{}_lr{}".format(opts.expri,opts.savepath,opts.batch_size,opts.learn_rate) | |
def save_model(model,name): | |
save_path = make_path() | |
if not os.path.isdir(os.path.join(config['checkpoint_base'], save_path)): | |
os.makedirs(os.path.join(config['checkpoint_base'], save_path), exist_ok=True) | |
model_name = os.path.join(config['checkpoint_base'], save_path, name) | |
torch.save(model.state_dict(), model_name) | |
""" ==================== Tools ======================== """ | |
def is_dist_avail_and_initialized(): | |
if not dist.is_available(): | |
return False | |
if not dist.is_initialized(): | |
return False | |
return True | |
def get_rank(): | |
if not is_dist_avail_and_initialized(): | |
return 0 | |
return dist.get_rank() | |
def makedir(path): | |
if not os.path.exists(path): | |
os.makedirs(path, 0o777) | |
# def visualizer(): | |
# if get_rank() == 0: | |
# # filewriter_path = config['visual_base']+opts.savepath+'/' | |
# save_path = make_path() | |
# filewriter_path = os.path.join(config['visual_base'], save_path) | |
# if opts.clear_visualizer and os.path.exists(filewriter_path): # 删掉以前的summary,以免重合 | |
# shutil.rmtree(filewriter_path) | |
# makedir(filewriter_path) | |
# writer = SummaryWriter(filewriter_path, comment='visualizer') | |
# return writer | |