def call(self, x, **kwargs):
"""
Build computation graph
Applies a convolution operation to 'x' for each dilation specified
in 'self.dilations' with shared kernel and bias weights
Processes 'x' through the attention layer mechanism as
att = softmax(conv2d(relu(conv2d(x))))
"""
# Evaluate the dilated conv ops
with tf.name_scope("dilated_conv_layers"):
outs = []
for i, op in enumerate(self.conv_ops):
with tf.name_scope("dilation_rate_%i_%i" % self.dilations[i]):
# Perform convolution with the same kernel across dilations
out = op(x, self.kernel)
if self.use_bias:
# Add bias if specified
cf = self.data_format == 'channels_first'
out = nn.bias_add(out,
self.bias,
data_format='NCHW' if cf else "NHWC")
if self.use_bn:
# Add BN layer if specified
bn = tf.keras.layers.BatchNormalization(
axis=self.channel_axis)
out = bn(out)
if self.activation is not None:
# Apply activation if specified
out = self.activation(out)
outs.append(out)
outs = tf.stack(outs, -1)
# Compute attention mechanism
with tf.name_scope("attention_mechanism"):
# Layer 1 (standard conv, any number of filters)
at1 = self.att_conv_1(x, self.att_K1)
if self.use_bias:
at1 = nn.bias_add(at1,
self.att_B1,
data_format='NCHW' if cf else "NHWC")
# Layer 2 (1x1 conv, 1 feature map pr. dilation)
at2 = self.att_conv_2(self.attention_activation(at1), self.att_K2)
if self.use_bias:
at2 = nn.bias_add(at2,
self.att_B2,
data_format='NCHW' if cf else "NHWC")
at_map = tf.nn.softmax(at2, name="attention_map")
# Compute attention weighted map
with tf.name_scope("weight_map"):
at_map = tf.expand_dims(at_map, axis=-2)
output = tf.reduce_sum(tf.multiply(at_map, outs), axis=-1)
if self.activation is not None:
# Activation function on output if specified
return self.activation(output)
return output
def core_ops_dense(inputs, kernel, bias, activation, dtype, units):
"""Add a GPU-compatible core ops dense function.
Adapted from the Tensorflow source code.
"""
rank = inputs.shape.rank
if rank is not None and rank > 2:
# Broadcasting is required for the inputs.
outputs = tf.tensordot(inputs, kernel, [[rank - 1], [0]])
# Reshape the output back to the original ndim of the input.
if not tf.executing_eagerly():
shape = inputs.shape.as_list()
output_shape = shape[:-1] + [units]
outputs.set_shape(output_shape)
else:
inputs = tf.cast(inputs, dtype)
if K.is_sparse(inputs):
outputs = tf.sparse.sparse_tensor_dense_matmul(inputs, kernel)
else:
outputs = tf.linalg.matmul(inputs, kernel)
if bias is not None:
outputs = nn.bias_add(outputs, bias)
if activation is not None:
return activation(outputs) # pylint: disable=not-callable
return outputs
def call(self, input):
masked_kernel = tf.math.multiply(self.mask, self.kernel)
x = nn.conv3d(input,
masked_kernel,
strides=[1, self.strides, self.strides, self.strides, 1],
padding=self.padding)
x = nn.bias_add(x, self.bias)
return x
def ConvReluMaxPool(X, weights, bias):
conv = nn.conv2d(X, weights, strides=[1, 1, 1, 1], padding="VALID", data_format="NHWC")
conv_bias = nn.bias_add(conv, bias, data_format="NHWC")
relu = nn.relu(conv_bias)
maxpool = nn.max_pool2d(relu, ksize=[1, 2, 2, 1], strides=[1, 2, 2, 1], padding="VALID", data_format="NHWC")
return maxpool
def Convolution(X, weights, bias):
conv = nn.conv2d(X,
weights,
strides=[1, 1, 1, 1],
padding="VALID",
data_format="NHWC")
conv_bias = nn.bias_add(conv, bias, data_format="NHWC")
return conv_bias
def prop_down(self, h, q=1, sig=True):
""" downwards mean-field propagation """
if K.is_sparse(h):
outputs = tf.matmul(h, tf.transpose(self.kernel), a_is_sparse=True)
else:
outputs = tf.matmul(h, tf.transpose(self.kernel))
if self.use_bias:
outputs = nn.bias_add(outputs, self.bias_visible)
return self.visible_activation(q*outputs)
def prop_up(self, v, q=1):
""" upwards mean-field propagation """
if K.is_sparse(v):
outputs = tf.matmul(v, self.kernel, a_is_sparse=True)
else:
outputs = tf.matmul(v, self.kernel)
if self.use_bias:
outputs = nn.bias_add(outputs, self.bias_hidden)
return self.hidden_activation(q*outputs)
def fully_connected(data, nInput, nOutput):
weights = tensorflow.Variable(
tensorflow.truncated_normal([nInput, nOutput], stddev=0.05))
biases = tensorflow.Variable(tensorflow.constant(0.05, shape=[nOutput]))
layer = tensorflow.matmul(data, weights)
layer = nn.bias_add(layer, biases)
layer = nn.relu(layer)
return layer
def convolution_layer(data, nChannels, filter_size, nFilters):
kernelShape = [filter_size, filter_size, nChannels, nFilters]
weights = tensorflow.Variable(
tensorflow.truncated_normal(kernelShape, stddev=0.05))
biases = tensorflow.Variable(tensorflow.constant(0.05, shape=[nFilters]))
layer = nn.conv2d(data, weights, [1, 1, 1, 1], 'SAME')
layer = nn.bias_add(layer, biases)
layer = nn.max_pool(layer, [1, 2, 2, 1], [1, 2, 2, 1], 'SAME')
layer = nn.relu(layer)
return layer
def vggBlock(X, weights1, bias1, weights2, bias2):
conv1 = nn.conv2d(X,
weights1,
strides=[1, 1, 1, 1],
padding="VALID",
data_format="NCHW")
conv1_bias = nn.bias_add(conv1, bias1, data_format="NCHW")
relu1 = nn.relu(conv1_bias)
conv2 = nn.conv2d(relu1,
weights2,
strides=[1, 1, 1, 1],
padding="SAME",
data_format="NCHW")
conv2_bias = nn.bias_add(conv2, bias2, data_format="NCHW")
relu2 = nn.relu(conv2_bias)
maxpool = nn.max_pool2d(relu2,
ksize=[1, 2, 2, 1],
strides=[1, 2, 2, 1],
padding="VALID",
data_format="NHWC")
return maxpool
def resnetBlock(X, weights1, bias1, weights2, bias2, mean1, variance1, offset1,
scale1, mean2, variance2, offset2, scale2):
conv1 = nn.conv2d(X,
weights1,
strides=[1, 1, 1, 1],
padding="VALID",
data_format="NCHW")
conv1_bias = nn.bias_add(conv1, bias1, data_format="NCHW")
bn1 = nn.batch_normalization(conv1_bias, mean1, variance1, offset1, scale1,
EPSILON)
relu1 = nn.relu(bn1)
conv2 = nn.conv2d(relu1,
weights2,
strides=[1, 1, 1, 1],
padding="SAME",
data_format="NCHW")
conv2_bias = nn.bias_add(conv2, bias2, data_format="NCHW")
bn2 = nn.batch_normalization(conv2_bias, mean2, variance2, offset2, scale2,
EPSILON)
return bn2
def ResizeConvReluMaxPool(X, weights, bias):
resize = image.resize(X, [N + 2, N + 2])
conv = nn.conv2d(resize,
weights,
strides=[1, 1, 1, 1],
padding="VALID",
data_format="NHWC")
conv_bias = nn.bias_add(conv, bias, data_format="NHWC")
relu = nn.relu(conv_bias)
maxpool = nn.max_pool2d(relu,
ksize=[1, 2, 2, 1],
strides=[1, 2, 2, 1],
padding="VALID",
data_format="NHWC",
name="output")
return maxpool
def build_model(self):
with tf.name_scope('Input'):
content_img = tf.constant(self.content_img_value,
name='content_image')
style_img = tf.constant(self.style_img_value, name='style_image')
self.magical_img = tf.Variable(initial_value=content_img,
name='magical_image')
input_tensor = tf.concat([
tf.expand_dims(content_img, axis=0),
tf.expand_dims(style_img, axis=0),
tf.expand_dims(self.magical_img, axis=0)
],
axis=0)
with tf.name_scope('conv1'):
conv1_1 = relu(bias_add(
conv2d(input_tensor,
tf.constant(self.vgg16_weights['block1_conv1'][0]),
strides=[1, 1, 1, 1],
padding='SAME'), self.vgg16_weights['block1_conv1'][1]),
name='conv1_1')
conv1_2 = relu(bias_add(
conv2d(conv1_1,
tf.constant(self.vgg16_weights['block1_conv2'][0]),
strides=[1, 1, 1, 1],
padding='SAME'), self.vgg16_weights['block1_conv2'][1]),
name='conv1_2')
pool1 = max_pool(conv1_2, [1, 2, 2, 1], [1, 2, 2, 1],
'SAME',
name='pool1')
with tf.name_scope('conv2'):
conv2_1 = relu(bias_add(
conv2d(pool1,
tf.constant(self.vgg16_weights['block2_conv1'][0]),
strides=[1, 1, 1, 1],
padding='SAME'), self.vgg16_weights['block2_conv1'][1]),
name='conv2_1')
conv2_2 = relu(bias_add(
conv2d(conv2_1,
tf.constant(self.vgg16_weights['block2_conv2'][0]),
strides=[1, 1, 1, 1],
padding='SAME'), self.vgg16_weights['block2_conv2'][1]),
name='conv2_2')
pool2 = max_pool(conv2_2, [1, 2, 2, 1], [1, 2, 2, 1],
'SAME',
name='pool2')
with tf.name_scope('conv3'):
conv3_1 = relu(bias_add(
conv2d(pool2,
tf.constant(self.vgg16_weights['block3_conv1'][0]),
strides=[1, 1, 1, 1],
padding='SAME'), self.vgg16_weights['block3_conv1'][1]),
name='conv3_1')
conv3_2 = relu(bias_add(
conv2d(conv3_1,
tf.constant(self.vgg16_weights['block3_conv2'][0]),
strides=[1, 1, 1, 1],
padding='SAME'), self.vgg16_weights['block3_conv2'][1]),
name='conv3_2')
conv3_3 = relu(bias_add(
conv2d(conv3_2,
tf.constant(self.vgg16_weights['block3_conv3'][0]),
strides=[1, 1, 1, 1],
padding='SAME'), self.vgg16_weights['block3_conv3'][1]),
name='conv3_3')
pool3 = max_pool(conv3_3, [1, 2, 2, 1], [1, 2, 2, 1],
'SAME',
name='pool3')
with tf.name_scope('conv4'):
conv4_1 = relu(bias_add(
conv2d(pool3,
tf.constant(self.vgg16_weights['block4_conv1'][0]),
strides=[1, 1, 1, 1],
padding='SAME'), self.vgg16_weights['block4_conv1'][1]),
name='conv4_1')
conv4_2 = relu(bias_add(
conv2d(conv4_1,
tf.constant(self.vgg16_weights['block4_conv2'][0]),
strides=[1, 1, 1, 1],
padding='SAME'), self.vgg16_weights['block4_conv2'][1]),
name='conv4_2')
conv4_3 = relu(bias_add(
conv2d(conv4_2,
tf.constant(self.vgg16_weights['block4_conv3'][0]),
strides=[1, 1, 1, 1],
padding='SAME'), self.vgg16_weights['block4_conv3'][1]),
name='conv4_3')
pool4 = max_pool(conv4_3, [1, 2, 2, 1], [1, 2, 2, 1],
'SAME',
name='pool4')
with tf.name_scope('conv5'):
conv5_1 = relu(bias_add(
conv2d(pool4,
tf.constant(self.vgg16_weights['block5_conv1'][0]),
strides=[1, 1, 1, 1],
padding='SAME'), self.vgg16_weights['block5_conv1'][1]),
name='conv5_1')
conv5_2 = relu(bias_add(
conv2d(conv5_1,
tf.constant(self.vgg16_weights['block5_conv2'][0]),
strides=[1, 1, 1, 1],
padding='SAME'), self.vgg16_weights['block5_conv2'][1]),
name='conv5_2')
conv5_3 = relu(bias_add(
conv2d(conv5_2,
tf.constant(self.vgg16_weights['block5_conv3'][0]),
strides=[1, 1, 1, 1],
padding='SAME'), self.vgg16_weights['block5_conv3'][1]),
name='conv5_3')
pool5 = max_pool(conv5_3, [1, 2, 2, 1], [1, 2, 2, 1],
'SAME',
name='pool5')
c_loss = self.content_loss([input_tensor])
s_loss = self.style_loss([conv1_1, conv2_1, conv3_1, conv4_1, conv5_1])
self.loss = self.content_weight * c_loss + self.style_weight * s_loss
self.train_op = tf.train.AdamOptimizer(learning_rate=self.lr).minimize(
self.loss)
if self.graph_write:
writer = tf.summary.FileWriter('logs',
graph=tf.get_default_graph())
writer.flush()
writer.close()
def call(self, inputs, training=None):
conv = conv2d(inputs, self.w, self.strides, self.padding)
if self.use_bias:
return bias_add(conv, self.b)
return conv
def transformerShift(
inputs,
memory,
conv_units,
short_units,
compression_rate: int = 2,
method: str = 'conv',
use_bias: bool = True,
kernel=None,
bias=None,
activation=None,
):
"""Compressive Transformer left-shift operation with fixed short and convolution memories
# Arguments
inputs: A +3D Tensor with shape: `(timesteps, ..., features)`.
memory: A +3D tensor with shape: `(conv_units+short_units, ..., features)`.
conv_units: int > 0, Number of convolution memory units in the layer.
short_units: int > 0, Number of short-term memory units in the layer.
compression_rate: int >= 1 (default - 2). Compression rate for the selected compression method.
method: str [`conv`, `mean`, `max`], Defines a memory compression function (default - conv).
Available compression functions:
- `conv`: 1D convolution over memory sequence.
- `mean`: 1D Mean keras.layers.AveragePooling1D.
- `max`: 1D Max Pooling via keras.MaxPooling1D.
Only for method=`conv`:
use_bias: Boolean (default - True), whether the layer uses a bias vector.
kernel: A 3D Tensor with shape: `(compression_rate, features, features)`.
bias: A Tensor of shape `(features,)`
activation: (default - None) Activation functon applied to the output of the compression function ( see keras.activations).
# Inputs Shape
+3D tensor with shape: `(batch_size, timesteps, ..., features)`,
+3D tensor with shape: `(conv_units+short_units, ..., features)`.
# Output Shape
+3D tensor with shape: `(batch_size, conv_units+short_units, ..., features)`.
# References
- ["Compressive Transformers for Long-Range Sequence Modelling" by Rae et. al.](https://arxiv.org/abs/1911.05507)
"""
if method not in ['conv', 'mean', 'max']:
raise ValueError(f'Compression method `{method}` is not supported.')
if len(inputs.shape) < 3:
raise ValueError(
'The dimension of the input vector'
' should be at least 3D: `(batch_size, time_steps, ..., features)`'
)
if len(memory.shape) < 2:
raise ValueError(
'The dimension of the memory vector'
' should be at least 2D: `(conv_units+short_units, ..., features)`'
)
if len(inputs.shape) < 3:
raise ValueError(
'The dimensions of the first tensor of the inputs to `transformerShift` '
f'should be at least 3D: `(batch_size, timesteps, ..., features)`. Found `{inputs.shape}`.'
)
if len(memory.shape) < 2:
raise ValueError(
'The dimensions of the second tensor of the memory to `transformerShift` '
f'should be at least 2D: `(conv_units+short_units, ..., features)`. Found `{memory.shape}`.'
)
if tensor_shape.dimension_value(inputs.shape[-1]) is None:
raise ValueError(
'The last dimension of the first tensor of the inputs to `transformerShift` '
'should be defined. Found `None`.')
if tensor_shape.dimension_value(memory.shape[-1]) is None:
raise ValueError(
'The last dimension of the second tensor of the inputs (memory) to `transformerShift` '
'should be defined. Found `None`.')
if tensor_shape.dimension_value(inputs.shape[-1]) is not \
tensor_shape.dimension_value(memory.shape[-1]):
raise ValueError(
'The last dimension of both input tensors to `transformerShift` '
f'should match. Found `{tensor_shape.dimension_value(inputs.shape[-1])}` and '
f'`{tensor_shape.dimension_value(memory.shape[-1])}`.')
if compression_rate < 1:
raise ValueError('Compression rate cannot be less than 1.')
if method in ['max', 'mean']:
compression_rate = np.int(np.floor(compression_rate))
activation = activations.get(activation)
timesteps = inputs.shape[1]
memories = []
batch_memory = tf.reshape(memory, (1, ) + memory.shape)
for i, batch in enumerate(inputs):
# Forgetting the oldest
conv_mem = batch_memory[:, :conv_units]
short_mem = batch_memory[:, conv_units:]
old_memory = short_mem[:, :timesteps, ...]
# Compressing the oldest memories
if method is 'conv':
if kernel.shape[1:-1] != batch.shape[1:]:
raise ValueError(
'The dimension of the kernel should be 3D: `(compression_rate, features, features)` '
'if shape of the input is 3D: `(batch_size, time_steps, features)`. '
f'Found `{kernel.shape[1:-1]}`!=`{batch.shape[1:]}`')
compression = nn.conv1d(old_memory,
filters=kernel,
stride=compression_rate,
padding='SAME',
name='Compression')
if use_bias:
if bias.shape[-1] != compression.shape[-1]:
raise ValueError(
'The dimension of the bias'
' should be equal to the number of features. '
f'Found `{bias.shape[-1]}`!=`{kernel.shape[-1]}`')
compression = nn.bias_add(compression, bias)
elif method is 'mean':
compression = nn.avg_pool1d(old_memory,
compression_rate,
strides=compression_rate,
padding='SAME',
name='Compression')
elif method is 'max':
compression = nn.max_pool1d(old_memory,
compression_rate,
strides=compression_rate,
padding='SAME',
name='Compression')
else:
compression = nn.avg_pool1d(old_memory,
compression_rate,
strides=compression_rate,
padding='SAME',
name='Compression')
if compression.shape[-1] > conv_mem.shape[-1]:
compression = nn.avg_pool1d(compression,
compression_rate,
strides=compression_rate,
padding='SAME')
if activation is not None:
compression = activation(compression)
# Update Short-Term Memory
short_mem = K.concatenate([
short_mem[:, self.units:],
tf.reshape(batch, (1, ) + batch.shape)
],
axis=1)
# Update Compressed Memory
conv_mem = K.concatenate([conv_mem[:, 1:, ...], compression], axis=1)
batch_memory = K.concatenate([conv_mem, short_mem], axis=1)
memories.append(batch_memory)
return K.concatenate(memories, axis=0)
