File size: 1,884 Bytes
d643072 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 |
import re
def tracker(args, result_dict, label="", pattern="epoch_step", metric="FID"):
if args.report_to == "wandb":
import wandb
wandb_name = f"[{args.log_metric}]_{args.name}"
wandb.init(project=args.tracker_project_name, name=wandb_name, resume="allow", id=wandb_name, tags="metrics")
run = wandb.run
if pattern == "step":
pattern = "sample_steps"
elif pattern == "epoch_step":
pattern = "step"
custom_name = f"custom_{pattern}"
run.define_metric(custom_name)
# define which metrics will be plotted against it
run.define_metric(f"{metric}_{label}", step_metric=custom_name)
steps = []
results = []
def extract_value(regex, exp_name):
match = re.search(regex, exp_name)
if match:
return match.group(1)
else:
return "unknown"
for exp_name, result_value in result_dict.items():
if pattern == "step":
regex = r".*step(\d+)_scale.*"
custom_x = extract_value(regex, exp_name)
elif pattern == "sample_steps":
regex = r".*step(\d+)_size.*"
custom_x = extract_value(regex, exp_name)
else:
regex = rf"{pattern}(\d+(\.\d+)?)"
custom_x = extract_value(regex, exp_name)
custom_x = 1 if custom_x == "unknown" else custom_x
assert custom_x != "unknown"
steps.append(float(custom_x))
results.append(result_value)
sorted_data = sorted(zip(steps, results))
steps, results = zip(*sorted_data)
for step, result in sorted(zip(steps, results)):
run.log({f"{metric}_{label}": result, custom_name: step})
else:
print(f"{args.report_to} is not supported")
|