diff --git a/tests/test_main.py b/tests/test_main.py index ebbc430..1e3f60c 100644 --- a/tests/test_main.py +++ b/tests/test_main.py @@ -5,7 +5,7 @@ import tensorflow as tf from transformerx.layers import ( - MultiHeadAttention, PositionalEncoding, PositionWiseFFN, AddNorm, + MultiHeadAttention, AbsolutePositionalEncoding, PositionWiseFFN, AddNorm, TransformerEncoderBlock, TransformerEncoder, TransformerDecoderBlock, DotProductAttention, ) from transformerx.txplot import Plot @@ -25,7 +25,7 @@ def test_transpose_qkv(): depth, num_steps = 32, 50 -pos_encoding = PositionalEncoding(depth, 0) +pos_encoding = AbsolutePositionalEncoding(depth, 0) X = pos_encoding(tf.zeros((2, num_steps, depth)), training=False) P = pos_encoding.P[:, : X.shape[1], :] plotter = Plot() diff --git a/transformerx/layers/__init__.py b/transformerx/layers/__init__.py index e59726e..2533de8 100644 --- a/transformerx/layers/__init__.py +++ b/transformerx/layers/__init__.py @@ -1,7 +1,7 @@ from .addnorm import AddNorm from .dot_product_attention import DotProductAttention from .multihead_attention import MultiHeadAttention -from .positional_encoding import PositionalEncoding +from .positional_encoding import AbsolutePositionalEncoding from .positionwise_ffn import PositionWiseFFN from .transformer_decoder import TransformerDecoder from .transformer_decoder_block import TransformerDecoderBlock diff --git a/transformerx/layers/addnorm.py b/transformerx/layers/addnorm.py index ddce7e7..d24c56b 100644 --- a/transformerx/layers/addnorm.py +++ b/transformerx/layers/addnorm.py @@ -94,29 +94,29 @@ def __init__(self, norm_shape: Tuple[int], dropout_rate: float = 0): self.dropout = tf.keras.layers.Dropout(dropout_rate) self.ln = tf.keras.layers.LayerNormalization(norm_shape) - def call(self, X: tf.Tensor, Y: tf.Tensor, **kwargs): + def call(self, x: tf.Tensor, residual: tf.Tensor, **kwargs): """Call AddNorm layer. Parameters ---------- - X : + x : Input tensor - Y : - Input tensor 2 + residual : + Residual input tensor Returns ------- output : Added and normalized tensor """ - if not isinstance(X, tf.Tensor): + if not isinstance(x, tf.Tensor): raise TypeError( f"Expected a tensor for the " - f"argument 'X', but received: {X}" + f"argument 'x', but received: {x}" ) - if not isinstance(Y, tf.Tensor): + if not isinstance(residual, tf.Tensor): raise TypeError( f"Expected a tensor for the " - f"argument 'Y', but received: {Y}" + f"argument 'residual', but received: {residual}" ) - return self.ln(self.dropout(Y, **kwargs) + X) \ No newline at end of file + return self.ln(self.dropout(residual, **kwargs) + x) \ No newline at end of file diff --git a/transformerx/layers/multihead_attention.py b/transformerx/layers/multihead_attention.py index bc9ab60..bc9999c 100644 --- a/transformerx/layers/multihead_attention.py +++ b/transformerx/layers/multihead_attention.py @@ -103,7 +103,7 @@ def __init__( def split_heads(self, X: tf.Tensor) -> tf.Tensor: """Transpose tensors for parallel computation of attention heads. - First transposition produces a tensor of shape X: (batch_size, num_heads, no. of queries or key-value pairs, + First transposition produces a tensor of shape x: (batch_size, num_heads, no. of queries or key-value pairs, depth / num_heads). Next it is rearranged to a new order (batch_size * num_heads, no. of queries or key-value pairs, depth / num_heads) which is then passed to the last rearrangement and returned. @@ -114,19 +114,19 @@ def split_heads(self, X: tf.Tensor) -> tf.Tensor: The tensor to be transposed and prepared for the multi-head attention layer (i.e. queries, keys, and values) Returns ------- - X : tf.Tensor + x : tf.Tensor Transposed tensor of shape ((batch_size * num_heads, no. of queries or key-value pairs, depth / num_heads) """ - # X = tf.reshape(X, shape=(X.shape[0], X.shape[1], self.num_heads, -1)) + # x = tf.reshape(x, shape=(x.shape[0], x.shape[1], self.num_heads, -1)) X = rearrange(X, "b h (heads hidden) -> b h heads hidden", heads=self.num_heads) - # print("X reshaped: ", X.shape) - # X = tf.transpose(X, perm=(0, 2, 1, 3)) + # print("x reshaped: ", x.shape) + # x = tf.transpose(x, perm=(0, 2, 1, 3)) X = rearrange(X, "b d1 d2 d3 -> b d2 d1 d3") - # print("X transposed: ", X.shape) - # return tf.reshape(X, shape=(-1, X.shape[2], X.shape[3])) + # print("x transposed: ", x.shape) + # return tf.reshape(x, shape=(-1, x.shape[2], x.shape[3])) X = rearrange(X, "b d1 d2 d3 -> (b d1) d2 d3") - # print("X reshaped2: ", X.shape) + # print("x reshaped2: ", x.shape) return X def inverse_transpose_qkv(self, X): diff --git a/transformerx/layers/positional_encoding.py b/transformerx/layers/positional_encoding.py index 47e3298..44d5ad5 100644 --- a/transformerx/layers/positional_encoding.py +++ b/transformerx/layers/positional_encoding.py @@ -2,27 +2,70 @@ import tensorflow as tf -class PositionalEncoding(tf.keras.layers.Layer): - def __init__(self, num_hiddens, dropout, max_len=1000): +class AbsolutePositionalEncoding(tf.keras.layers.Layer): + """Absolute positional encoding object. + + Generate a sinusoid for each dimension of the positional encoding where wavelengths form a geometric progression + from :math:`2π` to :math:`(10000)2π`. + + Notes + ----- + Absolute Position Encodings are a type of position embeddings for [Transformer-based models] where positional + encodings are added to the input embeddings at the bottoms of the encoder and decoder stacks. The positional + encodings have the same dimension :math:`d_model` as the embeddings, so that the two can be summed. In the original + implementation, sine and cosine functions of different frequencies are used: + + .. math:: + PE(pos, 2i) = \sin(pos^{2i/d_{model}}) + + PE(pos, 2i+1) = \cos(pos^{2i/d_{model}}) + + where is the position and is the dimension. + + + Parameters + ---------- + depth : + Length of the positional encoding's hidden units. + dropout_rate : + Float between 0 and 1. Fraction of the input units to drop. + max_len : + Maximum length of the steps to calculate sinusoid + + Examples + -------- + >>> depth, num_steps = 32, 50 + >>> pos_encoding = AbsolutePositionalEncoding(depth, dropout_rate=0.2) + >>> x = tf.zeros((1, num_steps, depth)) + >>> print(x.shape) + (1, 50, 32) + >>> X = pos_encoding(x, training=False) + >>> P = pos_encoding.P[:, : X.shape[1], :] + >>> print(X.shape) + (1, 50, 32) + >>> print(P.shape) + (1, 50, 32) + """ + + def __init__(self, depth, dropout_rate=0, max_len=1000): super().__init__() - self.dropout = tf.keras.layers.Dropout(dropout) + self.dropout = tf.keras.layers.Dropout(dropout_rate) # Create a long enough P - self.P = np.zeros((1, max_len, num_hiddens)) - print("P.shape", self.P.shape) + self.P = np.zeros((1, max_len, depth)) X = np.arange(max_len, dtype=np.float32).reshape(-1, 1) / np.power( - 10000, np.arange(0, num_hiddens, 2, dtype=np.float32) / num_hiddens + 10000, np.arange(0, depth, 2, dtype=np.float32) / depth ) self.P[:, :, 0::2] = tf.sin( - X + X ) # x[low::stride] -> positions: 0, 2, 4, ... of all rows and columns self.P[:, :, 1::2] = tf.cos( - X + X ) # x[low::stride] -> positions: 1, 3, 5 , ... of all rows and columns def call(self, X, **kwargs): - # print("X.shape[1]: ", X.shape[1]) - # print("self.P[:, : X.shape[1], :]: ", self.P[:, : X.shape[1], :].shape) + # print("x.shape[1]: ", x.shape[1]) + # print("self.P[:, : x.shape[1], :]: ", self.P[:, : x.shape[1], :].shape) X = X + self.P[:, : X.shape[1], :] return self.dropout(X, **kwargs) diff --git a/transformerx/layers/transformer_decoder.py b/transformerx/layers/transformer_decoder.py index 27bcc3c..8a686d6 100644 --- a/transformerx/layers/transformer_decoder.py +++ b/transformerx/layers/transformer_decoder.py @@ -1,6 +1,6 @@ import tensorflow as tf -from transformerx.layers.positional_encoding import PositionalEncoding +from transformerx.layers.positional_encoding import AbsolutePositionalEncoding from transformerx.layers.transformer_decoder_block import TransformerDecoderBlock @@ -21,7 +21,7 @@ def __init__( self.depth = depth self.n_blocks = n_blocks self.embedding = tf.keras.layers.Embedding(vocab_size, depth) - self.pos_encoding = PositionalEncoding(depth, dropout) + self.pos_encoding = AbsolutePositionalEncoding(depth, dropout) self.blocks = [ TransformerDecoderBlock( depth, diff --git a/transformerx/layers/transformer_encoder.py b/transformerx/layers/transformer_encoder.py index a4bb9c6..60d491f 100644 --- a/transformerx/layers/transformer_encoder.py +++ b/transformerx/layers/transformer_encoder.py @@ -1,6 +1,6 @@ import tensorflow as tf -from transformerx.layers.positional_encoding import PositionalEncoding +from transformerx.layers.positional_encoding import AbsolutePositionalEncoding from transformerx.layers.transformer_encoder_block import TransformerEncoderBlock @@ -22,7 +22,7 @@ def __init__( self.depth = depth self.n_blocks = n_blocks self.embedding = tf.keras.layers.Embedding(vocab_size, depth) - self.pos_encoding = PositionalEncoding(depth, dropout) + self.pos_encoding = AbsolutePositionalEncoding(depth, dropout) self.blocks = [ TransformerEncoderBlock( depth, diff --git a/transformerx/utils.py b/transformerx/utils.py index 331f71a..393e869 100644 --- a/transformerx/utils.py +++ b/transformerx/utils.py @@ -6,7 +6,7 @@ def masked_softmax(X, valid_lens): """Perform softmax operation by masking elements on the last axis.""" - # X: 3D tensor, valid_lens: 1D or 2D tensor + # x: 3D tensor, valid_lens: 1D or 2D tensor def _sequence_mask(X, valid_len, value=0): maxlen = X.shape[1] mask = tf.range(start=0, limit=maxlen, dtype=tf.float32)[None, :] < tf.cast(