BigDL-Nano TensorFlow SparseEmbedding and SparseAdam#

In this guide we demonstrates how to use SparseEmbedding and SparseAdam to obtain stroger performance with sparse gradient.

Step 0: Prepare Environment#

We recommend using conda to prepare the environment. Please refer to the install guide for more details.

conda create py37 python==3.7.10 setuptools==58.0.4
conda activate py37
# nightly bulit version
pip install --pre --upgrade bigdl-nano[tensorflow]
# set env variables for your conda environment
source bigdl-nano-init
pip install tensorflow-datasets

Step 1: Import BigDL-Nano#

The optimizations in BigDL-Nano are delivered through BigDL-Nano’s Model and Sequential classes. For most cases, you can just replace your tf.keras.Model to and tf.keras.Sequential to to benefits from BigDL-Nano.

from import Model, Sequential

Step 2: Load the data#

We demonstrate with imdb_reviews, a large movie review dataset.

import tensorflow_datasets as tfds
(raw_train_ds, raw_val_ds, raw_test_ds), info = tfds.load(
    split=['train[:80%]', 'train[80%:]', 'test'],

Step 3: Parepre the Data#

In particular, we remove

import tensorflow as tf
from tensorflow.keras.layers import TextVectorization
import string
import re

def custom_standardization(input_data):
    lowercase = tf.strings.lower(input_data)
    stripped_html = tf.strings.regex_replace(lowercase, "<br />", " ")
    return tf.strings.regex_replace(
        stripped_html, f"[{re.escape(string.punctuation)}]", ""

max_features = 20000
embedding_dim = 128
sequence_length = 500

vectorize_layer = TextVectorization(

# Let's make a text-only dataset (no labels):
text_ds = x, y: x)
# Let's call `adapt`:

def vectorize_text(text, label):
    text = tf.expand_dims(text, -1)
    return vectorize_layer(text), label

# Vectorize the data.
train_ds =
val_ds =
test_ds =

# Do async prefetching / buffering of the data for best performance on GPU.
train_ds = train_ds.cache().prefetch(buffer_size=10)
val_ds = val_ds.cache().prefetch(buffer_size=10)
test_ds = test_ds.cache().prefetch(buffer_size=10)

Step 4: Build Model# is a slightly modified version of tf.keras.Embedding layer, this embedding layer only applies regularizer to the output of the embedding layer, so that the gradient to embeddings is sparse. is a variant of the Adam optimizer that handles sparse updates more efficiently. Here we create two models, one using normal Embedding layer and Adam optimizer, the other using SparseEmbedding and SparseAdam.

from tensorflow.keras import layers
from import Embedding
from import SparseAdam

from tensorflow.keras import layers
from import Embedding
from import SparseAdam

inputs = tf.keras.Input(shape=(None,), dtype="int64")

# Embedding layer can only be used as the first layer in a model,
# you need to provide the argument inputShape (a Single Shape, does not include the batch dimension).
x = Embedding(max_features, embedding_dim)(inputs)
x = layers.Dropout(0.5)(x)

# Conv1D + global max pooling
x = layers.Conv1D(128, 7, padding="valid", activation="relu", strides=3)(x)
x = layers.Conv1D(128, 7, padding="valid", activation="relu", strides=3)(x)
x = layers.GlobalMaxPooling1D()(x)

# We add a vanilla hidden layer:
x = layers.Dense(128, activation="relu")(x)
x = layers.Dropout(0.5)(x)

# We project onto a single unit output layer, and squash it with a sigmoid:
predictions = layers.Dense(1, activation="sigmoid", name="predictions")(x)

model = Model(inputs, predictions)

# Compile the model with binary crossentropy loss and an SparseAdam optimizer.
model.compile(loss="binary_crossentropy", optimizer=SparseAdam(), metrics=["accuracy"])

Step 5: Training#

# Fit the model using the train and val datasets., validation_data=val_ds, epochs=3)


You can find the detailed result of training from here