迁移TensorFlow1.x脚本遇hub.text_embedding_column TF2兼容问题求解决

Hey there, I’ve dealt with this exact frustration when migrating TF1.x training scripts to TF2, so I know exactly how to help you out!

为什么原来的代码报错？ #

As you saw from the help() output, hub.text_embedding_column wasn’t adapted for TF2 in versions like TF2.1 and TF Hub 0.7.0—it’s built for the old TF1 Estimator API and doesn’t play nice with TF2’s eager execution or Keras-focused workflow.

替代方案：用TF Hub的Keras层（推荐） #

The cleanest way to use the Universal Sentence Encoder (USE) in TF2 is to wrap it as a Keras layer with hub.KerasLayer—this integrates seamlessly with modern TF2 workflows. Here’s how to do it step by step:

1. Load the USE model as a Keras layer

   import tensorflow as tf
   import tensorflow_hub as hub
   
   # Load USE as a trainable or non-trainable Keras layer
   use_embedding_layer = hub.KerasLayer(
       "https://tfhub.dev/google/universal-sentence-encoder/4",
       input_shape=[],  # Accepts scalar string inputs
       dtype=tf.string,
       trainable=False  # Set to True if you want to fine-tune the model
   )
   

2. Build your Keras model
   You can directly chain this embedding layer with your prediction head (classification/regression) like so:

   # Define input layer matching your text column name
   text_input = tf.keras.Input(shape=[], dtype=tf.string, name="test_col")
   
   # Generate embeddings from text
   embeddings = use_embedding_layer(text_input)
   
   # Add your prediction layers (example for binary classification)
   output = tf.keras.layers.Dense(1, activation="sigmoid")(embeddings)
   
   # Assemble the full model
   model = tf.keras.Model(inputs=text_input, outputs=output)
   
   # Compile the model
   model.compile(optimizer="adam", loss="binary_crossentropy", metrics=["accuracy"])
   

3. Train with your dataset
   If you’re using pandas data or tf.data.Dataset, feeding the model is straightforward:

   # Example with pandas DataFrame
   import pandas as pd
   
   sample_data = pd.DataFrame({
       "test_col": ["This is a sample sentence", "Another text snippet for training"],
       "label": [0, 1]
   })
   
   # Train the model
   model.fit(x=sample_data["test_col"], y=sample_data["label"], epochs=5)
   

如果你坚持要用特征列（比如和其他特征结合） #

If you need to keep using feature columns (e.g., combining text embeddings with numerical/categorical features), you can wrap the USE layer into a custom numeric feature column:

# Define a helper function to generate embeddings
def generate_embedding(input_tensor):
    return use_embedding_layer(input_tensor)

# Create a numeric feature column for the embeddings (USE outputs 512-dimensional vectors)
embedding_feature_column = tf.feature_column.numeric_column(
    key="test_col",
    shape=(512,),
    dtype=tf.float32,
    normalizer_fn=generate_embedding
)

# Now you can combine this with other feature columns and use it in a model
# For example, with a Keras model using feature columns:
feature_layer = tf.keras.layers.DenseFeatures([embedding_feature_column])
input_layer = tf.keras.Input(shape={}, dtype=tf.string, name="test_col")
features = feature_layer(input_layer)
output = tf.keras.layers.Dense(1, activation="sigmoid")(features)
feature_model = tf.keras.Model(inputs=input_layer, outputs=output)

This approach lets you mix text embeddings with other feature types while staying in the TF2 ecosystem.

内容的提问来源于stack exchange，提问作者Omar Cotugno