Convert Transformer NMT model usin Trax2Keras error

[OmarAlsaqa]

Description

I tried many times to convert diffrent models to keras but I always face a problem. I found on Gitter many answers like this one can any one help me.

Environment information

Running on Colab

$ pip freeze | grep trax
1.3.7

For bugs: reproduction and error logs

trax.fastmath.set_backend("tensorflow-numpy")

# define the output directory
output_dir = '/content/drive/MyDrive/Colab Notebooks/Transformer_FR_pretrained_336_Keras'

# define the training loop
training_loop = training.Loop(model,
                              train_task,
                              eval_tasks=[eval_task],
                              output_dir=output_dir)

# Run 1 step to adapt the backend
training_loop.run(1)

Step 1: Total number of trainable weights: 80370196 Step 1: Ran 1 train steps in 95.61 secs Step 1: train CrossEntropyLossWithLogSoftmax | 1.28223217 Step 1: eval CrossEntropyLossWithLogSoftmax | 1.65389895 Step 1: eval WeightedCategoryAccuracy | 0.67606580

ba_size = 1
# To convert the model to Keras, simply run:
keras_layer = trax.AsKeras(model, batch_size=ba_size)
# This will be a trax.trax2keras.AsKeras object
print(keras_layer)

<trax.trax2keras.AsKeras object at 0x7fb99cb9aed0>

import tensorflow as tf

# Create a full Keras model using the layer you loaded from trax.
inputs = tf.keras.Input(shape=(ba_size,), dtype='int32')
# Use default Keras syntax to link an input to the layer
hidden = keras_layer(inputs)
# Get the outputs from the trax-loaded layer
outputs = hidden
# Finally, wrap everything with a Keras Model
keras_model = tf.keras.Model(inputs=inputs, outputs=outputs)

WARNING:tensorflow:Entity <bound method TracebackException.format of <traceback.TracebackException object at 0x7fb99c8ef810>> appears to be a generator function. It will not be converted by AutoGraph.

WARNING:tensorflow:Entity <bound method TracebackException.format of <traceback.TracebackException object at 0x7fb99c8ef810>> appears to be a generator function. It will not be converted by AutoGraph.

WARNING: Entity <bound method TracebackException.format of <traceback.TracebackException object at 0x7fb99c8ef810>> appears to be a generator function. It will not be converted by AutoGraph.

---------------------------------------------------------------------------

StagingError                              Traceback (most recent call last)

<ipython-input-37-205c4cd80595> in <module>()
      4 inputs = tf.keras.Input(shape=(ba_size,), dtype='int32')
      5 # Use default Keras syntax to link an input to the layer
----> 6 hidden = keras_layer(inputs)
      7 # Get the outputs from the trax-loaded layer
      8 outputs = hidden

4 frames

/usr/local/lib/python3.7/dist-packages/tensorflow/python/autograph/impl/api.py in wrapper(*args, **kwargs)
    668       except Exception as e:  # pylint:disable=broad-except
    669         if hasattr(e, 'ag_error_metadata'):
--> 670           raise e.ag_error_metadata.to_exception(e)
    671         else:
    672           raise

StagingError: in user code:

    /usr/local/lib/python3.7/dist-packages/trax/trax2keras.py:184 call  *
        outputs, new_state = self._trax_layer.pure_fn(inputs, weights=weights,
    /usr/local/lib/python3.7/dist-packages/trax/layers/base.py:548 pure_fn  *
        raise LayerError(name, 'pure_fn',

    LayerError: Exception passing through layer Serial (in pure_fn):
      layer created in file [...]/trax/models/transformer.py, line 384
      layer input shapes: ShapeDtype{shape:(1, 1), dtype:int32}

      File [...]/autograph/operators/control_flow.py, line 1165, in if_stmt
        _py_if_stmt(cond, body, orelse)

      File [...]/autograph/operators/control_flow.py, line 1218, in _py_if_stmt
        return body() if cond else orelse()

      File [...]//tmp/tmp859vrwqw.py, line 44, in if_body_1
        outputs = ag__.converted_call(ag__.ld(self).forward, (ag__.ld(x),), None, fscope)

      File [...]/autograph/impl/api.py, line 461, in converted_call
        result = converted_f(*effective_args)

      File [...]//tmp/tmp1brwo2o6.py, line 11, in tf__forward
        ag__.converted_call(ag__.ld(self)._validate_forward_inputs, (ag__.ld(xs),), None, fscope)

      File [...]/autograph/impl/api.py, line 461, in converted_call
        result = converted_f(*effective_args)

      File [...]//tmp/tmps9xfnep4.py, line 20, in tf___validate_forward_inputs
        ag__.if_stmt(ag__.and_((lambda : ag__.not_(ag__.converted_call(ag__.ld(isinstance), (ag__.ld(xs), (ag__.ld(tuple), ag__.ld(list))), None, fscope))), (lambda : (ag__.ld(self)._n_in != 1))), if_body, else_body, get_state, set_state, (), 0)

      File [...]/autograph/operators/control_flow.py, line 1165, in if_stmt
        _py_if_stmt(cond, body, orelse)

      File [...]/autograph/operators/control_flow.py, line 1218, in _py_if_stmt
        return body() if cond else orelse()

      File [...]//tmp/tmps9xfnep4.py, line 16, in if_body
        raise ag__.converted_call(ag__.ld(TypeError), (f'Serial.forward input must be a tuple or list; instead got {ag__.converted_call(ag__.ld(type), (ag__.ld(xs),), None, fscope)}.',), None, fscope)

    TypeError: Serial.forward input must be a tuple or list; instead got <class 'tensorflow.python.ops.numpy_ops.np_arrays.ndarray'>.

[elx42]

Hi,

This is how I manage to transfer my nmt model from trax to keras. Maybe this helps you out on your issue. I am still struggling how to best host this model then in an efficient way in tf_serving. If you find out sth, keep me posted. :)

All the best, Matthias

model = trax.models.Transformer(
        input_vocab_size=32768,
        d_model=512, d_ff=1024,
        n_heads=8, n_encoder_layers=6, n_decoder_layers=6,
        max_len=1024,
        mode='eval' )

model.init_from_file( model_folder+'/model.pkl.gz'  )

keras_layer = trax.AsKeras(model, batch_size=1)
inputs  = tf.keras.Input(shape=(1024,), dtype='int32')
hidden = keras_layer((inputs, inputs))

outputs = hidden
keras_model = tf.keras.Model(inputs=inputs, outputs=outputs)

[AndreasAntoniou]

Hi,

This is how I manage to transfer my nmt model from trax to keras. Maybe this helps you out on your issue. I am still struggling how to best host this model then in an efficient way in tf_serving. If you find out sth, keep me posted. :)

All the best, Matthias

model = trax.models.Transformer(
        input_vocab_size=32768,
        d_model=512, d_ff=1024,
        n_heads=8, n_encoder_layers=6, n_decoder_layers=6,
        max_len=1024,
        mode='eval' )

model.init_from_file( model_folder+'/model.pkl.gz'  )

keras_layer = trax.AsKeras(model, batch_size=1)
inputs  = tf.keras.Input(shape=(1024,), dtype='int32')
hidden = keras_layer((inputs, inputs))

outputs = hidden
keras_model = tf.keras.Model(inputs=inputs, outputs=outputs)

Do you have any updates on this?