Layer.py 4.65 KB
Newer Older
Tiago de Freitas Pereira's avatar
Tiago de Freitas Pereira committed
1
2
3
4
5
#!/usr/bin/env python
# vim: set fileencoding=utf-8 :
# @author: Tiago de Freitas Pereira <tiago.pereira@idiap.ch>
# @date: Wed 11 May 2016 17:38 CEST

Tiago de Freitas Pereira's avatar
Tiago de Freitas Pereira committed
6
import tensorflow as tf
7
8
from bob.learn.tensorflow.initialization import Xavier
from bob.learn.tensorflow.initialization import Constant
Tiago de Freitas Pereira's avatar
Tiago de Freitas Pereira committed
9
10
11
12
13
14
15
16


class Layer(object):

    """
    Layer base class
    """

17
18
19
20
    def __init__(self, name,
                 activation=None,
                 weights_initialization=Xavier(),
                 bias_initialization=Constant(),
21
                 batch_norm=False,
22
                 use_gpu=False):
Tiago de Freitas Pereira's avatar
Tiago de Freitas Pereira committed
23
24
25
26
        """
        Base constructor

        **Parameters**
27
28
          name: Name of the layer
          activation: Tensorflow activation operation (https://www.tensorflow.org/versions/r0.10/api_docs/python/nn.html)
29
30
          weights_initialization: Initialization for the weights
          bias_initialization: Initialization for the biases
31
32
          use_gpu: I think this is not necessary to explain
          seed: Initialization seed set in Tensor flow
Tiago de Freitas Pereira's avatar
Tiago de Freitas Pereira committed
33
        """
Tiago de Freitas Pereira's avatar
Tiago de Freitas Pereira committed
34
        self.name = name
35
36
        self.weights_initialization = weights_initialization
        self.bias_initialization = bias_initialization
Tiago de Freitas Pereira's avatar
Tiago de Freitas Pereira committed
37
        self.use_gpu = use_gpu
38
        self.batch_norm = batch_norm
Tiago de Freitas Pereira's avatar
Tiago de Freitas Pereira committed
39

Tiago de Freitas Pereira's avatar
Tiago de Freitas Pereira committed
40
41
        self.input_layer = None
        self.activation = activation
Tiago de Freitas Pereira's avatar
Tiago de Freitas Pereira committed
42

43
44
45
        # Batch normalization variables
        self.beta = None
        self.gamma = None
46
47
        self.batch_mean = None
        self.batch_var = None
48

Tiago de Freitas Pereira's avatar
Tiago de Freitas Pereira committed
49
    def create_variables(self, input_layer):
Tiago de Freitas Pereira's avatar
Tiago de Freitas Pereira committed
50
51
        NotImplementedError("Please implement this function in derived classes")

52
    def get_graph(self, training_phase=True):
Tiago de Freitas Pereira's avatar
Tiago de Freitas Pereira committed
53
        NotImplementedError("Please implement this function in derived classes")
54

55
56
57
    def variable_exist(self, var):
        return var in [v.name.split("/")[0] for v in tf.all_variables()]

58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
    def batch_normalize(self, x, phase_train):
        """
        Batch normalization on convolutional maps.
        Args:
            x:           Tensor, 4D BHWD input maps
            n_out:       integer, depth of input maps
            phase_train: boolean tf.Variable, true indicates training phase
            scope:       string, variable scope
            affn:      whether to affn-transform outputs
        Return:
            normed:      batch-normalized maps
        Ref: http://stackoverflow.com/questions/33949786/how-could-i-use-batch-normalization-in-tensorflow/33950177
        """
        from tensorflow.python.ops import control_flow_ops

73
        name = "batch_norm_" + str(self.name)
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
        reuse = self.variable_exist(name)

        #if reuse:
            #import ipdb; ipdb.set_trace();

        with tf.variable_scope(name, reuse=reuse):

            phase_train = tf.convert_to_tensor(phase_train, dtype=tf.bool)
            n_out = int(x.get_shape()[-1])

            self.beta = tf.get_variable(name + '_beta',
                                        initializer=tf.constant(0.0, shape=[n_out], dtype=x.dtype),
                                        trainable=True,
                                        dtype=x.dtype)

            self.gamma = tf.get_variable(name + '_gamma',
                                         initializer=tf.constant(1.0, shape=[n_out], dtype=x.dtype),
                                         trainable=True,
                                         dtype=x.dtype)

            if len(x.get_shape()) == 2:
                self.batch_mean, self.batch_var = tf.nn.moments(x, [0], name='moments_{0}'.format(name))
            else:
                self.batch_mean, self.batch_var = tf.nn.moments(x, range(len(x.get_shape())-1), name='moments_{0}'.format(name))

            ema = tf.train.ExponentialMovingAverage(decay=0.9)

            def mean_var_with_update():
                ema_apply_op = ema.apply([self.batch_mean, self.batch_var])
                with tf.control_dependencies([ema_apply_op]):
                    return tf.identity(self.batch_mean), tf.identity(self.batch_var)

            mean, var = control_flow_ops.cond(phase_train,
                                                        mean_var_with_update,
                                                        lambda: (ema.average(self.batch_mean), ema.average(self.batch_var)),
                                                        name=name + "mean_var")

            normed = tf.nn.batch_normalization(x, mean, var, self.beta, self.gamma, 1e-3)
112
        return normed
113
114
115
116
117
118
119
120
121
122
123

    def get_varible_by_name(self, var):
        """
        Doing this because of that https://github.com/tensorflow/tensorflow/issues/1325
        """

        for v in tf.all_variables():
            if (len(v.name.split("/")) > 1) and (var in v.name.split("/")[1]):
                return v

        return None