-
Notifications
You must be signed in to change notification settings - Fork 1
/
logger.py
101 lines (86 loc) · 3.36 KB
/
logger.py
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
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
# Logger classes originally from
# https://github.com/SeanNaren/deepspeech.pytorch/blob/master/logger.py
import os
import torch
def to_np(x):
return x.cpu().numpy()
class VisdomLogger(object):
def __init__(self, id, num_epochs):
from visdom import Visdom
self.viz = Visdom()
self.opts = dict(
title=id,
ylabel="",
xlabel="Epoch",
legend=[
"Training Loss",
"Validation Loss",
"Lower Bound",
"Discriminative Loss",
],
)
self.viz_window = None
self.epochs = torch.arange(1, num_epochs + 1)
self.visdom_plotter = True
def update(self, epoch, values):
x_axis = self.epochs[0 : epoch + 1]
y_axis = torch.stack(
(
values["train_loss_results"][: epoch + 1],
values["val_loss_results"][: epoch + 1],
values["lower_bound_results"][: epoch + 1],
values["discrim_loss_results"][: epoch + 1],
),
dim=1,
)
self.viz_window = self.viz.line(
X=x_axis,
Y=y_axis,
opts=self.opts,
win=self.viz_window,
update="replace" if self.viz_window else None,
)
def load_previous_values(self, start_epoch, package):
# Add all values except the epoch we're starting from
self.update(start_epoch - 1, package)
class TensorBoardLogger(object):
def __init__(self, id, log_dir, log_params):
os.makedirs(log_dir, exist_ok=True)
from torch.utils.tensorboard import SummaryWriter
self.id = id
self.tensorboard_writer = SummaryWriter(log_dir)
self.log_params = log_params
def update(self, epoch, values, parameters=None):
tr_loss, val_loss, lower_bound, discrim_loss = (
values["train_loss_results"][epoch + 1],
values["val_loss_results"][epoch + 1],
values["lower_bound_results"][epoch + 1],
values["discrim_loss_results"][epoch + 1],
)
values = {
"Avg Train Loss": tr_loss,
"Avg Val Loss": val_loss,
"Avg Lower Bound": lower_bound,
"Avg Discriminative Loss": discrim_loss,
}
self.tensorboard_writer.add_scalars(self.id, values, epoch + 1)
if self.log_params:
for tag, value in parameters():
tag = tag.replace(".", "/")
self.tensorboard_writer.add_histogram(tag, to_np(value), epoch + 1)
self.tensorboard_writer.add_histogram(
tag + "/grad", to_np(value.grad), epoch + 1
)
def load_previous_values(self, start_epoch, values):
train_loss_results = values["train_loss_results"][:start_epoch]
val_loss_results = values["val_loss_results"][:start_epoch]
lb_results = values["lower_bound_results"][:start_epoch]
dl_results = values["discrim_loss_results"][:start_epoch]
for i in range(start_epoch):
values = {
"Avg Train Loss": train_loss_results[i],
"Avg Val Loss": val_loss_results[i],
"Avg Lower Bound": lb_results[i],
"Avg Discriminative Loss": dl_results[i],
}
self.tensorboard_writer.add_scalars(self.id, values, i + 1)