44 lines
1.2 KiB
Python
44 lines
1.2 KiB
Python
import threading
|
|
|
|
import numpy as np
|
|
|
|
from utils.util import to_np
|
|
|
|
|
|
def histogram_logger(tb_writer, iter_index, net=None):
|
|
if net is not None:
|
|
for tag, value in net.named_parameters():
|
|
tag = tag.replace('.', '/')
|
|
tb_writer.add_histogram(tag, to_np(value), iter_index)
|
|
tb_writer.add_histogram(tag + '/grad', to_np(value.grad), iter_index)
|
|
|
|
|
|
def value_logger(tb_writer, iter_index, loss_info):
|
|
for tag, value in loss_info.items():
|
|
tb_writer.add_scalar(tag, value, iter_index)
|
|
|
|
|
|
class TBImageRecorder(threading.Thread):
|
|
"""
|
|
TBImageRecorder
|
|
"""
|
|
|
|
def __init__(self, tb_writer, func, queue):
|
|
super(TBImageRecorder, self).__init__()
|
|
self._tb_writer = tb_writer
|
|
self._func = func
|
|
self._queue = queue
|
|
|
|
def run(self):
|
|
while True:
|
|
msgs, iter_index = self._queue.get()
|
|
if msgs:
|
|
img_info = self._func(*msgs)
|
|
|
|
print("logging the images")
|
|
for tag, images in img_info.items():
|
|
if images is not None:
|
|
self._tb_writer.add_image(tag, np.clip(images, 0, 255).astype(np.uint8), iter_index)
|
|
else:
|
|
break
|