fcn.tf/vgg16_model.py at master · MtZero/fcn.tf · 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
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
"""vgg16 model.

Related papers:
https://arxiv.org/pdf/1409.1556.pdf
"""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import tensorflow as tf
import numpy as np

import TensorflowUtils as utils

class vgg16(object):
    """vgg16 model."""

    # def __init__(self, is_training, data_format, batch_norm_decay, batch_norm_epsilon):
    def __init__(self):
        """vgg16 constructor.

        Args:
          is_training: if build training or inference model.
          data_format: the data_format used during computation.
                       one of 'channels_first' or 'channels_last'.
        """
        # self._batch_norm_decay = batch_norm_decay
        # self._batch_norm_epsilon = batch_norm_epsilon
        # self._is_training = is_training
        # assert data_format in ('channels_first', 'channels_last')
        # self._data_format = data_format
        self._kernel_size = 3
        self._stride = 1
        self.pool_size = 2
        self.pool_stride = 2

    def forward_pass(self, x):
        raise NotImplementedError(
            'forward_pass() is implemented in ResNet sub classes')

    def _vgg16_modified(self, x, weights):
        with tf.name_scope('vgg16') as name_scope:
            orig_x = x

            layers = (
                'conv1_1', 'relu1_1', 'conv1_2', 'relu1_2', 'pool1',

                'conv2_1', 'relu2_1', 'conv2_2', 'relu2_2', 'pool2',

                'conv3_1', 'relu3_1', 'conv3_2', 'relu3_2', 'conv3_3',
                'relu3_3', 'pool3',

                'conv4_1', 'relu4_1', 'conv4_2', 'relu4_2', 'conv4_3',
                'relu4_3', 'pool4',

                'conv5_1', 'relu5_1', 'conv5_2', 'relu5_2', 'conv5_3',
                'relu5_3'
            )
            # x = self._batch_norm(x)
            # x = self._relu(x)
            net = {}
            current = x
            for i, name in enumerate(layers):
                kind = name[:4]
                if kind == 'conv':
                    kernels, bias = weights[i][0][0][0][0]
                    # matconvnet: weights are [width, height, in_channels, out_channels]
                    # tensorflow: weights are [height, width, in_channels, out_channels]
                    kernels = utils.get_variable(np.transpose(kernels, (1, 0, 2, 3)), name=name + "_w")
                    bias = utils.get_variable(bias.reshape(-1), name=name + "_b")
                    current = self._conv(current, kernels, bias, name)
                elif kind == 'relu':
                    current = self._relu(current, name=name)
                elif kind == 'pool':
                    current = self._max_pool(current, self.pool_size, self.pool_stride, name)
                net[name] = current

            return net


    def _conv(self, x, W, bias, name):
        """Convolution."""
        with tf.variable_scope(name) as scope:
            conv = tf.nn.conv2d(x, W, strides=[1, 1, 1, 1], padding="SAME")
            return tf.nn.bias_add(conv, bias)
            # bn = self._batch_norm(x)


    def _batch_norm(self, x):
        if self._data_format == 'channels_first':
            data_format = 'NCHW'
        else:
            data_format = 'NHWC'
        return tf.contrib.layers.batch_norm(
            x,
            decay=self._batch_norm_decay,
            center=True,
            scale=True,
            epsilon=self._batch_norm_epsilon,
            is_training=self._is_training,
            fused=True,
            data_format=data_format)

    def _relu(self, x, name):
        return tf.nn.relu(x, name=name)

    def _fully_connected(self, x, out_dim):
        with tf.name_scope('fully_connected') as name_scope:
            x = tf.layers.dense(x, out_dim)

        tf.logging.info('image after unit %s: %s', name_scope, x.get_shape())
        return x

    def _avg_pool(self, x, pool_size, stride):
        with tf.name_scope('avg_pool') as name_scope:
            x = tf.layers.average_pooling2d(
                x, pool_size, stride, 'SAME', data_format=self._data_format)

        tf.logging.info('image after unit %s: %s', name_scope, x.get_shape())
        return x

    def _max_pool(self, x, pool_size, stride, name):
        with tf.name_scope(name) as name_scope:
            x = tf.layers.max_pooling2d(x, pool_size, stride, 'SAME')
        tf.logging.info('image after unit %s: %s', name_scope, x.get_shape())
        return x

    def conv2d_transpose_strided(self, x, W, b, output_shape=None, stride=2):
        if output_shape is None:
            output_shape = x.get_shape().as_list()
            output_shape[1] *= 2
            output_shape[2] *= 2
            output_shape[3] = W.get_shape().as_list()[2]
        conv = tf.nn.conv2d_transpose(x, W, output_shape, strides=[1, stride, stride, 1], padding="SAME")
        return tf.nn.bias_add(conv, b)