tensorflow: Tensorflow profiler crashes when using string categorical layer

Tensorflow profiler crashes when using string categorical layer, not allowing to profile models with those layers.

System information

• OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Ubuntu

• TensorFlow installed from (source or binary): binary

• TensorFlow version (use command below)😦’v1.12.1-35610-gd8c49c2fde’, ‘2.4.0-dev20200701’) tf-estimator-nightly 2.4.0.dev2020063001 tf-nightly 2.4.0.dev20200701

• Python version: 3.7

• CUDA/cuDNN version:

• GPU model and memory: ±----------------------------------------------------------------------------+ | NVIDIA-SMI 440.33.01 Driver Version: 440.33.01 CUDA Version: 10.2 | |-------------------------------±---------------------±---------------------+ | GPU Name Persistence-M| Bus-Id Disp.A | Volatile Uncorr. ECC | | Fan Temp Perf Pwr:Usage/Cap| Memory-Usage | GPU-Util Compute M. | |===============================+======================+======================| | 0 Tesla V100-SXM2… On | 00000000:00:1B.0 Off | 0 | | N/A 38C P0 50W / 300W | 15450MiB / 16160MiB | 0% Default | ±------------------------------±---------------------±---------------------+ | 1 Tesla V100-SXM2… On | 00000000:00:1C.0 Off | 0 | | N/A 37C P0 54W / 300W | 15450MiB / 16160MiB | 0% Default | ±------------------------------±---------------------±---------------------+ | 2 Tesla V100-SXM2… On | 00000000:00:1D.0 Off | 0 | | N/A 36C P0 55W / 300W | 15522MiB / 16160MiB | 0% Default | ±------------------------------±---------------------±---------------------+ | 3 Tesla V100-SXM2… On | 00000000:00:1E.0 Off | 0 | | N/A 37C P0 55W / 300W | 15522MiB / 16160MiB | 0% Default | ±------------------------------±---------------------±---------------------+

== cuda libs =================================================== /usr/local/cuda-9.2/doc/man/man7/libcudart.so.7 /usr/local/cuda-9.2/doc/man/man7/libcudart.7 /usr/local/cuda-9.2/lib64/libcudart.so.9.2.148 /usr/local/cuda-9.2/lib64/libcudart_static.a /usr/local/cuda-9.0/doc/man/man7/libcudart.so.7 /usr/local/cuda-9.0/doc/man/man7/libcudart.7 /usr/local/cuda-9.0/lib64/libcudart.so.9.0.176 /usr/local/cuda-9.0/lib64/libcudart_static.a /usr/local/cuda-10.1/targets/x86_64-linux/lib/libcudart.so.10.1.243 /usr/local/cuda-10.1/targets/x86_64-linux/lib/libcudart_static.a /usr/local/cuda-10.1/doc/man/man7/libcudart.so.7 /usr/local/cuda-10.1/doc/man/man7/libcudart.7 /usr/local/cuda-10.2/targets/x86_64-linux/lib/libcudart.so.10.2.89 /usr/local/cuda-10.2/targets/x86_64-linux/lib/libcudart_static.a /usr/local/cuda-10.2/doc/man/man7/libcudart.so.7 /usr/local/cuda-10.2/doc/man/man7/libcudart.7 /usr/local/cuda-10.0/doc/man/man7/libcudart.so.7 /usr/local/cuda-10.0/doc/man/man7/libcudart.7 /usr/local/cuda-10.0/lib64/libcudart.so.10.0.130 /usr/local/cuda-10.0/lib64/libcudart_static.a /usr/local/cuda-8.0/doc/man/man7/libcudart.so.7 /usr/local/cuda-8.0/doc/man/man7/libcudart.7 /usr/local/cuda-8.0/lib64/libcudart.so.8.0.61 /usr/local/cuda-8.0/lib64/libcudart_static.a

== tensorflow installed from info ==================

== python version ============================================== (major, minor, micro, releaselevel, serial) (3, 6, 10, ‘final’, 0)

Describe the current behavior Profiler crashes, doesn’t allow to profile a model

Describe the expected behavior I expect to be able to profile models with string categorical encoding

Standalone code to reproduce the issue

import tensorflow as tf
import numpy as np
import pandas as pd
from tensorflow import keras
from tensorflow.keras import layers

from datetime import datetime
from packaging import version
import os

from tensorflow.keras.layers.experimental.preprocessing import Normalization
from tensorflow.keras.layers.experimental.preprocessing import CategoryEncoding
from tensorflow.keras.layers.experimental.preprocessing import StringLookup

#tf.debugging.set_log_device_placement(True)

file_url = "https://storage.googleapis.com/applied-dl/heart.csv"
dataframe = pd.read_csv(file_url)

val_dataframe = dataframe.sample(frac=0.2, random_state=1337)
train_dataframe = dataframe.drop(val_dataframe.index)

print(
    "Using %d samples for training and %d for validation"
    % (len(train_dataframe), len(val_dataframe))
)

def dataframe_to_dataset(dataframe):
    dataframe = dataframe.copy()
    labels = dataframe.pop("target")
    ds = tf.data.Dataset.from_tensor_slices((dict(dataframe), labels))
    ds = ds.shuffle(buffer_size=len(dataframe))
    return ds


train_ds = dataframe_to_dataset(train_dataframe)
val_ds = dataframe_to_dataset(val_dataframe)

for x, y in train_ds.take(1):
    print("Input:", x)
    print("Target:", y)


def encode_numerical_feature(feature, name, dataset):
    normalizer = Normalization()

    feature_ds = dataset.map(lambda x, y: x[name])
    feature_ds = feature_ds.map(lambda x: tf.expand_dims(x, -1))

    normalizer.adapt(feature_ds)

    encoded_feature = normalizer(feature)
    return encoded_feature


def encode_string_categorical_feature(feature, name, dataset):
    index = StringLookup()

    feature_ds = dataset.map(lambda x, y: x[name])
    feature_ds = feature_ds.map(lambda x: tf.expand_dims(x, -1))

    index.adapt(feature_ds)

    encoded_feature = index(feature)

    encoder = CategoryEncoding(output_mode="binary")

    feature_ds = feature_ds.map(index)

    encoder.adapt(feature_ds)

    encoded_feature = encoder(encoded_feature)
    return encoded_feature


def encode_integer_categorical_feature(feature, name, dataset):
    encoder = CategoryEncoding(output_mode="binary")

    feature_ds = dataset.map(lambda x, y: x[name])
    feature_ds = feature_ds.map(lambda x: tf.expand_dims(x, -1))

    encoder.adapt(feature_ds)

    encoded_feature = encoder(feature)
    return encoded_feature

def build_models():
    sex = keras.Input(shape=(1,), name="sex", dtype="int64")
    cp = keras.Input(shape=(1,), name="cp", dtype="int64")
    fbs = keras.Input(shape=(1,), name="fbs", dtype="int64")
    restecg = keras.Input(shape=(1,), name="restecg", dtype="int64")
    exang = keras.Input(shape=(1,), name="exang", dtype="int64")
    ca = keras.Input(shape=(1,), name="ca", dtype="int64")

    thal = keras.Input(shape=(1,), name="thal", dtype="string")

    age = keras.Input(shape=(1,), name="age")
    trestbps = keras.Input(shape=(1,), name="trestbps")
    chol = keras.Input(shape=(1,), name="chol")
    thalach = keras.Input(shape=(1,), name="thalach")
    oldpeak = keras.Input(shape=(1,), name="oldpeak")
    slope = keras.Input(shape=(1,), name="slope")

    all_inputs = [
        sex,
        cp,
        fbs,
        restecg,
        exang,
        ca,
        thal,
        age,
        trestbps,
        chol,
        thalach,
        oldpeak,
        slope,
    ]

    numerical_inputs = [
        age,
        trestbps,
        chol,
        thalach,
        oldpeak,
        slope,
    ]


    sex_encoded = encode_integer_categorical_feature(sex, "sex", train_ds)
    cp_encoded = encode_integer_categorical_feature(cp, "cp", train_ds)
    fbs_encoded = encode_integer_categorical_feature(fbs, "fbs", train_ds)
    restecg_encoded = encode_integer_categorical_feature(restecg, "restecg", train_ds)
    exang_encoded = encode_integer_categorical_feature(exang, "exang", train_ds)
    ca_encoded = encode_integer_categorical_feature(ca, "ca", train_ds)

    thal_encoded = encode_string_categorical_feature(thal, "thal", train_ds)

    age_encoded = encode_numerical_feature(age, "age", train_ds)
    trestbps_encoded = encode_numerical_feature(trestbps, "trestbps", train_ds)
    chol_encoded = encode_numerical_feature(chol, "chol", train_ds)
    thalach_encoded = encode_numerical_feature(thalach, "thalach", train_ds)
    oldpeak_encoded = encode_numerical_feature(oldpeak, "oldpeak", train_ds)
    slope_encoded = encode_numerical_feature(slope, "slope", train_ds)

    all_features = layers.concatenate(
        [
            sex_encoded,
            cp_encoded,
            fbs_encoded,
            restecg_encoded,
            exang_encoded,
            slope_encoded,
            ca_encoded,
            thal_encoded,
            age_encoded,
            trestbps_encoded,
            chol_encoded,
            thalach_encoded,
            oldpeak_encoded,
        ]
    )

    numerical_features = layers.concatenate(
        [
            age_encoded,
            trestbps_encoded,
            chol_encoded,
            thalach_encoded,
            oldpeak_encoded,
            slope_encoded
        ]
    )


    x = layers.Dense(32, activation="relu")(all_features)
    x = layers.Dropout(0.5)(x)
    output = layers.Dense(1, activation="sigmoid")(x)
    model_cat_string = keras.Model(all_inputs, output)
    model_cat_string.compile("adam", "binary_crossentropy", metrics=["accuracy"])
    
    x = layers.Dense(32, activation="relu")(numerical_features)
    x = layers.Dropout(0.5)(x)
    output = layers.Dense(1, activation="sigmoid")(x)
    model_only_num = keras.Model(numerical_inputs, output)
    model_only_num.compile("adam", "binary_crossentropy", metrics=["accuracy"])

    return (model_only_num, model_cat_string)



model_only_numerical, model_with_categorical_string_layers = build_models()


logs = "logs/" + datetime.now().strftime("%Y%m%d-%H%M%S")
tboard_callback = tf.keras.callbacks.TensorBoard(log_dir=logs, histogram_freq=1, profile_batch='1,20')


#model with only numerical features, profiler works fine
model_only_numerical.summary()
model_only_numerical.fit(train_ds.batch(32), 
                     epochs=50, 
                     validation_data=val_ds.batch(32),
                    callbacks=[tboard_callback]
                    )

#model with categorical string encoding layers, profiler crashes!
model_with_categorical_string_layers.summary()
model_with_categorical_string_layers.fit(train_ds.batch(32), 
                     epochs=50, 
                     validation_data=val_ds.batch(32),
                    callbacks=[tboard_callback]
                    )