Với Keras, bạn có thể xây dựng mô hình Transformer sử dụng 6 lớp Encoder và Decoder. Dưới đây là một thiết kế cơ bản cho mô hình Transformer bằng Keras. Để giảm độ phức tạp, mã ví dụ chỉ mô phỏng mô hình Sub-Transformer (Encoder hoặc Decoder), và không thực hiện gradient checkpointing hoặc tính toán hiệu quả hơn, điều mà một mô hình lớn thực sự có thể cần.

Mã giả đơn giản không bao gồm keras-nlp. Tôi khuyên bạn nên thử tự học và nâng cấp chúng để hiểu rõ hơn về cách hoạt động của mô hình Transformer.

import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers

class TransformerBlock(layers.Layer):
    def __init__(self, embed_dim, num_heads, ff_dim, rate=0.1):
        super(TransformerBlock, self).__init__()
        self.att = layers.MultiHeadAttention(num_heads=num_heads, key_dim=embed_dim)
        self.ffn = keras.Sequential(
            [layers.Dense(ff_dim, activation="relu"), layers.Dense(embed_dim),]
        )
        self.layernorm1 = layers.LayerNormalization(epsilon=1e-6)
        self.layernorm2 = layers.LayerNormalization(epsilon=1e-6)
        self.dropout1 = layers.Dropout(rate)
        self.dropout2 = layers.Dropout(rate)

    def call(self, inputs, training):
        attn_output = self.att(inputs, inputs)
        attn_output = self.dropout1(attn_output, training=training)
        out1 = self.layernorm1(inputs + attn_output)
        ffn_output = self.ffn(out1)
        ffn_output = self.dropout2(ffn_output, training=training)
        return self.layernorm2(out1 + ffn_output)

embed_dim = 32  # Embedding size for each token
num_heads = 2  # Number of attention heads
ff_dim = 32  # Hidden layer size in feed forward network inside transformer

inputs = layers.Input(shape=(None,))
embedding_layer = layers.Embedding(input_dim=10000, output_dim=embed_dim)
x = embedding_layer(inputs)
transformer_block = TransformerBlock(embed_dim, num_heads, ff_dim)
x = transformer_block(x)
outputs = layers.Dense(2)(x)

model = keras.Model(inputs=inputs, outputs=outputs)

Đó chỉ là một trang phục kiểu Transformers, điều đó không đảm bảo rằng chúng ta đã hiểu rõ mô hình Transformer. Để khám phá sâu hơn, chào mừng bạn đến với trituenhantao.io - nơi bạn có thể nắm bắt các khái niệm AI một cách intuitively. Thế giới của Transformer đang chờ đón bạn!