lecrapaud 0.4.0__py3-none-any.whl

lecrapaud/speed_tests/test-gpu-bilstm.ipynb

@@ -0,0 +1,261 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Speed Test\n",
+    "\n",
+    "### **Conclusion**\n",
+    "- **~1.8X faster on GPU compared to CPU on Apple Silicon. On Intel processors, 2X faster with CPU (but metal is not supported on intel processor)** \n",
+    "- **CuDNN** (NVIDIA’s deep learning library) significantly improves performance for LSTM-based models compared to vanilla RNNs. (See LSTM vs RNN lstm in keras documentation)\n",
+    "\n",
+    "---\n",
+    "\n",
+    "### **Model Details**\n",
+    "- **Model Type:** BiLSTM  \n",
+    "- **Total Parameters:** 2,658,945 (~10.14 MB)  \n",
+    "- **Training Data:** 25,000 sequences  \n",
+    "- **Features:** 1 (number of occurrences in the corpus)\n",
+    "\n",
+    "---\n",
+    "\n",
+    "### **Rules of Thumb**\n",
+    "1. **Sample Size:**  \n",
+    "   `nb_samples > 10 × nb_features × nb_classes`  \n",
+    "   (with `nb_classes = 5` for regression).\n",
+    "\n",
+    "2. **Parameter-Sample Ratio:**  \n",
+    "   `nb_parameters < nb_samples / 10`  \n",
+    "   (or even `nb_samples / 50` for deep learning).\n",
+    "\n",
+    "---\n",
+    "\n",
+    "### **GPU Performance**\n",
+    "\n",
+    "#### **tf.keras**  \n",
+    "- Run 1: **146s** (~179ms/step) — loss: `0.4118`, accuracy: `0.8075`, val_loss: `0.3565`, val_accuracy: `0.8464`  \n",
+    "- Run 2 (Apple Silicon): **29s** (~32ms/step) — loss: `0.4167`, accuracy: `0.8018`, val_loss: `0.3612`, val_accuracy: `0.8474`\n",
+    "\n",
+    "#### **Keras**  \n",
+    "- Run 1: **119s** (~146ms/step) — loss: `0.4102`, accuracy: `0.8142`, val_loss: `0.3496`, val_accuracy: `0.8467`  \n",
+    "- Run 2 (Apple Silicon): **25s** (~30ms/step) — loss: `0.4167`, accuracy: `0.8018`, val_loss: `0.3611`, val_accuracy: `0.8474`\n",
+    "\n",
+    "---\n",
+    "\n",
+    "### **CPU Performance**\n",
+    "\n",
+    "#### **tf.keras**  \n",
+    "- Run 1: **73s** (~89ms/step) — loss: `0.4184`, accuracy: `0.8078`, val_loss: `0.3444`, val_accuracy: `0.8507`  \n",
+    "- Run 2 (Apple Silicon): **47s** (~59ms/step) — loss: `0.5113`, accuracy: `0.7238`, val_loss: `0.3405`, val_accuracy: `0.8538`\n",
+    "\n",
+    "#### **Keras**  \n",
+    "- Run 1: **81s** (~97ms/step) — loss: `0.4027`, accuracy: `0.8147`, val_loss: `0.3395`, val_accuracy: `0.8508`  \n",
+    "- Run 2 (Apple Silicon): **47s** (~58ms/step) — loss: `0.5113`, accuracy: `0.7238`, val_loss: `0.3405`, val_accuracy: `0.853\n",
+    "\n",
+    "---\n",
+    "\n",
+    "### **Key Observations**\n",
+    "1. **CPU vs GPU:**  \n",
+    "   - GPU training is approximately **2X faster** than CPU for this BiLSTM model.  \n",
+    "\n",
+    "2. **Accuracy and Loss:**  \n",
+    "   - Both CPU and GPU achieve similar final accuracy and loss values, demonstrating functional equivalence.  \n",
+    "\n",
+    "3. **Recommendation:**  \n",
+    "   - Use CPU for smaller datasets or models like BiLSTM.  \n",
+    "   - For larger datasets or deeper architectures, GPU with CuDNN can provide significant speedups.\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import tensorflow as tf\n",
+    "\n",
+    "devices = tf.config.list_physical_devices()\n",
+    "print(\"\\nDevices: \", devices)\n",
+    "\n",
+    "gpus = tf.config.list_physical_devices(\"GPU\")\n",
+    "if gpus:\n",
+    "    details = tf.config.experimental.get_device_details(gpus[0])\n",
+    "    print(\"GPU details: \", details)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# get data\n",
+    "import tensorflow as tf\n",
+    "from tensorflow import keras\n",
+    "import numpy as np\n",
+    "\n",
+    "np.random.seed(42)  # for reproducibility\n",
+    "\n",
+    "max_features = 20000\n",
+    "maxlen = 100  # cut texts after this number of words (among top max_features most common words)\n",
+    "batch_size = 32\n",
+    "\n",
+    "(x_train, y_train), (x_test, y_test) = keras.datasets.imdb.load_data(\n",
+    "    num_words=max_features\n",
+    ")\n",
+    "\n",
+    "x_train = tf.keras.preprocessing.sequence.pad_sequences(x_train, maxlen=maxlen)\n",
+    "x_test = tf.keras.preprocessing.sequence.pad_sequences(x_test, maxlen=maxlen)\n",
+    "print(x_train.shape, x_train[:2], y_train.shape, y_train[:2])\n",
+    "print(\n",
+    "    f\"{x_train.shape[0]} train samples and {y_test.shape[0]} tests samples, for a total of {x_train.shape[0] + x_test.shape[0]} samples.\"\n",
+    ")\n",
+    "y_train = np.array(y_train)\n",
+    "y_test = np.array(y_test)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from tensorflow import keras\n",
+    "\n",
+    "print(f\"Using tf.keras version {keras.__version__}\")\n",
+    "\n",
+    "# Define the input layer\n",
+    "inputs = keras.Input(shape=(maxlen,))  # maxlen is the input length for each sequence\n",
+    "\n",
+    "# Embedding layer\n",
+    "x = keras.layers.Embedding(max_features, 128)(\n",
+    "    inputs\n",
+    ")  # max_features is the vocabulary size\n",
+    "\n",
+    "# Bidirectional LSTM layer\n",
+    "x = keras.layers.Bidirectional(keras.layers.LSTM(64))(x)\n",
+    "\n",
+    "# Dropout layer for regularization\n",
+    "x = keras.layers.Dropout(0.5)(x)\n",
+    "\n",
+    "# Output layer (for binary classification)\n",
+    "outputs = keras.layers.Dense(1, activation=\"sigmoid\")(x)\n",
+    "\n",
+    "# Define the model with inputs and outputs\n",
+    "model = keras.Model(inputs=inputs, outputs=outputs, name=\"BiLSTM\")\n",
+    "\n",
+    "# Print model summary\n",
+    "model.summary()\n",
+    "\n",
+    "# Compile the model\n",
+    "model.compile(optimizer=\"adam\", loss=\"binary_crossentropy\", metrics=[\"accuracy\"])\n",
+    "\n",
+    "# Train the model\n",
+    "print(\"Train...\")\n",
+    "model.fit(\n",
+    "    x_train, y_train, batch_size=batch_size, epochs=1, validation_data=(x_test, y_test)\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import keras\n",
+    "\n",
+    "print(f\"Using keras version {keras.__version__}\")\n",
+    "\n",
+    "model = keras.models.Sequential(name=\"BiLSTM\")\n",
+    "model.add(keras.layers.Embedding(max_features, 128, input_length=maxlen))\n",
+    "model.add(keras.layers.Bidirectional(keras.layers.LSTM(64)))\n",
+    "model.add(keras.layers.Dropout(0.5))\n",
+    "model.add(keras.layers.Dense(1, activation=\"sigmoid\"))\n",
+    "model.summary()\n",
+    "# try using different optimizers and different optimizer configs\n",
+    "model.compile(\"adam\", \"binary_crossentropy\", metrics=[\"accuracy\"])\n",
+    "\n",
+    "print(\"Train...\")\n",
+    "model.fit(\n",
+    "    x_train, y_train, batch_size=batch_size, epochs=1, validation_data=[x_test, y_test]\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# # Comparing GPU performance of LSTM using CuDNN (keras.layers.LSTM) and not using CuDNN (keras.layers.LSTMCell)\n",
+    "\n",
+    "# from tensorflow import keras\n",
+    "# from keras.models import Sequential\n",
+    "# from keras.layers import Dense, Dropout, Embedding, LSTM, Input, Bidirectional\n",
+    "\n",
+    "# keras.utils.set_random_seed(42)\n",
+    "\n",
+    "\n",
+    "# def build_model(allow_cudnn_kernel=True):\n",
+    "#     # CuDNN is only available at the layer level, and not at the cell level.\n",
+    "#     # This means `LSTM(units)` will use the CuDNN kernel,\n",
+    "#     # while RNN(LSTMCell(units)) will run on non-CuDNN kernel.\n",
+    "#     if allow_cudnn_kernel:\n",
+    "#         # The LSTM layer with default options uses CuDNN.\n",
+    "#         model = Sequential(name=\"BiLSTM with CuDNN\")\n",
+    "#         lstm_layer = keras.layers.LSTM(64)\n",
+    "#     else:\n",
+    "#         # Wrapping a LSTMCell in a RNN layer will not use CuDNN.\n",
+    "#         model = Sequential(name=\"BiLSTM without CuDNN\")\n",
+    "#         lstm_layer = keras.layers.RNN(keras.layers.LSTMCell(64))\n",
+    "#     model.add(Embedding(max_features, 128, input_length=maxlen))\n",
+    "#     model.add(Bidirectional(lstm_layer))\n",
+    "#     model.add(Dropout(0.5))\n",
+    "#     model.add(Dense(1, activation=\"sigmoid\"))\n",
+    "#     model.summary()\n",
+    "#     return model\n",
+    "\n",
+    "\n",
+    "# model = build_model(allow_cudnn_kernel=True)\n",
+    "# model.compile(\"adam\", \"binary_crossentropy\", metrics=[\"accuracy\"])\n",
+    "\n",
+    "# print(\"Train model using CuDNN kernel...\")\n",
+    "# model.fit(\n",
+    "#     x_train, y_train, batch_size=batch_size, epochs=1, validation_data=[x_test, y_test]\n",
+    "# )\n",
+    "\n",
+    "\n",
+    "# model_noncudnn = build_model(allow_cudnn_kernel=False)\n",
+    "# model_noncudnn.compile(\"adam\", \"binary_crossentropy\", metrics=[\"accuracy\"])\n",
+    "\n",
+    "# print(\"Train model not using CuDNN kernel...\")\n",
+    "# model_noncudnn.fit(\n",
+    "#     x_train, y_train, batch_size=batch_size, epochs=1, validation_data=[x_test, y_test]\n",
+    "# )"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": ".venv",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.12.8"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

lecrapaud/speed_tests/test-gpu-resnet.ipynb

@@ -0,0 +1,166 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Speed Test\n",
+    "\n",
+    "### **Conclusion**\n",
+    "- **12x faster with GPU compared to CPU. (53s vs 680s) on Apple Silicon (test jan 24)**\n",
+    "\n",
+    "---\n",
+    "\n",
+    "### **Model Details**\n",
+    "- **Model Type:** Very Deep CNN (ResNet, 50 layers)\n",
+    "- **Total Parameters:** 23,792,612 (~90.76 MB)\n",
+    "- **Training Data:** 50,000 samples\n",
+    "- **Input Features:** `32 × 32 × 3 ≈ 3,000`\n",
+    "\n",
+    "---\n",
+    "\n",
+    "### **Rules of Thumb**\n",
+    "1. **Sample Size:**  \n",
+    "   `nb_samples > 10 × nb_features × nb_classes`  \n",
+    "   (with `nb_classes = 5` for regression).\n",
+    "\n",
+    "2. **Parameter-Sample Ratio:**  \n",
+    "   `nb_parameters < nb_samples / 10`  \n",
+    "   (or even `nb_samples / 50` for deep learning).\n",
+    "\n",
+    "---\n",
+    "\n",
+    "### **GPU Performance**\n",
+    "\n",
+    "#### **tf.keras**  \n",
+    "- Run 1: **221s** (~256ms/step) — loss: `4.6651`, accuracy: `0.0858`\n",
+    "- Run 2 (with Apple Silicon M2 GPU): **47s** (~53ms/step) - loss: `4.8064` - accuracy: `0.0620`\n",
+    "\n",
+    "#### **Keras**  \n",
+    "- Run 1: **225s** (~263ms/step) — loss: `4.7491`, accuracy: `0.0738`  \n",
+    "- Run 2 (with Apple Silicon M2 GPU): **47s** (~53ms/step) - loss: `4.6743` - accuracy: `0.0761`\n",
+    "\n",
+    "---\n",
+    "\n",
+    "### **CPU Performance**\n",
+    "\n",
+    "#### **tf.keras**  \n",
+    "- Run 1: **1440s** (~1841ms/step) — accuracy: `0.0346`, loss: `5.2095`  \n",
+    "- Run 2 (with Apple Silicon M2 CPU): **541s** (~679ms/step) — accuracy: `0.0534`, loss: `4.9699`  \n",
+    "\n",
+    "#### **Keras**  \n",
+    "- Run 1: **1440s** (~1841ms/step) — accuracy: `0.0210`, loss: `5.6996`  \n",
+    "- Run 2 (with Apple Silicon M2 CPU): **535s** (~671ms/step) — accuracy: `0.0534`, loss: `5.0862`\n",
+    "\n",
+    "\n",
+    "---\n",
+    "\n",
+    "### **Key Takeaways**\n",
+    "- GPUs offer a **6x speed improvement** compared to CPUs for deep CNN training.\n",
+    "- CPU training becomes impractically slow for large models and datasets.\n",
+    "- GPU is highly recommended for deep learning tasks, especially for very deep architectures like ResNet.\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import tensorflow as tf\n",
+    "\n",
+    "devices = tf.config.list_physical_devices()\n",
+    "print(\"\\nDevices: \", devices)\n",
+    "\n",
+    "gpus = tf.config.list_physical_devices(\"GPU\")\n",
+    "if gpus:\n",
+    "    details = tf.config.experimental.get_device_details(gpus[0])\n",
+    "    print(\"GPU details: \", details)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# get data\n",
+    "from tensorflow import keras\n",
+    "\n",
+    "cifar = keras.datasets.cifar100\n",
+    "(x_train, y_train), (x_test, y_test) = cifar.load_data()\n",
+    "print(x_train.shape, x_train[:2], y_train.shape, y_train[:2])\n",
+    "print(\n",
+    "    f\"{x_train.shape[0]} train samples and {y_test.shape[0]} tests samples, for a total of {x_train.shape[0] + x_test.shape[0]} samples.\"\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import tensorflow as tf\n",
+    "\n",
+    "model = tf.keras.applications.ResNet50(\n",
+    "    include_top=True,\n",
+    "    weights=None,\n",
+    "    input_shape=(32, 32, 3),\n",
+    "    classes=100,\n",
+    ")\n",
+    "model.summary()\n",
+    "\n",
+    "loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=False)\n",
+    "\n",
+    "model.compile(optimizer=\"adam\", loss=loss_fn, metrics=[\"accuracy\"])\n",
+    "\n",
+    "model.fit(x_train, y_train, epochs=1, batch_size=64)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import keras\n",
+    "\n",
+    "model = keras.applications.ResNet50(\n",
+    "    include_top=True,\n",
+    "    weights=None,\n",
+    "    input_shape=(32, 32, 3),\n",
+    "    classes=100,\n",
+    ")\n",
+    "model.summary()\n",
+    "\n",
+    "loss_fn = keras.losses.SparseCategoricalCrossentropy(from_logits=False)\n",
+    "\n",
+    "model.compile(optimizer=\"adam\", loss=loss_fn, metrics=[\"accuracy\"])\n",
+    "\n",
+    "model.fit(x_train, y_train, epochs=1, batch_size=64)"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": ".venv",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.12.8"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

lecrapaud/speed_tests/test-gpu-transformers.ipynb

@@ -0,0 +1,254 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Speed Test\n",
+    "\n",
+    "### **Conclusion**\n",
+    "- **~same on CPU compared to GPU** for this model and setup.  \n",
+    "\n",
+    "---\n",
+    "\n",
+    "### **Model Details**\n",
+    "- **Model Type:** Transformer (encoder only)  \n",
+    "- **Total Parameters:** 4,934,529 (18.82 MB)\n",
+    "- **Training Data:** 25,000 sequences  \n",
+    "- **Features:** 1 (number of occurrences in the corpus)\n",
+    "\n",
+    "---\n",
+    "\n",
+    "### **Rules of Thumb**\n",
+    "1. **Sample Size:**  \n",
+    "   `nb_samples > 10 × nb_features × nb_classes`  \n",
+    "   (with `nb_classes = 5` for regression).\n",
+    "\n",
+    "2. **Parameter-Sample Ratio:**  \n",
+    "   `nb_parameters < nb_samples / 10`  \n",
+    "   (or even `nb_samples / 50` for deep learning).\n",
+    "\n",
+    "---\n",
+    "\n",
+    "### **GPU Performance**\n",
+    "\n",
+    "#### **tf.keras**  \n",
+    "- Run 1 (Apple Silicon) : 412s 525ms/step - loss: 7.6209 - accuracy: 0.4999 - val_loss: 7.6246 - val_accuracy: 0.5000\n",
+    "\n",
+    "#### **Keras**  \n",
+    "- Run 1 (Apple Silicon) : 419s 534ms/step - loss: 7.6191 - accuracy: 0.5002 - val_loss: 7.6246 - val_accuracy: 0.5000\n",
+    "\n",
+    "\n",
+    "---\n",
+    "\n",
+    "### **CPU Performance**\n",
+    "\n",
+    "#### **tf.keras**  \n",
+    "- Run 1 (Apple Silicon) : 439s 560ms/step - accuracy: 0.4992 - loss: 8.0630 - val_accuracy: 0.5000 - val_loss: 8.0590\n",
+    "\n",
+    "#### **Keras**  \n",
+    "- Run 1 (Apple Silicon) : 417s 531ms/step - accuracy: 0.5006 - loss: 7.9368 - val_accuracy: 0.5000 - val_loss: 7.9712\n",
+    "\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import tensorflow as tf\n",
+    "\n",
+    "devices = tf.config.list_physical_devices()\n",
+    "print(\"\\nDevices: \", devices)\n",
+    "\n",
+    "gpus = tf.config.list_physical_devices(\"GPU\")\n",
+    "if gpus:\n",
+    "    details = tf.config.experimental.get_device_details(gpus[0])\n",
+    "    print(\"GPU details: \", details)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# get data\n",
+    "import tensorflow as tf\n",
+    "from tensorflow import keras\n",
+    "import numpy as np\n",
+    "\n",
+    "np.random.seed(42)  # for reproducibility\n",
+    "\n",
+    "max_features = 20000\n",
+    "maxlen = 100  # cut texts after this number of words (among top max_features most common words)\n",
+    "batch_size = 32\n",
+    "\n",
+    "(x_train, y_train), (x_test, y_test) = keras.datasets.imdb.load_data(\n",
+    "    num_words=max_features\n",
+    ")\n",
+    "\n",
+    "x_train = tf.keras.preprocessing.sequence.pad_sequences(x_train, maxlen=maxlen)\n",
+    "x_test = tf.keras.preprocessing.sequence.pad_sequences(x_test, maxlen=maxlen)\n",
+    "print(x_train.shape, x_train[:2], y_train.shape, y_train[:2])\n",
+    "print(\n",
+    "    f\"{x_train.shape[0]} train samples and {y_test.shape[0]} tests samples, for a total of {x_train.shape[0] + x_test.shape[0]} samples.\"\n",
+    ")\n",
+    "y_train = np.array(y_train)\n",
+    "y_test = np.array(y_test)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from tensorflow import keras\n",
+    "\n",
+    "print(f\"Using tf.keras version {keras.__version__}\")\n",
+    "\n",
+    "ff_dim = 512\n",
+    "head_size = 128\n",
+    "num_heads = 16\n",
+    "num_layers = 2\n",
+    "\n",
+    "\n",
+    "def transformer_encoder(inputs, num_layers, head_size, num_heads, ff_dim, dropout=0):\n",
+    "    for _ in range(num_layers):\n",
+    "        # Attention and Normalization\n",
+    "        x = keras.layers.LayerNormalization(epsilon=1e-6)(inputs)\n",
+    "        x = keras.layers.MultiHeadAttention(\n",
+    "            key_dim=head_size, num_heads=num_heads, dropout=dropout\n",
+    "        )(x, x)\n",
+    "        x = keras.layers.Add()([x, inputs])\n",
+    "\n",
+    "        # Feed Forward Part\n",
+    "        y = keras.layers.LayerNormalization(epsilon=1e-6)(x)\n",
+    "        y = keras.layers.Dense(ff_dim, activation=\"relu\")(y)\n",
+    "        y = keras.layers.Dropout(dropout)(y)\n",
+    "        y = keras.layers.Dense(inputs.shape[-1])(y)\n",
+    "        inputs = keras.layers.Add()([y, x])\n",
+    "\n",
+    "    return inputs\n",
+    "\n",
+    "\n",
+    "# Define the input layer\n",
+    "inputs = keras.layers.Input(\n",
+    "    shape=(maxlen,)\n",
+    ")  # maxlen is the input length for each sequence\n",
+    "\n",
+    "# Embedding layer\n",
+    "x = keras.layers.Embedding(max_features, 128)(\n",
+    "    inputs\n",
+    ")  # max_features is the vocabulary size\n",
+    "x = transformer_encoder(x, num_layers, head_size, num_heads, ff_dim)\n",
+    "x = keras.layers.GlobalAveragePooling1D()(x)\n",
+    "x = keras.layers.LayerNormalization(epsilon=1e-6)(x)\n",
+    "# Define output layer based on target type\n",
+    "outputs = keras.layers.Dense(1, activation=\"linear\")(x)\n",
+    "\n",
+    "# Build the model\n",
+    "model = keras.Model(inputs=inputs, outputs=outputs, name=\"transformer\")\n",
+    "\n",
+    "model.summary()\n",
+    "# try using different optimizers and different optimizer configs\n",
+    "model.compile(\"adam\", \"binary_crossentropy\", metrics=[\"accuracy\"])\n",
+    "\n",
+    "print(\"Train...\")\n",
+    "batch_size = 32\n",
+    "epoch = 100\n",
+    "model.fit(\n",
+    "    x_train, y_train, batch_size=batch_size, epochs=1, validation_data=[x_test, y_test]\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import keras\n",
+    "\n",
+    "print(f\"Using keras version {keras.__version__}\")\n",
+    "\n",
+    "ff_dim = 512\n",
+    "head_size = 128\n",
+    "num_heads = 16\n",
+    "num_layers = 2\n",
+    "\n",
+    "\n",
+    "def transformer_encoder(inputs, num_layers, head_size, num_heads, ff_dim, dropout=0):\n",
+    "    for _ in range(num_layers):\n",
+    "        # Attention and Normalization\n",
+    "        x = keras.layers.LayerNormalization(epsilon=1e-6)(inputs)\n",
+    "        x = keras.layers.MultiHeadAttention(\n",
+    "            key_dim=head_size, num_heads=num_heads, dropout=dropout\n",
+    "        )(x, x)\n",
+    "        x = keras.layers.Add()([x, inputs])\n",
+    "\n",
+    "        # Feed Forward Part\n",
+    "        y = keras.layers.LayerNormalization(epsilon=1e-6)(x)\n",
+    "        y = keras.layers.Dense(ff_dim, activation=\"relu\")(y)\n",
+    "        y = keras.layers.Dropout(dropout)(y)\n",
+    "        y = keras.layers.Dense(inputs.shape[-1])(y)\n",
+    "        inputs = keras.layers.Add()([y, x])\n",
+    "\n",
+    "    return inputs\n",
+    "\n",
+    "\n",
+    "# Define the input layer\n",
+    "inputs = keras.layers.Input(\n",
+    "    shape=(maxlen,)\n",
+    ")  # maxlen is the input length for each sequence\n",
+    "\n",
+    "# Embedding layer\n",
+    "x = keras.layers.Embedding(max_features, 128)(\n",
+    "    inputs\n",
+    ")  # max_features is the vocabulary size\n",
+    "x = transformer_encoder(x, num_layers, head_size, num_heads, ff_dim)\n",
+    "x = keras.layers.GlobalAveragePooling1D()(x)\n",
+    "x = keras.layers.LayerNormalization(epsilon=1e-6)(x)\n",
+    "# Define output layer based on target type\n",
+    "outputs = keras.layers.Dense(1, activation=\"linear\")(x)\n",
+    "\n",
+    "# Build the model\n",
+    "model = keras.Model(inputs=inputs, outputs=outputs, name=\"transformer\")\n",
+    "\n",
+    "model.summary()\n",
+    "# try using different optimizers and different optimizer configs\n",
+    "model.compile(\"adam\", \"binary_crossentropy\", metrics=[\"accuracy\"])\n",
+    "\n",
+    "print(\"Train...\")\n",
+    "batch_size = 32\n",
+    "epoch = 100\n",
+    "model.fit(\n",
+    "    x_train, y_train, batch_size=batch_size, epochs=1, validation_data=[x_test, y_test]\n",
+    ")"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": ".venv",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.12.8"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}