#!/usr/bin/env python
# ******************************************************************************
# Copyright 2022 Brainchip Holdings Ltd.
#
# 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.
# ******************************************************************************
__all__ = ["Attention", "string_to_softmax", "QuantizedAttention"]
import keras
import tensorflow as tf
from .layers_base import (register_quantize_target, register_aligned_inputs, apply_buffer_bitwidth,
QuantizedLayer)
from .reshaping import QuantizedReshape, QuantizedPermute
from .shiftmax import shiftmax, QuantizedShiftmax
from .recorders import TensorRecorder
from .quantizers import OutputQuantizer
from ..tensors import FixedPoint, round_log2, pow2
[docs]
def string_to_softmax(s):
"""
Convert a string to a softmax function.
Available options are 'softmax' for standard softmax, 'shiftmax' for
shiftmax.
Args:
s (str): string to convert.
Returns:
A softmax function.
"""
if s == "softmax":
return tf.nn.softmax
if s in ("softmax2", "shiftmax"):
return shiftmax
raise NotImplementedError("softmax should be in ['softmax', 'shiftmax']"
f" but received {s}.")
[docs]
@keras.saving.register_keras_serializable()
class Attention(keras.layers.Layer):
"""Dot-product attention layer with configurable softmax.
Inputs are a tuple of tensors:
- a query tensor of shape [batch, tokens, hidden],
- a key tensor of shape [batch, tokens, hidden],
- a value tensor of shape [batch, tokens, hidden].
The calculation follows the steps:
1. Split query, key, value per attention heads
q, k, v : [batch, tokens, hidden] -> [batch, token, num_heads, dim]
2. Calculate cross-token scores as a query-key dot product:
scores = tf.matmul(query, key, transpose_b=True)
scores : [batch, num_heads, token, token]
3. Rescale score by dividing by the squared-root of dim.
4. Use scores to calculate a mask
mask = softmax(scores)
5. Combine mask with value
output = tf.matmul(mask, value)
output: [batch, num_heads, token, dim]
6. Merge heads to get back to 2D
output: [batch, num_heads, token, dim] -> [batch, token, hidden]
Args:
num_heads (int): the number of attention heads
softmax (str, optional): 'softmax' or 'shiftmax'. Defaults to 'softmax'
"""
def __init__(self, num_heads, softmax="softmax", **kwargs):
super().__init__(**kwargs)
self.num_heads = num_heads
self.softmax = softmax
self.softmax_op = string_to_softmax(self.softmax)
def build(self, input_shape):
super().build(input_shape)
assert len(input_shape) == 3
self.hidden_size = input_shape[0][-1]
if self.hidden_size % self.num_heads != 0:
raise ValueError(
f"Embedding dimension = {self.hidden_size} should be divisible"
f" by number of heads = {self.num_heads}"
)
self.dim = self.hidden_size // self.num_heads
# Attention replace score / scale by a shift.
# For that, we need to calculate the shift_scale
scale = tf.math.sqrt(tf.cast(self.dim, dtype=tf.float32))
self.scale_shift = round_log2(scale)
def separate_heads(self, x):
x = keras.layers.Reshape((-1, self.num_heads, self.dim))(x)
return keras.layers.Permute((2, 1, 3))(x)
def call(self, inputs):
# Separate 2D embeddings per head to obtain 3D inputs
query = self.separate_heads(inputs[0])
key = self.separate_heads(inputs[1])
value = self.separate_heads(inputs[2])
# Dot product query and key for each head and pairs of tokens
score = tf.matmul(query, key, transpose_b=True)
# Rescale the corresponding score, dividing it by 2**scale_shift
scaled_score = score * pow2(-self.scale_shift)
# Apply the configurable softmax operation
mask = self.softmax_op(scaled_score, axis=-1)
# Combine each score with value to obtain new embeddings per tokens
output = tf.matmul(mask, value)
# Join heads to get back to 2D embeddings per token
output = keras.layers.Permute((2, 1, 3))(output)
output = keras.layers.Reshape((-1, self.hidden_size))(output)
return output, mask
def get_config(self):
config = super().get_config()
softmax = self.softmax
if self.softmax == 'softmax2':
# softmax2 is the legacy name, use shiftmax now
softmax = 'shiftmax'
config.update(
{
"num_heads": self.num_heads,
"softmax": softmax
}
)
return config
[docs]
@register_quantize_target(Attention)
@register_aligned_inputs
@keras.saving.register_keras_serializable()
class QuantizedAttention(QuantizedLayer, Attention):
"""An Attention layer that operates on quantized inputs
Args:
num_heads (int): the number of attention heads
quant_config (dict, optional): the serialized quantization configuration. Defaults to None.
softmax (str, optional): 'softmax' or 'shiftmax'. Defaults to 'shiftmax'
"""
arg_constraints = {'softmax': 'shiftmax'}
def __init__(self, num_heads, quant_config=None, softmax='shiftmax', **kwargs):
super().__init__(num_heads=num_heads, softmax=softmax, quant_config=quant_config, **kwargs)
out_quant_cfg = self.quant_config.get("output_quantizer", False)
if out_quant_cfg:
self.out_quantizer = OutputQuantizer(
name="output_quantizer", **out_quant_cfg)
else:
self.out_quantizer = None
# Override softmax operation
softmax_quant_conf = self.quant_config.get("softmax", None)
self.softmax_op = QuantizedShiftmax(
quant_config=softmax_quant_conf, name="quantized_softmax")
self.buffer_bitwidth = apply_buffer_bitwidth(self.quant_config, signed=True)
# Add objects that will store the shift values.
self.values_shift = TensorRecorder(self.name + "/value_shift")
def separate_heads(self, x):
x = QuantizedReshape((-1, self.num_heads, self.dim))(x)
return QuantizedPermute((2, 1, 3))(x)
def call(self, inputs):
if any(not isinstance(x, FixedPoint) for x in inputs):
# If any of the inputs is not a FixedPoint, raise an error
raise ValueError("QuantizedAttention only accepts FixedPoint inputs")
# Separate 2D embeddings per head to obtain 3D inputs
query = self.separate_heads(inputs[0])
key = self.separate_heads(inputs[1])
# Expand the values to a higher bitwidth to avoid saturation and align them
value, vshift = inputs[2].expand(self.buffer_bitwidth)
self.values_shift(vshift)
value = self.separate_heads(value)
# Promote query to avoid saturation
query = query.promote(self.buffer_bitwidth)
# Dot product query and key for each head and pairs of tokens
score = tf.matmul(query, key, transpose_b=True)
# Rescale the corresponding score, dividing it by 2**scale_shift
scaled_score = score >> self.scale_shift
# Apply the configurable softmax operation
mask = self.softmax_op(scaled_score)
# Promote mask to make sure we don't overflow
mask = mask.promote(self.buffer_bitwidth)
# Combine each score with value to obtain new embeddings per tokens
output = tf.matmul(mask, value)
# Join heads to get back to 2D embeddings per token
output = QuantizedPermute((2, 1, 3))(output)
output = QuantizedReshape((-1, self.hidden_size))(output)
# Refine output bitwidth precision if needed
if self.out_quantizer is not None:
output = self.out_quantizer(output)
return output, mask