cnn.py

import tensorflow as tf
import os
# print(tf.__file__)
os.environ["CUDA_VISIBLE_DEVICES"]="1"

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.is_train = tf.placeholder(tf.bool, name="is_train")
        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,self.input_size,4], 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):

        x_norm = tf.layers.batch_normalization(input, training=self.is_train)
        conv1 = tf.layers.conv2d(
            inputs = x_norm,
            filters = 4,
            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 = conv1,
            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,
            filters = 4,
            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 = conv3,
            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,
            filters = 4,
            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=self.is_train,
            name = 'pool3_dropout')

        flat = tf.layers.flatten(
            inputs = pool3_dropout, 
            name = 'flat')

        fc1 = tf.layers.dense(
            inputs = flat,
            units = 56,
            activation = tf.nn.relu,
            name = 'fc1')

        fc1_dropout = tf.layers.dropout(
            inputs = fc1,
            rate = dropout_rate,
            training=self.is_train,
            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:
            l=tf.losses.mean_squared_error( self.label , self.output )
            # cross_entropy = tf.nn.softmax_cross_entropy_with_logits_v2(
                # labels = self.label, logits = self.output, name = 'cross_entropy')
            l_mean=tf.reduce_mean(l, name = 'mse_mean')
            # cross_entropy_mean = tf.reduce_mean(cross_entropy, name = 'cross_entropy_mean')
        return l_mean

    def optimizer_initializer(self):
        update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
        with tf.control_dependencies(update_ops):
            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,self.is_train: True})
        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,self.is_train: False})
        return output, validate_loss

    def test(self, data):

        output = self.sess.run(self.output, feed_dict = {self.input: data, self.dropout_rate: 0.0,self.is_train: False})

        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 gradients(self,data,label):
        grads = tf.gradients(self.loss,tf.trainable_variables())
        grads_and_vars = list(zip(grads,tf.trainable_variables()))
        res = self.sess.run(grads_and_vars,feed_dict = {self.input: data, self.label: label,self.dropout_rate: 0.2,self.is_train: False})
        return res
    
    def load(self, filepath):

        if os.path.splitext(filepath)[1] != '.ckpt':
            filepath += '.ckpt'

        self.saver.restore(self.sess, filepath)