import tensorflow as tf
import os
from tensorflow.python.tools import freeze_graph
from tensorflow.python.framework import graph_util
from tensorflow.python.saved_model import builder as saved_model_builder
from tensorflow.python.saved_model import signature_def_utils
from tensorflow.python.saved_model import signature_constants
from tensorflow.python.saved_model import tag_constants
from tensorflow.python.saved_model import utils as saved_model_utils
class CNN(object):
def __init__(self, input_size, num_classes, optimizer):
self.num_classes = num_classes
self.input_size = input_size
self.optimizer = optimizer
self.learning_rate = tf.placeholder(tf.float32,
shape=[],
name='learning_rate')
self.dropout_rate = tf.placeholder(tf.float32,
shape=[],
name='dropout_rate')
self.input = tf.placeholder(tf.float32, [None] + self.input_size,
name='input')
self.label = tf.placeholder(tf.float32, [None, self.num_classes],
name='label')
self.output = self.network_initializer()
self.loss = self.loss_initializer()
self.optimization = self.optimizer_initializer()
self.saver = tf.train.Saver()
self.sess = tf.Session()
self.sess.run(tf.global_variables_initializer())
def network(self, input, dropout_rate):
conv1 = tf.layers.conv2d(inputs=input,
filters=64,
kernel_size=[3, 3],
padding='same',
activation=tf.nn.relu,
name='conv1')
conv2 = tf.layers.conv2d(inputs=conv1,
filters=64,
kernel_size=[3, 3],
padding='same',
activation=tf.nn.relu,
name='conv2')
pool1 = tf.layers.max_pooling2d(inputs=conv2,
pool_size=[2, 2],
strides=[2, 2],
name='pool1')
pool1_dropout = tf.layers.dropout(inputs=pool1,
rate=dropout_rate,
training=True,
name='pool1_dropout')
conv3 = tf.layers.conv2d(inputs=pool1_dropout,
filters=128,
kernel_size=[3, 3],
padding='same',
activation=tf.nn.relu,
name='conv3')
conv4 = tf.layers.conv2d(inputs=conv3,
filters=128,
kernel_size=[3, 3],
padding='same',
activation=tf.nn.relu,
name='conv4')
pool2 = tf.layers.max_pooling2d(inputs=conv4,
pool_size=[2, 2],
strides=[2, 2],
name='pool2')
pool2_dropout = tf.layers.dropout(inputs=pool2,
rate=dropout_rate,
training=True,
name='pool2_dropout')
conv5 = tf.layers.conv2d(inputs=pool2_dropout,
filters=256,
kernel_size=[3, 3],
padding='same',
activation=tf.nn.relu,
name='conv5')
pool3 = tf.layers.max_pooling2d(inputs=conv5,
pool_size=[2, 2],
strides=[2, 2],
name='pool3')
pool3_dropout = tf.layers.dropout(inputs=pool3,
rate=dropout_rate,
training=True,
name='pool3_dropout')
flat = tf.layers.flatten(inputs=pool3_dropout, name='flat')
fc1 = tf.layers.dense(inputs=flat,
units=256,
activation=tf.nn.relu,
name='fc1')
fc1_dropout = tf.layers.dropout(inputs=fc1,
rate=dropout_rate,
training=True,
name='fc1_dropout')
fc2 = tf.layers.dense(inputs=fc1_dropout,
units=self.num_classes,
activation=None,
name='fc2')
# Give output node a
output = tf.identity(fc2, name='output')
return output
def network_initializer(self):
with tf.variable_scope('cnn') as scope:
ouput = self.network(input=self.input,
dropout_rate=self.dropout_rate)
return ouput
def loss_initializer(self):
with tf.variable_scope('loss') as scope:
cross_entropy = tf.nn.softmax_cross_entropy_with_logits_v2(
labels=self.label, logits=self.output, name='cross_entropy')
cross_entropy_mean = tf.reduce_mean(cross_entropy,
name='cross_entropy_mean')
return cross_entropy_mean
def optimizer_initializer(self):
if self.optimizer == 'Adam':
optimizer = tf.train.AdamOptimizer(
learning_rate=self.learning_rate).minimize(self.loss)
else:
optimizer = tf.train.GradientDescentOptimizer(
learning_rate=self.learning_rate).minimize(self.loss)
return optimizer
def train(self, data, label, learning_rate, dropout_rate):
_, train_loss = self.sess.run(
[self.optimization, self.loss],
feed_dict={
self.input: data,
self.label: label,
self.learning_rate: learning_rate,
self.dropout_rate: dropout_rate
})
return train_loss
def validate(self, data, label):
output, validate_loss = self.sess.run([self.output, self.loss],
feed_dict={
self.input: data,
self.label: label,
self.dropout_rate: 0.0
})
return output, validate_loss
def test(self, data):
output = self.sess.run(self.output,
feed_dict={
self.input: data,
self.dropout_rate: 0.0
})
return output
def save(self, directory, filename):
if not os.path.exists(directory):
os.makedirs(directory)
filepath = os.path.join(directory, filename + '.ckpt')
self.saver.save(self.sess, filepath)
return filepath
def save_signature(self, directory):
signature = signature_def_utils.build_signature_def(
inputs={
'input':
saved_model_utils.build_tensor_info(self.input),
'dropout_rate':
saved_model_utils.build_tensor_info(self.dropout_rate)
},
outputs={
'output': saved_model_utils.build_tensor_info(self.output)
},
method_name=signature_constants.PREDICT_METHOD_NAME)
signature_map = {
signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY: signature
}
model_builder = saved_model_builder.SavedModelBuilder(directory)
model_builder.add_meta_graph_and_variables(
self.sess,
tags=[tag_constants.SERVING],
signature_def_map=signature_map,
clear_devices=True)
model_builder.save(as_text=False)
def save_as_pb(self, directory, filename):
if not os.path.exists(directory):
os.makedirs(directory)
# Save check point for graph frozen later
ckpt_filepath = self.save(directory=directory, filename=filename)
pbtxt_filename = filename + '.pbtxt'
pbtxt_filepath = os.path.join(directory, pbtxt_filename)
pb_filepath = os.path.join(directory, filename + '.pb')
# This will only save the graph but the variables will not be saved.
# You have to freeze your model first.
tf.train.write_graph(graph_or_graph_def=self.sess.graph_def,
logdir=directory,
name=pbtxt_filename,
as_text=True)
# Freeze graph
# Method 1
freeze_graph.freeze_graph(input_graph=pbtxt_filepath,
input_saver='',
input_binary=False,
input_checkpoint=ckpt_filepath,
output_node_names='cnn/output',
restore_op_name='save/restore_all',
filename_tensor_name='save/Const:0',
output_graph=pb_filepath,
clear_devices=True,
initializer_nodes='')
# Method 2
'''
graph = tf.get_default_graph()
input_graph_def = graph.as_graph_def()
output_node_names = ['cnn/output']
output_graph_def = graph_util.convert_variables_to_constants(self.sess, input_graph_def, output_node_names)
with tf.gfile.GFile(pb_filepath, 'wb') as f:
f.write(output_graph_def.SerializeToString())
'''
return pb_filepath
def load(self, filepath):
if os.path.splitext(filepath)[1] != '.ckpt':
filepath += '.ckpt'
self.saver.restore(self.sess, filepath)
