在Google Colab本地运行时用GPU调用Universal Sentence Encoder报错求助

问题描述 #

学习TensorFlow机器学习课程时，使用TensorFlow Hub的Universal Sentence Encoder v4进行文本分类。该模型在Google Colab GPU环境、本地无GPU运行时均正常，但配置Colab连接本地RTX 3060 GPU后，出现JIT编译失败的错误。本地环境基于Anaconda，已通过conda安装tensorflow_gpu、cudatoolkit和cudnn，不清楚错误原因及调试方向。

代码片段 #

import tensorflow_hub as hub
tf_hub_embedding = hub.KerasLayer('https://tfhub.dev/google/universal-sentence-encoder/4',trainable=False,name='USE')

rand_sent = random.choice(train_sents)
print(f'Random sent: {rand_sent}\n')
print(f'Embedded sent: {tf_hub_embedding([rand_sent])[0][:30]}\n')
print(f'Embed length: {len(tf_hub_embedding([rand_sent])[0])}')

错误信息 #

Random sent: Data of a Japanese study of patients with unresectable sacral chordoma showed comparable high control rates after hypofractionated carbon ion therapy only .

---------------------------------------------------------------------------
UnknownError                              Traceback (most recent call last)
Input In [55], in <cell line: 3>()
      1 rand_sent = random.choice(train_sents)
      2 print(f'Random sent: {rand_sent}\n')
----> 3 print(f'Embedded sent: {tf_hub_embedding([rand_sent])[0][:30]}\n')
      4 print(f'Embed length: {len(tf_hub_embedding([rand_sent])[0])}')

File ~\anaconda3\lib\site-packages\keras\utils\traceback_utils.py:67, in filter_traceback.<locals>.error_handler(*args, **kwargs)
     65 except Exception as e:  # pylint: disable=broad-except
     66   filtered_tb = _process_traceback_frames(e.__traceback__)
---> 67   raise e.with_traceback(filtered_tb) from None
     68 finally:
     69   del filtered_tb

File ~\anaconda3\lib\site-packages\tensorflow_hub\keras_layer.py:229, in KerasLayer.call(self, inputs, training)
    223 # ...but we may also have to pass a Python boolean for `training`, which
    224 # is the logical "and" of this layer's trainability and what the surrounding
    225 # model is doing (analogous to tf.keras.layers.BatchNormalization in TF2).
    226 # For the latter, we have to look in two places: the `training` argument,
    227 # or else Keras' global `learning_phase`, which might actually be a tensor.
    228 if not self._has_training_argument:
---> 229   result = f()
    230 else:
    231   if self.trainable:

UnknownError: Exception encountered when calling layer "USE" (type KerasLayer).

Graph execution error:

JIT compilation failed.
     [[{{node EncoderDNN/EmbeddingLookup/EmbeddingLookupUnique/embedding_lookup/mod}}]] [Op:__inference_restored_function_body_36706]

Call arguments received by layer "USE" (type KerasLayer):
  • inputs=["'Data of a Japanese study of patients with unresectable sacral chordoma showed comparable high control rates after hypofractionated carbon ion therapy only .'"]
  • training=None

调试方向与解决方法 #

• 校验版本兼容性：RTX3060属于安培架构，要求CUDA版本≥11.0，对应TensorFlow版本需匹配（如TF2.5+对应CUDA11.2）。用conda list查看tensorflow_gpu、cudatoolkit、cudnn的版本，确保三者版本适配，若不匹配则重新安装对应版本。
• 禁用JIT编译：错误核心是JIT编译失败，在代码开头添加import tensorflow as tf; tf.config.optimizer.set_jit(False)，强制关闭即时编译，绕过该问题。
• 确认GPU识别状态：运行print(tf.config.list_physical_devices('GPU'))，检查输出是否包含RTX3060。若未识别，更新GPU驱动至最新版本，或确认conda环境中tensorflow_gpu安装正确（而非CPU版TensorFlow）。
• 清理模型缓存：删除TensorFlow Hub的本地缓存目录（默认路径为~/.cache/tensorflow/hub），重新加载模型，避免缓存损坏导致的异常。
• 对齐Colab与本地环境版本：确保Colab使用的TensorFlow版本与本地Anaconda环境一致，连接本地运行时选择正确的环境，避免跨版本兼容性问题。

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