conv-nmt/temporal_encoder_decoder.py at master · oneTimePad/conv-nmt · 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
import tensorflow as tf

class TemporalEncoderBlock(tf.layers.Layer):
    """Temporal Convolution Encoder Block
    """

    def __init__(self, num_channels,
                 kernel_size=2,
                 dropout=0.2,
                 trainable=True,
                 name=None,
                 dtype=None,
                 activation=tf.nn.relu6,
                 activity_regularizer=None,
                 **kwargs):
        super(TemporalEncoderBlock, self).__init__(trainable=trainable,
            dtype=dtype,
            activity_regularizer=activity_regularizer,
            name=name,
            **kwargs)

        self.temporal_conv = tf.layers.Conv1D(filters=num_channels,
            kernel_size=kernel_size,
            activation=activation,
            padding="SAME")

        self.batch_norm = tf.layers.BatchNormalization()

        self.pooled = tf.layers.MaxPooling1D(pool_size=2, strides=1)

    def call(self, inputs, training=True):

        temporal_conv = self.temporal_conv(inputs)
        batch_norm = self.batch_norm(temporal_conv)
        pooled = self.pooled(batch_norm)
        return pooled

class TemporalDecoderBlock(tf.layers.Layer):

    def __init__(self, num_channels,
                 kernel_size=2,
                 trainable=True,
                 name=None,
                 dtype=None,
                 activation=None,
                 activity_regularizer=None,
                 **kwargs):
        super(TemporalDecoderBlock, self).__init__(trainable=trainable,
            dtype=dtype,
            activity_regularizer=activity_regularizer,
            name=name,
            **kwargs)

        self.temporal_conv = tf.layers.Conv1D(filters=num_channels,
            kernel_size=kernel_size,
            activation=activation,
            padding="SAME")

        self.unpooled = tf.keras.layers.UpSampling1D(size=2)

    def call(self, inputs, training=True):

        temporal_conv = self.temporal_conv(inputs)
        unpooled = self.unpooled(temporal_conv)
        return unpooled

class TemporalEncoder(tf.layers.Layer):

    """Temporal Encoder Network"""

    def __init__(self, num_blocks,
                 trainable = True,
                 name = None,
                 dtype = None,
                 activity_regularizer = None,
                 **kwargs):
        super(TemporalEncoder, self).__init__(trainable = trainable,
                                        dtype = dtype,
                                        activity_regularizer = activity_regularizer,
                                        name = name,
                                        **kwargs)

        self._blocks = []

        for i in range(num_blocks):
            self._blocks.extend([TemporalEncoderBlock(num_channels=(32 * (2 ^ (i + 1))))])

    def call(self, inputs, training=True):
        output = inputs
        for block in self._blocks:
            output = block(output)
        return output

class TemporalDecoder(tf.layers.Layer):

    """Temporal Decoder Network"""

    def __init__(self, num_encoder_blocks,
                  num_blocks,
                  trainable = True,
                  name = None,
                  dtype = None,
                  activity_regularizer = None,
                  **kwargs):
         super(TemporalDecoder, self).__init__(trainable = trainable,
                                         dtype = dtype,
                                         activity_regularizer = activity_regularizer,
                                         name = name,
                                         **kwargs)

         self._blocks = []
         last_filter_size = 32 * (num_encoder_blocks)
         for i in range(num_blocks):
             self._blocks.extend([TemporalDecoderBlock(num_channels=last_filter_size / (2 * (i + 1)))])

    def call(self, inputs, training=True):
        output = inputs
        for block in self._blocks:
            output = block(output)
        return output