56 lines
1.5 KiB
Python
56 lines
1.5 KiB
Python
"""
|
|
usage:
|
|
python tabulate_rl_loss.py --log_dir /home/zhengzangwei/projs/Open-Sora-dev/logs/loss --ckpt_name epoch0-global_step9000
|
|
|
|
save the processed json to:
|
|
Open-Sora-dev/evaluation_results/rectified_flow/<ckpt_name>_loss.json
|
|
"""
|
|
|
|
import argparse
|
|
import json
|
|
import os
|
|
from ast import literal_eval
|
|
|
|
|
|
def parse_args():
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--log_dir", type=str)
|
|
args = parser.parse_args()
|
|
return args
|
|
|
|
|
|
if __name__ == "__main__":
|
|
args = parse_args()
|
|
|
|
files = os.listdir(args.log_dir)
|
|
files = [
|
|
"img_0.log",
|
|
"img_1.log",
|
|
"img_2.log",
|
|
"144p_vid.log",
|
|
"240p_vid.log",
|
|
"360p_vid.log",
|
|
"480p_vid.log",
|
|
"720p_vid.log",
|
|
]
|
|
|
|
loss_info = {}
|
|
|
|
for fname in files:
|
|
path = os.path.join(args.log_dir, fname)
|
|
with open(path, "r", encoding="utf-8") as f:
|
|
content = f.readlines()
|
|
eval_line = content[-1].split("losses:")[-1].strip()
|
|
loss_dict = literal_eval(eval_line)
|
|
for key, loss in loss_dict.items():
|
|
resolution, frame = key
|
|
if resolution not in loss_info:
|
|
loss_info[resolution] = {}
|
|
loss_info[resolution][frame] = format(loss, ".4f")
|
|
|
|
# Convert and write JSON object to file
|
|
output_file_path = os.path.join(args.log_dir, "loss.json")
|
|
with open(output_file_path, "w") as outfile:
|
|
json.dump(loss_info, outfile, indent=4, sort_keys=True)
|
|
print(f"results saved to: {output_file_path}")
|