deployml/models/tensorflow_model.py at master · ikhlestov/deployml · GitHub

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
import numpy as np
import tensorflow as tf


class Model:
    input_node_name = 'input'
    output_node_name = 'output'
    input_data_type = tf.float32

    def __init__(self):
        self.input = tf.placeholder(
            self.input_data_type,
            shape=(None, 224, 224, 3),
            name=self.input_node_name
        )
        self.labels = tf.placeholder(
            tf.int32,
            shape=(None, 100),
            name='labels'
        )
        self.kernel_count = 0
        self._build_net()
        self._initialize_session()

    def _initialize_session(self):
        """Initialize session, variables, saver"""
        config = tf.ConfigProto()
        # restrict model GPU memory utilization to min required
        config.gpu_options.allow_growth = True
        self.sess = tf.Session(config=config)
        self.sess.run(tf.global_variables_initializer())

    def get_kernel_name(self):
        self.kernel_count += 1
        return 'kernel_{}'.format(self.kernel_count)

    def __kernel(self, in_features, out_features):
        return tf.get_variable(name=self.get_kernel_name(),
                               shape=(3, 3, in_features, out_features))

    def _build_net(self):
        # some constants
        strides = (1, 1, 1, 1)
        padding = "VALID"
        x = self.input

        # first block
        with tf.name_scope("first_block"):
            x = tf.nn.conv2d(x, filter=self.__kernel(3, 6), strides=strides, padding=padding)
            x = tf.nn.relu(tf.layers.batch_normalization(x))
            x = tf.nn.conv2d(x, filter=self.__kernel(6, 9), strides=strides, padding=padding)
            x = tf.nn.relu(tf.layers.batch_normalization(x))
            x = tf.nn.pool(x, (2, 2), 'AVG', padding=padding, strides=(2, 2))

        # second block
        with tf.name_scope("second_block"):
            x = tf.nn.conv2d(x, filter=self.__kernel(9, 12), strides=strides, padding=padding)
            x = tf.nn.relu(tf.layers.batch_normalization(x))
            x = tf.nn.conv2d(x, filter=self.__kernel(12, 16), strides=strides, padding=padding)
            x = tf.nn.relu(tf.layers.batch_normalization(x))
            x = tf.nn.pool(x, (2, 2), 'AVG', padding=padding, strides=(2, 2))

        # third block
        with tf.name_scope("third_block"):
            x = tf.nn.conv2d(x, filter=self.__kernel(16, 32), strides=strides, padding=padding)
            x = tf.nn.relu(tf.layers.batch_normalization(x))
            x = tf.nn.conv2d(x, filter=self.__kernel(32, 64), strides=strides, padding=padding)
            x = tf.nn.relu(tf.layers.batch_normalization(x))
            x = tf.nn.pool(x, (2, 2), 'AVG', padding=padding, strides=(2, 2))

        # forth block
        with tf.name_scope("forth_block"):
            x = tf.nn.conv2d(x, filter=self.__kernel(64, 128), strides=strides, padding=padding)
            x = tf.nn.relu(tf.layers.batch_normalization(x))
            x = tf.nn.conv2d(x, filter=self.__kernel(128, 256), strides=strides, padding=padding)
            # x = self.large_block(x, strides)
            x = tf.nn.relu(tf.layers.batch_normalization(x))

        # transition to classes
        with tf.name_scope("transition_to_classes"):
            x = tf.nn.pool(x, (20, 20), 'AVG', padding=padding, strides=(1, 1))
            x = tf.layers.flatten(x)
            x = tf.layers.dense(x, 512)
            x = tf.layers.batch_normalization(x)
            x = tf.nn.relu(x)
            x = tf.layers.dense(x, 100)
        x = tf.nn.sigmoid(x, name=self.output_node_name)
        self.output = x

        with tf.name_scope("training"):
            loss = tf.losses.softmax_cross_entropy(onehot_labels=self.labels, logits=x)
            train_step = tf.train.AdamOptimizer(0.01).minimize(loss)

    def predict(self, inputs):
        feed_dict = {self.input: inputs}
        pred = self.sess.run(self.output, feed_dict=feed_dict)
        return pred

    def large_block(self, x, strides):
        x = tf.nn.conv2d(x, filter=self.__kernel(256, 512), strides=strides, padding='SAME')
        x = tf.nn.conv2d(x, filter=self.__kernel(512, 512), strides=strides, padding='SAME')
        x = tf.nn.conv2d(x, filter=self.__kernel(512, 512), strides=strides, padding='SAME')
        x = tf.nn.conv2d(x, filter=self.__kernel(512, 256), strides=strides, padding='SAME')
        return x


if __name__ == '__main__':
    sample_image = np.random.random((1, 224, 224, 3))
    model = Model()
    preds = model.predict(sample_image)