trax
trax copied to clipboard
google
Convert Transformer NMT model usin Trax2Keras error
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'>.
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)
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?
