CODEWITHSUNDEEP
pythontensorflowkeras
Jules
Jules

Reputation: 445

Error loading pretrained BERT on TPU using Keras

I am working with TensorFlow and Keras on a TPU and I need to load the pretrained BERT model and convert it to a TPU model. I first load the model, and it works perfectly.

def _load_bert(self):

    logging.info('Loading BERT from %s', self.bert_path)
    config_path = os.path.join(self.bert_path, 'bert_config.json')
    checkpoint_path = os.path.join(self.bert_path, 'bert_model.ckpt')

    bert = keras_bert.load_trained_model_from_checkpoint(
        config_path,
        checkpoint_path
    )

    self._bert = bert
    self._seq_length = int(self._bert.input[0].shape[1])

And then I am trying to convert it to a TPU model doing the following:

self._bert = tensorflow.contrib.tpu.keras_to_tpu_model(
    self._bert,
    strategy=tensorflow.contrib.tpu.TPUDistributionStrategy(
        tensorflow.contrib.cluster_resolver.TPUClusterResolver(
            self.tpu, zone=self.tpu_zone, project=self.gcp_project
        )
    )
)

And then for some reason I got the following error:

ValueError: ('Expected `model` argument to be a `Model` instance, got ', <keras.engine.training.Model object at xxxxxxxxxx >)

It works when I don't convert it to a TPU model. Does anyone have a clue?

Upvotes: 1

Views: 226

Answers (1)

rigo
rigo

Reputation: 325

try compiling your keras.engine.training.Model object with a strategy before running it or training it. 

resolver = tf.contrib.cluster_resolver.TPUClusterResolver(self.tpu, 
                                                           zone=self.tpu_zone, 
                                                           project=self.gcp_project)
tf.contrib.distribute.initialize_tpu_system(resolver)
strategy = tf.contrib.distribute.TPUStrategy(resolver)
with strategy.scope():
  model = self._bert
  model.compile(optimizer=tf.keras.optimizers.Adagrad(learning_rate=0.1), 
                 loss='mean_squared_error', metrics=['acc','mse'])

Upvotes: 3

Related Questions