-
Notifications
You must be signed in to change notification settings - Fork 0
/
logger.py
131 lines (110 loc) · 4.2 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
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
"""
File: logger.py
Modified by: Senthil Purushwalkam
Code referenced from https://gist.github.com/gyglim/1f8dfb1b5c82627ae3efcfbbadb9f514
Email: spurushw<at>andrew<dot>cmu<dot>edu
Github: https://github.com/senthilps8
Description:
"""
import tensorflow as tf
from torch.autograd import Variable
import numpy as np
import scipy.misc
import os
try:
from StringIO import StringIO # Python 2.7
except ImportError:
from io import BytesIO # Python 3.x
class Logger(object):
def __init__(self, log_dir, name=None):
"""Create a summary writer logging to log_dir."""
if name is None:
name = 'temp'
self.name = name
if name is not None:
try:
os.makedirs(os.path.join(log_dir, name))
except:
pass
self.writer = tf.summary.FileWriter(os.path.join(log_dir, name),
filename_suffix=name)
else:
self.writer = tf.summary.FileWriter(log_dir, filename_suffix=name)
def scalar_summary(self, tag, value, step):
"""Log a scalar variable."""
summary = tf.Summary(value=[tf.Summary.Value(tag=tag, simple_value=value)])
self.writer.add_summary(summary, step)
def image_summary(self, tag, images, step):
"""Log a list of images."""
img_summaries = []
for i, img in enumerate(images):
# Write the image to a string
try:
s = StringIO()
except:
s = BytesIO()
scipy.misc.toimage(img).save(s, format="png")
# Create an Image object
img_sum = tf.Summary.Image(encoded_image_string=s.getvalue(),
height=img.shape[0],
width=img.shape[1])
# Create a Summary value
img_summaries.append(tf.Summary.Value(tag='%s/%d' % (tag, i), image=img_sum))
# Create and write Summary
summary = tf.Summary(value=img_summaries)
self.writer.add_summary(summary, step)
def histo_summary(self, tag, values, step, bins=1000):
"""Log a histogram of the tensor of values."""
# Create a histogram using numpy
counts, bin_edges = np.histogram(values, bins=bins)
# Fill the fields of the histogram proto
hist = tf.HistogramProto()
hist.min = float(np.min(values))
hist.max = float(np.max(values))
hist.num = int(np.prod(values.shape))
hist.sum = float(np.sum(values))
hist.sum_squares = float(np.sum(values**2))
# Drop the start of the first bin
bin_edges = bin_edges[1:]
# Add bin edges and counts
for edge in bin_edges:
hist.bucket_limit.append(edge)
for c in counts:
hist.bucket.append(c)
# Create and write Summary
summary = tf.Summary(value=[tf.Summary.Value(tag=tag, histo=hist)])
self.writer.add_summary(summary, step)
self.writer.flush()
def to_np(self, x):
return x.data.cpu().numpy()
def to_var(self, x):
if torch.cuda.is_available():
x = x.cuda()
return Variable(x)
def model_param_histo_summary(self, model, step):
"""log histogram summary of model's parameters
and parameter gradients
"""
for tag, value in model.named_parameters():
if value.grad is None:
continue
tag = tag.replace('.', '/')
tag = self.name+'/'+tag
self.histo_summary(tag, self.to_np(value), step)
self.histo_summary(tag+'/grad', self.to_np(value.grad), step)
def bounding_box(self, img, boxes, step):
bb_img = tf.image.draw_bounding_boxes(
img,
boxes,
name=None
)
try:
s = StringIO()
except:
s = BytesIO()
scipy.misc.toimage(bb_img).save(s, format="png")
# Create an Image object
img_sum = tf.Summary.Image(encoded_image_string=s.getvalue(),
height=bb_img.shape[0],
width=bb_img.shape[1])
self.writer.add_summary(img_sum, step)