def _maybe_build(self, inputs): # Check input assumptions set before layer building, e.g. input rank. if not self.built: input_spec.assert_input_compatibility(self.input_spec, inputs, self.name) input_list = nest.flatten(inputs) input_shapes = None if all(hasattr(x, 'shape') for x in input_list): input_shapes = nest.map_structure(lambda x: x.shape, inputs) # Only call `build` if the user has manually overridden the build method. if not hasattr(self.build, '_is_default'): # Any setup work performed only once should happen in an `init_scope` # to avoid creating symbolic Tensors that will later pollute any eager # operations. with tf_utils.maybe_init_scope(self): self.build(input_shapes) # We must set self.built since user defined build functions are not # constrained to set self.built. self.built = True
def _build_from_signature(self, query, value, key=None): """Builds layers and variables. Once the method is called, self._built_from_signature will be set to True. Args: query: query tensor or TensorShape. value: value tensor or TensorShape. key: key tensor or TensorShape. """ self._built_from_signature = True if hasattr(query, "shape"): query_shape = tensor_shape.TensorShape(query.shape) else: query_shape = query if hasattr(value, "shape"): value_shape = tensor_shape.TensorShape(value.shape) else: value_shape = value if key is None: key_shape = value_shape elif hasattr(key, "shape"): key_shape = tensor_shape.TensorShape(key.shape) else: key_shape = key common_kwargs = dict(kernel_initializer=self._kernel_initializer, bias_initializer=self._bias_initializer, kernel_regularizer=self._kernel_regularizer, bias_regularizer=self._bias_regularizer, activity_regularizer=self._activity_regularizer, kernel_constraint=self._kernel_constraint, bias_constraint=self._bias_constraint) # Any setup work performed only once should happen in an `init_scope` # to avoid creating symbolic Tensors that will later pollute any eager # operations. with tf_utils.maybe_init_scope(self): free_dims = query_shape.rank - 1 einsum_equation, bias_axes, output_rank = _build_proj_equation( free_dims, bound_dims=1, output_dims=2) self._query_dense = einsum_dense.EinsumDense( einsum_equation, output_shape=_get_output_shape( output_rank - 1, [self._num_heads, self._key_dim]), bias_axes=bias_axes if self._use_bias else None, name="query", **common_kwargs) einsum_equation, bias_axes, output_rank = _build_proj_equation( key_shape.rank - 1, bound_dims=1, output_dims=2) self._key_dense = einsum_dense.EinsumDense( einsum_equation, output_shape=_get_output_shape( output_rank - 1, [self._num_heads, self._key_dim]), bias_axes=bias_axes if self._use_bias else None, name="key", **common_kwargs) einsum_equation, bias_axes, output_rank = _build_proj_equation( value_shape.rank - 1, bound_dims=1, output_dims=2) self._value_dense = einsum_dense.EinsumDense( einsum_equation, output_shape=_get_output_shape( output_rank - 1, [self._num_heads, self._value_dim]), bias_axes=bias_axes if self._use_bias else None, name="value", **common_kwargs) # Builds the attention computations for multi-head dot product attention. # These computations could be wrapped into the keras attention layer once # it support mult-head einsum computations. self._build_attention(output_rank) if self._output_shape: if not isinstance(self._output_shape, collections.abc.Sized): output_shape = [self._output_shape] else: output_shape = self._output_shape else: output_shape = [query_shape[-1]] einsum_equation, bias_axes, output_rank = _build_proj_equation( free_dims, bound_dims=2, output_dims=len(output_shape)) self._output_dense = einsum_dense.EinsumDense( einsum_equation, output_shape=_get_output_shape(output_rank - 1, output_shape), bias_axes=bias_axes if self._use_bias else None, name="attention_output", **common_kwargs)
def __init__(self, num_heads, key_dim, local_scope, num_timesteps, num_features, value_dim=None, dropout=0.0): super(MixedMultiHeadAttention, self).__init__() assert num_heads % 4 == 0 self._num_mixed_heads = int(num_heads / 4) self._key_dim = key_dim self._local_scope = local_scope self._num_timesteps = num_timesteps self._num_features = num_features self._value_dim = value_dim if value_dim else key_dim self._dropout = dropout self._query_dense = EinsumDense(equation="abe,cdef->acdbf", output_shape=(4, self._num_mixed_heads, self._num_timesteps, self._key_dim), bias_axes="f") self._key_dense = EinsumDense(equation="abe,cdef->acdfb", output_shape=(4, self._num_mixed_heads, self._key_dim, self._num_timesteps), bias_axes="f") self._value_dense = EinsumDense(equation="abe,cdef->acdbf", output_shape=(4, self._num_mixed_heads, self._num_timesteps, self._value_dim), bias_axes="f") self._softmax = Softmax() self._dropout_layer = Dropout(self._dropout) self._output_dense = EinsumDense(equation="abc,cd->abd", output_shape=(self._num_timesteps, self._num_features), bias_axes="d") with tf_utils.maybe_init_scope(self): g = np.zeros((self._num_timesteps, self._num_timesteps)) l = np.zeros((self._num_timesteps, self._num_timesteps)) f = np.zeros((self._num_timesteps, self._num_timesteps)) b = np.zeros((self._num_timesteps, self._num_timesteps)) for i in range(self._num_timesteps): for j in range(self._num_timesteps): if i - self._local_scope > j or j > i + self._local_scope: l[i, j] = np.NINF if i > j: f[i, j] = np.NINF if i < j: b[i, j] = np.NINF m = np.stack([g, l, f, b]) m = tf.convert_to_tensor(m, tf.float16) m = tf.expand_dims(m, 1) m = tf.expand_dims(m, 0) self._masks = m