在 GPU 上训练模型时,Tensorflow 在 CPU 上加载权重

时间:2021-06-29 10:54:57

标签: python tensorflow keras nlp tf.keras

我在 Colab 中编写了一个 Bert 模型,并使用 GPU 对其进行了训练,并下载了权重以进行进一步推理。对于预测,我不需要 GPU,我在没有 GPU 的本地机器上进行测试。但是我在本地 PC 中加载时出现以下错误,而 Colab 上没有错误。我不知道如何进一步处理。

File "/home/akash/anaconda3/lib/python3.7/site-packages/tensorflow/python/saved_model/load.py", line 909,
 in load_internal str(err) + "\n If trying to load on a different device from the "
FileNotFoundError: Op type not registered 'CaseFoldUTF8' in binary running on akash. Make sure
the Op and Kernel are registered in the binary running in this process.
Note that if you are loading a saved graph which used ops from tf.contrib, accessing (e.g.) 
`tf.contrib.resampler` should be done before importing the graph, as contrib ops are lazily registered 
when the module is first accessed. 

我已经加载使用,

self.classifier_model = self.build_classifier_model()
self.classifier_model.load_weights(BERT_HEADING)

pip list | grep 'tensorflow' 的输出

tensorflow                         2.5.0
tensorflow-addons                  0.13.0
tensorflow-datasets                4.3.0
tensorflow-estimator               2.5.0
tensorflow-hub                     0.12.0
tensorflow-metadata                1.1.0
tensorflow-model-optimization      0.6.0
tensorflow-text                    2.5.0

我的模型:

bert_model_name = 'small_bert/bert_en_uncased_L-8_H-512_A-8'
tfhub_handle_encoder = 'https://tfhub.dev/tensorflow/small_bert/bert_en_uncased_L-8_H-512_A-8/1'
tfhub_handle_preprocess = 'https://tfhub.dev/tensorflow/bert_en_uncased_preprocess/3'
bert_preprocess_model = hub.KerasLayer(tfhub_handle_preprocess)
bert_model = hub.KerasLayer(tfhub_handle_encoder)

def build_classifier_model():
  text_input = tf.keras.layers.Input(shape=(), dtype=tf.string, name='text')
  preprocessing_layer = hub.KerasLayer(tfhub_handle_preprocess, name='preprocessing')
  encoder_inputs = preprocessing_layer(text_input)
  encoder = hub.KerasLayer(tfhub_handle_encoder, trainable=True, name='BERT_encoder')
  outputs = encoder(encoder_inputs)
  net = outputs['pooled_output']
  net = tf.keras.layers.Dropout(0.1)(net)
  net = tf.keras.layers.Dense(updated_data_frame['heading'].nunique(), activation='softmax', name='classifier')(net)
  return tf.keras.Model(text_input, net)

classifier_model = build_classifier_model()

epochs = 5
steps_per_epoch = tf.data.experimental.cardinality(train_ds).numpy()
num_train_steps = steps_per_epoch * epochs
num_warmup_steps = int(0.1*num_train_steps)

init_lr = 3e-5
optimizer = optimization.create_optimizer(init_lr=init_lr,
                                          num_train_steps=num_train_steps,
                                          num_warmup_steps=num_warmup_steps,
                                          optimizer_type='adamw')
classifier_model.compile(optimizer=optimizer,
                         loss=loss,
                         metrics=['CategoricalAccuracy'])
print(f'Training model with {tfhub_handle_encoder}')
history = classifier_model.fit(x=train_ds,
                               validation_data=val_ds,
                               epochs=5)

saved_model_path = 'resume_headings.h5'
classifier_model.save_weights(saved_model_path)

reloaded_model= build_classifier_model() # <-- This was working fine on Colab but giving an error (detailed desc above)
reloaded_model.load_weights(saved_model_path)

1 个答案:

答案 0 :(得分:0)

您遇到的错误很可能是因为您没有在本地安装 tensorflow-text(或者您使用 BERT 的环境没有安装 tensorflow-text)。

我进行此观察是因为我可以看到您同时使用 PIP 和 CONDA 作为包管理器,而且很容易陷入陷阱(不知道您使用的是什么环境/虚拟环境)。

例如,您可以有一个环境为 TF 2.3 和 pandas 1.0.0,另一个环境为 TF 2.5 和 pandas 1.2.0。现在当然,如果您使用第二个环境,则不能使用 pandas 版本 > 1.0.0 附带的功能,因为依赖项实际上是相对于该特定环境的。我希望你能更清楚。

另外,请确保您在本地拥有与您训练 BERT 模型的 TensorFlow 版本完全相同的版本。

(使用 !pip install tensorflow-text 安装)

您也可以在这里看到一个非常相似的错误:https://github.com/tensorflow/hub/issues/705

同时,如果您想确保您的模型已加载到 GPU 上,请使用以下代码段/逻辑:

with tf.device('/GPU:0'):
    model = load_model()