forked from georgesung/ssd_tensorflow_traffic_sign_detection
-
Notifications
You must be signed in to change notification settings - Fork 0
/
viz_model.py
24 lines (21 loc) · 721 Bytes
/
viz_model.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
'''
Visualize the model using TensorBoard
'''
import tensorflow as tf
from settings import *
from model import SSDModel
FM_ONLY = False # Only want to see feature map sizes?
with tf.Graph().as_default(), tf.Session() as sess:
if FM_ONLY:
# Only want to see feature map sizes (e.g. loss function and vector concatenation not yet set up)
if MODEL == 'AlexNet':
from model import AlexNet as MyModel
else:
raise NotImplementedError('Model %s not supported' % MODEL)
_ = MyModel()
else:
# This includes the entire graph, e.g. loss function, optimizer, etc.
_ = SSDModel()
tf.summary.merge_all()
writer = tf.summary.FileWriter('./tensorboard_out', sess.graph)
tf.global_variables_initializer().run()