def get_logger(model_dir):
path = get_log_path(model_dir)
utils.create_folders_if_necessary(path)
logging.basicConfig(level=logging.INFO,
format="%(message)s",
handlers=[
logging.FileHandler(filename=path),
logging.StreamHandler(sys.stdout)
])
return logging.getLogger()
def get_txt_logger(model_dir):
path = os.path.join(model_dir, "log.txt")
utils.create_folders_if_necessary(path)
logging.basicConfig(level=logging.INFO,
format="%(message)s",
handlers=[
logging.FileHandler(filename=path,
encoding='utf-8'),
logging.StreamHandler(sys.stdout)
])
return logging.getLogger()
def save_model(model1, model2, model_dir):
if torch.cuda.is_available():
model2.cpu()
if model1 is not None:
model1.cpu()
if model1 is not None:
path1 = get_model_path(model_dir, 1)
utils.create_folders_if_necessary(path1)
torch.save(model1, path1)
if model2 is not None:
path2 = get_model_path(model_dir, 2)
utils.create_folders_if_necessary(path2)
torch.save(model2, path2)
if torch.cuda.is_available():
model2.cuda()
if model1 is not None:
model1.cuda()
def save_config_in_table(args, name=None):
"""Save arguments passed to a train script in a CSV table and return a hash."""
csv_path = os.path.join(get_storage_dir(), f"{name or 'config'}.csv")
utils.create_folders_if_necessary(csv_path)
if not os.path.isfile(csv_path):
with open(csv_path, "w"):
pass
# Get current CSV header
with open(csv_path) as csv_file:
csv_reader = csv.reader(csv_file)
csv_header = None
for row in csv_reader:
csv_header = row
break
# Write config (and header)
args = OrderedDict(sorted(vars(args).items(), key=lambda t: t[0]))
with open(csv_path, "a") as csv_file:
csv_writer = csv.writer(csv_file)
args_header = ["hash"] + list(args.keys())
if csv_header is None:
csv_writer.writerow(args_header)
else:
assert csv_header == args_header, "Argument names have changed. Please change the name of the current config table."
config_hash = hashlib.md5(pickle.dumps(list(
args.values()))).hexdigest()[:10]
csv_writer.writerow([config_hash] + list(args.values()))
return config_hash
def get_csv_writer(model_dir):
csv_path = get_csv_path(model_dir)
utils.create_folders_if_necessary(csv_path)
csv_file = open(csv_path, "a")
return csv_file, csv.writer(csv_file)
def save_status(status, model_dir):
path = get_status_path(model_dir)
utils.create_folders_if_necessary(path)
with open(path, "w") as file:
json.dump(status, file)
def save_model(model, model_dir):
path = get_model_path(model_dir)
utils.create_folders_if_necessary(path)
torch.save(model, path)
def save(self):
utils.create_folders_if_necessary(self.path)
json.dump(self.vocab, open(self.path, "w"))
def save_difficulty(difficulty, model_dir):
path = get_difficulty_path(model_dir)
utils.create_folders_if_necessary(path)
with open(path, "w") as file:
json.dump({"difficulty": difficulty}, file)
def get_csv_logger(model_dir):
csv_path = os.path.join(model_dir, "log.csv")
utils.create_folders_if_necessary(csv_path)
csv_file = open(csv_path, "a")
return csv_file, csv.writer(csv_file)
def save_status(status, model_dir):
path = get_status_path(model_dir)
utils.create_folders_if_necessary(path)
torch.save(status, path)
type=int,
default=64,
help="batch size for PPO (default: 256)")
parser.add_argument("--debug",
action="store_true",
default=False,
help="Records Q values during training")
parser.add_argument("--dense-reward",
action="store_true",
help="Use dense reward during training.")
args = parser.parse_args()
# Get model directory
model_dir = "storage/" + args.model + "_seed_" + str(args.seed)
if args.dense_reward: model_dir += "_denser"
utils.create_folders_if_necessary(model_dir)
# Store json args
logger = utils.get_logger(model_dir)
with open(model_dir + '/args.json', 'w') as outfile:
json.dump(vars(args), outfile)
# Set seed for all randomness sources
utils.seed(args.seed)
# Load training status
try:
status = utils.load_status(model_dir)
except OSError:
status = {"num_frames": 0, "difficulty": args.difficulty}
utils.save_status(status, model_dir)
def save_config(model_dir, config):
path = os.path.join(model_dir, "config.pickle")
utils.create_folders_if_necessary(path)
pickle.dump(config, open(path, "wb"))
def get_csv_writer(save_dir):
csv_path = get_csv_path(save_dir)
utils.create_folders_if_necessary(csv_path)
file = open(csv_path, "a")
return csv.writer(file)
