File size: 883 Bytes
a8639ac
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
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 &")