Spaces:

akhaliq
/

deeplab2

Runtime error

deeplab2 / model /layers /axial_layers.py

akhaliq3

spaces demo

506da10 almost 3 years ago

No virus

22.6 kB

	# coding=utf-8
	# Copyright 2021 The Deeplab2 Authors.
	#
	# Licensed under the Apache License, Version 2.0 (the "License");
	# you may not use this file except in compliance with the License.
	# You may obtain a copy of the License at
	#
	# http://www.apache.org/licenses/LICENSE-2.0
	#
	# Unless required by applicable law or agreed to in writing, software
	# distributed under the License is distributed on an "AS IS" BASIS,
	# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	# See the License for the specific language governing permissions and
	# limitations under the License.

	"""Implements Axial-Attention layers proposed in Axial-DeepLab.

	Axial-Attention factorizes 2D self-attention into two 1D self-attentions, so
	that it can be applied on large inputs. Axial-Attention is typically used to
	replace 3x3 convolutions in a bottleneck residual block.

	[1] Axial-Deeplab: Stand-Alone Axial-Attention for Panoptic Segmentation,
	ECCV 2020 Spotlight.
	Huiyu Wang, Yukun Zhu, Bradley Green, Hartwig Adam, Alan Yuille,
	Liang-Chieh Chen.
	"""

	import numpy as np
	import tensorflow as tf

	from deeplab2.model import utils
	from deeplab2.model.layers import activations
	from deeplab2.model.layers import positional_encodings


	class AxialAttention(tf.keras.layers.Layer):
	"""An axial-attention layer."""

	def __init__(self,
	query_shape=129,
	memory_flange=32,
	total_key_depth=512,
	total_value_depth=1024,
	num_heads=8,
	name='axial_attention',
	use_query_rpe_similarity=True,
	use_key_rpe_similarity=True,
	use_content_similarity=True,
	retrieve_value_rpe=True,
	retrieve_value_content=True,
	initialization_std_for_query_key_rpe=1.0,
	initialization_std_for_value_rpe=1.0,
	self_attention_activation='softmax',
	bn_layer=tf.keras.layers.BatchNormalization,
	conv_kernel_weight_decay=0.0):
	"""Initializes an axial-attention layer.

	This function is designed to support both global and local axial-attention
	in a unified way. If query_shape is larger than the length of input, a
	global attention is applied. If query_shape is smaller than the length of
	input, a local attention is applied. In this case, the input is divided into
	blocks of length query_shape, padded by memory_flange on both sides. Then,
	local attention is applied within each query block. The choice of
	query_shape does not affect the output value but affects computation
	efficiency and memory usage. In general, use global attention (large
	query_shape) if possible. Local axial-attention has not been supported yet.

	Args:
	query_shape: An integer, the block size for local axial attention.
	Defaults to 129 since 129 is usually the largest feature map where we do
	global attention (1025 with stride 8, or 2049 with stride 16).
	memory_flange: An integer, the memory flange padded to each query block in
	local attention. It has no effect in global attention. Defaults to 32,
	which is equivalent to a span of 65 in Aixal-DeepLab paper -- A pixel
	can see 32 pixels on its left and 32 pixels on its right.
	total_key_depth: An integer, the total depth of keys, which is also the
	depth of queries and the depth of key (query) positional encodings.
	total_value_depth: An integer, the total depth of the values, which is
	also the depth of value positional encodings.
	num_heads: An integer, the number of heads in multi-head attention.
	name: A string, the name of this axial attention layer.
	use_query_rpe_similarity: A boolean, whether to use the attention
	similarity between the queries and the relative positional encodings.
	use_key_rpe_similarity: A boolean, whether to use the attention similarity
	between the keys and the relative positional encodings.
	use_content_similarity: A boolean, whether to use the content similarity
	between the queries and the keys.
	retrieve_value_rpe: A boolean, whether to retrieve the relative positional
	encodings of the values.
	retrieve_value_content: A boolean, whether to retrieve the content of the
	values.
	initialization_std_for_query_key_rpe: A float, the initialization std for
	the relative positional encodings of the queries and keys.
	initialization_std_for_value_rpe: A float, the initialization std for the
	relative positional encodings of the values.
	self_attention_activation: A string, type of activation function for
	self-attention. Support 'sigmoid' and 'softmax'.
	bn_layer: A tf.keras.layers.Layer that computes the normalization
	(default: tf.keras.layers.BatchNormalization).
	conv_kernel_weight_decay: A float, the weight decay for convolution
	kernels.

	Returns:
	output: A [batch, length, total_value_depth] tensor.

	Raises:
	ValueError: If none of the three similarities (use_query_rpe_similarity,
	use_key_rpe_similarity, use_content_similarity) is used.
	ValueError: If neither of value content or value rpe is retrieved.
	ValueError: If self_attention_activation is not supported.
	ValueError: If total_key_depth is not divisible by num_heads.
	ValueError: If total_value_depth is not divisible by num_heads.
	"""
	# Validate the attention similarity choices.
	if not any([
	use_content_similarity, use_key_rpe_similarity, use_query_rpe_similarity
	]):
	raise ValueError(
	'Should use at least one similarity to compute attention.')

	# Validate the retrieve value choices.
	if not retrieve_value_content and not retrieve_value_rpe:
	raise ValueError('Should retrieve at least one of content or rpe.')

	if total_key_depth % num_heads:
	raise ValueError('Total_key_depth should be divisible by num_heads.')

	if total_value_depth % num_heads:
	raise ValueError('Total_value_depth should be divisible by num_heads.')

	super(AxialAttention, self).__init__(name=name)
	self._query_shape = query_shape
	self._memory_flange = memory_flange
	self._total_key_depth = total_key_depth
	self._total_value_depth = total_value_depth
	self._num_heads = num_heads
	self._use_query_rpe_similarity = use_query_rpe_similarity
	self._use_key_rpe_similarity = use_key_rpe_similarity
	self._use_content_similarity = use_content_similarity
	self._retrieve_value_rpe = retrieve_value_rpe
	self._retrieve_value_content = retrieve_value_content
	self._initialization_std_for_query_key_rpe = (
	initialization_std_for_query_key_rpe)
	self._initialization_std_for_value_rpe = initialization_std_for_value_rpe
	self._self_attention_activation = self_attention_activation
	self._conv_kernel_weight_decay = conv_kernel_weight_decay

	self._batch_norm_qkv = bn_layer(axis=-1, name='batch_norm_qkv')
	self._batch_norm_similarity = bn_layer(
	axis=[0, 2], name='batch_norm_similarity')
	self._batch_norm_retrieved_output = bn_layer(
	axis=[0, 2, 4], name='batch_norm_retrieved_output')

	self._key_depth_per_head = total_key_depth // num_heads
	self._attention_activate_fn = activations.get_activation(
	self_attention_activation)

	def build(self, input_shape):
	"""Builds axial-attention layer weights.

	Args:
	input_shape: An integer list of length 3, the shape of the input tensor.

	Raises:
	NotImplementedError: Local axial-attention has not been implemented. It is
	triggered if query_shape is less than input_shape.
	"""

	# Apply global attention if query_shape is larger than input_shape[1].
	if self._query_shape >= input_shape[1]:
	self._query_shape = input_shape[1]
	self._memory_flange = 0
	else:
	raise NotImplementedError('Local axial attention has not been '
	'implemented yet.')
	self._memory_shape = self._query_shape + 2 * self._memory_flange

	# Compute query key value with one convolution and an optional batch norm.
	# The initialization std is standard transformer initialization (without
	# batch norm), as used in SASA and ViT. In our case, we use batch norm by
	# default, so it does not require careful tuning. If one wants to remove
	# all batch norms in axial attention, this standard initialization should
	# still be good, but a more careful initialization is encouraged.
	self.qkv_kernel = self.add_weight(
	name='qkv_kernel',
	shape=[input_shape[-1],
	self._total_key_depth * 2 + self._total_value_depth],
	initializer=tf.keras.initializers.TruncatedNormal(
	stddev=input_shape[-1]**-0.5),
	regularizer=tf.keras.regularizers.l2(self._conv_kernel_weight_decay))

	if self._use_query_rpe_similarity:
	self._query_rpe = positional_encodings.RelativePositionalEncoding(
	self._query_shape,
	self._memory_shape,
	self._key_depth_per_head,
	self._num_heads,
	'query_rpe',
	initialization_std=self._initialization_std_for_query_key_rpe,
	conv_kernel_weight_decay=self._conv_kernel_weight_decay)

	if self._use_key_rpe_similarity:
	self._key_rpe = positional_encodings.RelativePositionalEncoding(
	self._query_shape,
	self._memory_shape,
	self._key_depth_per_head,
	self._num_heads,
	'key_rpe',
	initialization_std=self._initialization_std_for_query_key_rpe,
	conv_kernel_weight_decay=self._conv_kernel_weight_decay)

	if self._retrieve_value_rpe:
	self._value_rpe = positional_encodings.RelativePositionalEncoding(
	self._query_shape,
	self._memory_shape,
	self._total_value_depth // self._num_heads,
	self._num_heads,
	'value_rpe',
	initialization_std=self._initialization_std_for_value_rpe,
	conv_kernel_weight_decay=self._conv_kernel_weight_decay)

	def call(self, input_tensor, training=False):
	"""Performs a forward pass.

	Args:
	input_tensor: An input [batch, length, channel] tensor.
	training: A boolean flag indicating whether training behavior should be
	used (default: False).

	Returns:
	output: An output [batch, length, total_value_depth] tensor.
	"""
	# Alternatively, the einsum can be implemented as a 1x1 convolution.
	# However, it is not obvious which implementation is more efficient (without
	# careful benchmarking), so we use einsum for its flexibility and
	# consistency with other parts of the function.
	query_key_value = tf.einsum(
	'nlc,cd->nld', input_tensor, self.qkv_kernel, name='compute_qkv')
	query_key_value = self._batch_norm_qkv(query_key_value, training=training)

	# Split query key value.
	query, key, value = tf.split(
	query_key_value,
	[self._total_key_depth, self._total_key_depth, self._total_value_depth],
	axis=-1)

	# Reshape the query, key, and value.
	query = tf.reshape(query, [-1, self._query_shape, self._num_heads,
	self._key_depth_per_head])
	query = tf.transpose(a=query, perm=[0, 2, 1, 3])
	key = tf.reshape(key, [-1, np.prod(self._memory_shape), self._num_heads,
	self._key_depth_per_head])
	key = tf.transpose(a=key, perm=[0, 2, 1, 3])
	value = tf.reshape(value, [-1, np.prod(self._memory_shape), self._num_heads,
	self._total_value_depth // self._num_heads])

	# Gather all similarity logits into a list.
	similarity_logits = []

	# Compute the content similarity term: q * k.
	if self._use_content_similarity:
	content_similarity = tf.einsum(
	'bhld,bhmd->bhlm', query, key, name='content_similarity')
	similarity_logits.append(content_similarity)

	# Compute the query rpe similarity term: q * rpe.
	if self._use_query_rpe_similarity:
	query_rpe = self._query_rpe(None)
	query_rpe_similarity = tf.einsum(
	'bhld,hlmd->bhlm', query, query_rpe, name='query_rpe_similarity')
	similarity_logits.append(query_rpe_similarity)

	# Compute the key rpe similarity term: k * rpe.
	if self._use_key_rpe_similarity:
	key_rpe = self._key_rpe(None)
	key_rpe_similarity = tf.einsum(
	'bhmd,hlmd->bhlm', key, key_rpe, name='key_rpe_similarity')
	similarity_logits.append(key_rpe_similarity)

	# Apply an optional batch norm to the similarities and sum them.
	similarity_logits = tf.stack(similarity_logits)
	similarity_logits = self._batch_norm_similarity(similarity_logits,
	training=training)
	similarity_logits = tf.reduce_sum(input_tensor=similarity_logits, axis=0)

	# Apply an attention activation function, e.g. softmax.
	weights = self._attention_activate_fn(similarity_logits)

	# Gather retrieved values or rpes into a list.
	retrieve_list = []

	# Retrieve the content of the attended value.
	if self._retrieve_value_content:
	retrieved_content = tf.einsum(
	'bhlm,bmhd->bhld', weights, value, name='retrieve_value_content')
	retrieve_list.append(retrieved_content)

	# Retrieve the relative position of the attended value.
	if self._retrieve_value_rpe:
	value_rpe = self._value_rpe(None)
	retrieved_rpe = tf.einsum(
	'bhlm,hlmd->bhld', weights, value_rpe, name='retrieve_value_rpe')
	retrieve_list.append(retrieved_rpe)

	# Apply batch norms to retrieved contents and rpes respectively.
	retrieved_output = tf.stack(retrieve_list)
	retrieved_output = self._batch_norm_retrieved_output(retrieved_output,
	training=training)
	# Additive contents and rpes.
	retrieved_output = tf.reduce_sum(input_tensor=retrieved_output, axis=0)

	# Combine the heads by transposing and reshaping the tensor.
	retrieved_output = utils.transpose_and_reshape_for_attention_operation(
	retrieved_output)

	return retrieved_output


	class AxialAttention2D(tf.keras.layers.Layer):
	"""Sequentially applies height-axis and width-axis axial-attention."""

	def __init__(self,
	strides=1,
	filters=512,
	name='attention',
	key_expansion=1,
	value_expansion=2,
	query_shape=(129, 129),
	memory_flange=(32, 32),
	**kwargs):
	"""Initializes an AxialAttention2D layer.

	Args:
	strides: An integer, the stride for the output, usually 1 or 2.
	filters: An integer, the base number of channels for the layer.
	name: A string, the name of the attention layer.
	key_expansion: A float, the channel expansion ratio for keys.
	value_expansion: A float, the channel expansion ratio for values.
	query_shape: An integer, the maximum query shape for both the height axis
	and the width axis.
	memory_flange: An integer list of length 2. The memory flange for the
	height axis and the width axis.
	**kwargs: A dictionary of keyword arguments passed to height-axis,
	width-axis, and 2D global AxialAttention.

	Returns:
	output: A [batch, strided height, strided width, output_channels] tensor.
	"""
	super(AxialAttention2D, self).__init__(name=name)
	total_key_depth = int(round(filters * key_expansion))
	total_value_depth = int(round(filters * value_expansion))
	self._strides = strides
	self._total_key_depth = total_key_depth
	self._total_value_depth = total_value_depth
	self._height_axis = AxialAttention(
	total_key_depth=total_key_depth,
	total_value_depth=total_value_depth,
	query_shape=query_shape[0],
	memory_flange=memory_flange[0],
	name='height_axis',
	**kwargs)
	self._width_axis = AxialAttention(
	total_key_depth=total_key_depth,
	total_value_depth=total_value_depth,
	query_shape=query_shape[1],
	memory_flange=memory_flange[1],
	name='width_axis',
	**kwargs)

	def call(self, inputs, training=False):
	"""Performs a forward pass.

	Args:
	inputs: An input [batch, height, width, channel] tensor.
	training: A boolean flag indicating whether training behavior should be
	used (default: False).

	Returns:
	output: An output [batch, strided_height, strided_width,
	filters * value_expansion] tensor.
	"""
	_, height, width, channel = inputs.get_shape().as_list()

	# Transpose and reshape the width axis to the batch dimension.
	x = tf.transpose(a=inputs, perm=[0, 2, 1, 3])
	x = tf.reshape(x, [-1, height, channel])
	x = self._height_axis(x, training=training)
	# Reshape and transpose back to a 4D tensor.
	x = tf.reshape(x, [-1, width, height, self._total_value_depth])
	x = tf.transpose(a=x, perm=[0, 2, 1, 3])
	# Height axis striding.
	if self._strides > 1:
	x = x[:, ::self._strides, :, :]

	# Reshape the height axis to the batch dimension.
	_, strided_height, _, _ = x.get_shape().as_list()
	x = tf.reshape(x, [-1, width, self._total_value_depth])
	x = self._width_axis(x, training=training)
	# Reshape back to a 4D tensor.
	x = tf.reshape(x, [-1, strided_height, width, self._total_value_depth])
	# Width axis striding.
	if self._strides > 1:
	x = x[:, :, ::self._strides, :]

	return x


	class GlobalAttention2D(tf.keras.layers.Layer):
	"""A 2D global attention layer."""

	def __init__(self,
	strides=1,
	filters=512,
	name='attention',
	key_expansion=1,
	value_expansion=2,
	query_shape=(129, 129),
	memory_flange=(32, 32),
	double_global_attention=False,
	**kwargs):
	"""Initializes a GlobalAttention2D layer.

	Args:
	strides: An integer, the stride for the output, usually 1 or 2.
	filters: An integer, the base number of channels for the layer.
	name: A string, the name of the attention layer.
	key_expansion: A float, the channel expansion ratio for keys.
	value_expansion: A float, the channel expansion ratio for values.
	query_shape: An integer, the maximum query shape for both the height axis
	and the width axis.
	memory_flange: An integer list of length 2. The memory flange for the
	height axis and the width axis.
	double_global_attention: A boolean, whether to use two global attention
	layers. Two global attention layers match the parameter count to a
	seqentially applied height and width axial attention layer.
	**kwargs: A dictionary of keyword arguments passed to height-axis,
	width-axis, and 2D global AxialAttention.

	Returns:
	output: A [batch, strided height, strided width, output_channels] tensor.

	Raises:
	ValueError: If relative positional encoding is enforced in kwargs.
	"""
	if any([kwargs.get('use_query_rpe_similarity', False),
	kwargs.get('use_key_rpe_similarity', False),
	kwargs.get('retrieve_value_rpe', False)]):
	raise ValueError('GlobalAttention2D does not support relative positional '
	'encodings.')

	super(GlobalAttention2D, self).__init__(name=name)
	total_key_depth = int(round(filters * key_expansion))
	total_value_depth = int(round(filters * value_expansion))
	self._strides = strides
	self._double_global_attention = double_global_attention
	self._total_key_depth = total_key_depth
	self._total_value_depth = total_value_depth

	# Global attention does not support relative positional encodings.
	kwargs['use_query_rpe_similarity'] = False
	kwargs['use_key_rpe_similarity'] = False
	kwargs['retrieve_value_rpe'] = False
	self._kwargs = kwargs

	def build(self, input_shape):
	"""Builds global attention layers according to the 4D input_shape."""
	_, height, width, _ = input_shape
	# Implement 2D global attention as 1D axial-attention by flattening the 2D
	# inputs into 1D. We also disable the relative positional encodings in
	# axial attention, so that only content-based attention is used. The query
	# shape is set to height * width, so that the axial attention is global.
	self._global = AxialAttention(
	total_key_depth=self._total_key_depth,
	total_value_depth=self._total_value_depth,
	query_shape=height*width,
	memory_flange=0,
	name='global',
	**self._kwargs)

	# Use two global attention layers in one residual block. This option
	# ensures that global attention models have similar number of layers and
	# parameters as axial-attention models.
	if self._double_global_attention:
	self._global2 = AxialAttention(
	total_key_depth=self._total_key_depth,
	total_value_depth=self._total_value_depth,
	query_shape=height*width,
	memory_flange=0,
	name='global2',
	**self._kwargs)

	def call(self, inputs, training=False):
	"""Performs a forward pass.

	Args:
	inputs: An input [batch, height, width, channel] tensor.
	training: A boolean flag indicating whether training behavior should be
	used (default: False).

	Returns:
	output: An output [batch, strided_height, strided_width,
	filters * value_expansion] tensor.
	"""
	_, height, width, channel = inputs.get_shape().as_list()

	# Reshape the inputs so that the attention is global 2D.
	x = tf.reshape(inputs, [-1, height * width, channel])

	# Implement 2D global attention as 1D axial-attention by flattening the 2D
	# inputs into 1D. We also disable the relative positional encodings in
	# axial attention, so that only content-based attention is used.
	x = self._global(x, training=training)

	# Use two global attention layers in one residual block. This option
	# ensures that global attention models have the same number of layers and
	# parameters as axial-attention models.
	if self._double_global_attention:
	x = self._global2(x, training=training)
	x = tf.reshape(x, [-1, height, width, self._total_value_depth])
	if self._strides > 1:
	x = x[:, ::self._strides, ::self._strides, :]

	return x