Cannot assign a device for operation when using GPU with multiple embedding inputs in Tensorflow 2.3 with Keras

Question

I found some related issues for this problem with Tensorflow 1.X, but none with Tensorflow 2.X using Keras.

When using a single Embedding feature, everything works fine, but if I add multiple, I start getting colocation errors when running with a GPU. Everything works fine with just a CPU.

Any idea of a workaround? I simplified my code to the following minimal example:

import tensorflow as tf

def create_model():
  test_input = tf.keras.Input(shape=(None,), dtype='string', name='test')
  test2_input = tf.keras.Input(shape=(None,), dtype='string', name='test2')
  feature_layer_inputs = {}
  feature_layer_inputs['test'] = test_input
  feature_layer_inputs['test2'] = test2_input

  vocab_list = ['This', 'That', 'Thing']
  feature_col = tf.feature_column.categorical_column_with_vocabulary_list(
      key='test', vocabulary_list=vocab_list,
      num_oov_buckets=0)
  embed_col = tf.feature_column.embedding_column(
      categorical_column=feature_col, dimension=4, combiner='mean')
  first_embed_layer = tf.keras.layers.DenseFeatures(
      feature_columns=[embed_col], name="first_embed_layer")

  second_vocab_list = ['a', 'b', 'c']
  feature_col_two = tf.feature_column.categorical_column_with_vocabulary_list(
      key='test2', vocabulary_list=second_vocab_list,
      num_oov_buckets=0)
  embed_col_two = tf.feature_column.embedding_column(
      categorical_column=feature_col_two, dimension=4, combiner='mean')
  second_embed_layer = tf.keras.layers.DenseFeatures(
      feature_columns=[embed_col_two], name="second_embed_layer")
  
  x = first_embed_layer(feature_layer_inputs)
  y = second_embed_layer(feature_layer_inputs)
  x = tf.keras.layers.concatenate([x, y])
  
  hidden_layer = tf.keras.layers.Dense(units=35, use_bias=False,
      name="user-embeddings-layer")(x)

  model = tf.keras.Model(
    inputs=[v for v in feature_layer_inputs.values()],
    outputs=[hidden_layer]
  )

  model.compile(optimizer=tf.keras.optimizers.Adagrad(lr=.01),
                # loss=loss_func,
                loss="sparse_categorical_crossentropy",
                metrics=['accuracy'])
  return model

in_tensor = tf.constant(['This', 'That'])
other_tensor = tf.constant(['a', 'b'])

features = {
  'test': in_tensor,
  'test2': other_tensor,
}
y = tf.constant([1, 2])

model = create_model()
history = model.fit(x=features, y=y,
                    epochs=10, shuffle=False, 
                    batch_size=1,
                    verbose=1,
                    callbacks=[]

The full error is:

tensorflow.python.framework.errors_impl.InvalidArgumentError: Cannot assign a device for operation functional_1/first_embed_layer/test_embedding/test_embedding_weights/embedding_lookup_sparse/embedding_lookup: Could not satisfy explicit device specification '' because the node {{colocation_node functional_1/first_embed_layer/test_embedding/test_embedding_weights/embedding_lookup_sparse/embedding_lookup}} was colocated with a group of nodes that required incompatible device '/job:localhost/replica:0/task:0/device:GPU:0'. All available devices [/job:localhost/replica:0/task:0/device:CPU:0, /job:localhost/replica:0/task:0/device:XLA_CPU:0, /job:localhost/replica:0/task:0/device:XLA_GPU:0, /job:localhost/replica:0/task:0/device:GPU:0].

Answer 1

I found a workaround, which was to use the Sequential model API as opposed to the functional API, and to put both of my embedding columns into a single DenseFeatures layer.

This is less than ideal because now I have to search by variable dimension size to refer to the weights of one of the embedding columns (they don't seem to be located in the same location of the layer.weights[index] array reliably as you add more embedding columns). But, at least it runs.

I filed the following github issue as well: https://github.com/tensorflow/tensorflow/issues/42590

import tensorflow as tf

def create_model():
  # All models in this course are sequential.
  model = tf.keras.models.Sequential()

  vocab_list = ['This', 'That', 'Thing']
  feature_col = tf.feature_column.categorical_column_with_vocabulary_list(
      key='test', vocabulary_list=vocab_list,
      num_oov_buckets=0)
  embed_col = tf.feature_column.embedding_column(
      categorical_column=feature_col, dimension=4, combiner='mean')

  second_vocab_list = ['a', 'b', 'c']
  feature_col_two = tf.feature_column.categorical_column_with_vocabulary_list(
      key='test2', vocabulary_list=second_vocab_list,
      num_oov_buckets=0)
  embed_col_two = tf.feature_column.embedding_column(
      categorical_column=feature_col_two, dimension=8, combiner='mean')
  first_embed_layer = tf.keras.layers.DenseFeatures(
      feature_columns=[embed_col, embed_col_two], name="first_embed_layer")

  model.add(first_embed_layer)
  model.add(tf.keras.layers.Dense(units=35, use_bias=False,
      name="user-embeddings-layer"))
                           
  # Construct the layers into a model that TensorFlow can execute.  
  # Notice that the loss function for multi-class classification
  # is different than the loss function for binary classification.  
  model.compile(optimizer=tf.keras.optimizers.Adam(lr=.01),
                loss="sparse_categorical_crossentropy",
                metrics=['accuracy'])
  
  return model

in_tensor = tf.constant(['This', 'That'])
other_tensor = tf.constant(['a', 'b'])

features = {
  'test': in_tensor,
  'test2': other_tensor,
}
y = tf.constant([1, 2])

model = create_model()

history = model.fit(x=features, y=y,
                    epochs=10, shuffle=False, 
                    batch_size=1,
                    verbose=1,
                    callbacks=[]