Explore Objax to Tensorflow conversion about objax HOT 5 CLOSED

wrn_width = 2            # Width of WideResNet
wrn_depth = 28           # Depth of WideResNet
batch_size = 4

# Model
model = WideResNet(nin=3, nclass=10, depth=wrn_depth, width=wrn_width)
model_vars = model.vars()

# Prediction operation
predict_op = lambda x: objax.functional.softmax(model(x, training=False))
predict_op = objax.Jit(predict_op, model_vars)

# Run prediction on random batch
x = objax.random.normal((batch_size, 3, 32, 32))
pred_y = predict_op(x)
print(pred_y)

# Convert model to Tensorflow and run it on the same batch
predict_tf = jax2tf.convert(predict_op)
print(predict_tf(np.array(x)))

class Objax2Tf(tf.Module):

  def __init__(self, module: objax.Module):
    assert isinstance(module, objax.Module), 'Input argument to Objax2Tf must be an Objax module.'

    module_vars = module.vars()

    def wrapped_op(tensor_list: List[JaxArray], kwargs, *args):
        original_values = module_vars.tensors()
        try:
            module_vars.assign(tensor_list)
            return module(*args, **kwargs)
        finally:
            module_vars.assign(original_values)

    tf_function = jax2tf.convert(wrapped_op)
    self._all_vars = [tf.Variable(v) for v in module_vars.tensors()]
    self._call = tf.function(
        lambda *args, **kwargs: tf_function(self._all_vars, kwargs, *args),
        autograph=False)
    
  def __call__(self, *args, **kwargs):
    return self._call(*args, **kwargs)

predict_tf = Objax2Tf(predict_op)
print(predict_tf(np.array(x)))