code-completion / logger.py
jblitzar's picture
Upload folder using huggingface_hub
a8639ac verified
raw
history blame contribute delete
883 Bytes
import torch
from torch.utils.tensorboard import SummaryWriter
import os
import webbrowser
writer = None
def flush():
global writer
writer.flush()
writer = None
def log_data(data, i):
for key in data.keys():
writer.add_scalar(key, data[key], i)
def log_img(img, name):
writer.add_image(name, img)
def init_logger(net, data=None, dir="runs"):
net.eval()
global writer
if not writer or writer is None:
writer = SummaryWriter(dir)
if data is not None:
existing_files = [
f for f in os.listdir(dir) if f.startswith("events.out.tfevents.")
]
if not existing_files:
writer.add_graph(net, data)
# writer.close()
net.train()
os.system("tensorboard --logdir runs > /dev/null 2>&1 &")
# os.system("sleep 5; open -a /Applications/Safari.app http://localhost:6006 &")