File size: 1,233 Bytes
7efee70 |
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 |
import logging
import sys
import torch
import wandb
from .plot_color import Plot
from .metrics import Metric
class Logger:
def __init__(self, args, mds):
self.molecule = args.molecule
self.save_dir = args.save_dir
self.wandb = args.wandb
self.plot = Plot(args, mds)
self.metric = Metric(args, mds)
self.rmsd = float("inf")
def __call__(self, loss, rollout, policy):
metrics = self.metric()
if self.rmsd > metrics["rmsd"]:
self.rmsd = metrics["rmsd"]
torch.save(policy.state_dict(), f"{self.save_dir}/policy.pt")
if self.wandb:
if metrics["ets"] is not None:
wandb.log({
"rmsd": metrics["rmsd"],
"rmsd_std": metrics["rmsd_std"],
"thp": metrics["thp"],
"ets": metrics["ets"],
"ets_std": metrics["ets_std"],
"loss": loss
})
else:
wandb.log({
"rmsd": metrics["rmsd"],
"rmsd_std": metrics["rmsd_std"],
"thp": metrics["thp"],
"loss": loss
})
|