1

I want to have retrieval of GCS object/any S3 object as a part of the Model, as a first layer which will obtain features based on the filename, because it will lower the networking overhead, and I am trying to wrap the download into the tf.function, but no success. Here is MWE:

import tensorflow as tf
@tf.function
def load_file(a):
    if tf.is_tensor(a):
        a_path = tf.strings.substr(a, 0, 2) + "/" + a
    else:
        a_path = a[0:2] + "/" + a
    with tf.io.gfile.GFile("gs://some_bucket" + a_path) as f:
        return f.read()
load_file(tf.constant("file3"))

which raises error

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
Cell In [22], line 9
      7     with tf.io.gfile.GFile("gs://some_bucket" + a_path) as f:
      8         return f.read()
----> 9 load_file(tf.constant("file3"))

File /opt/conda/envs/wanna-hmic/lib/python3.9/site-packages/tensorflow/python/util/traceback_utils.py:153, in filter_traceback.<locals>.error_handler(*args, **kwargs)
    151 except Exception as e:
    152   filtered_tb = _process_traceback_frames(e.__traceback__)
--> 153   raise e.with_traceback(filtered_tb) from None
    154 finally:
    155   del filtered_tb

File /opt/conda/envs/wanna-hmic/lib/python3.9/site-packages/tensorflow/python/framework/func_graph.py:1147, in func_graph_from_py_func.<locals>.autograph_handler(*args, **kwargs)
   1145 except Exception as e:  # pylint:disable=broad-except
   1146   if hasattr(e, "ag_error_metadata"):
-> 1147     raise e.ag_error_metadata.to_exception(e)
   1148   else:
   1149     raise

TypeError: in user code:

File "/tmp/ipykernel_4006/3877294148.py", line 8, in load_file  *
    return f.read()

TypeError: __init__(): incompatible constructor arguments. The following argument types are supported:
    1. tensorflow.python.lib.io._pywrap_file_io.BufferedInputStream(filename: str, buffer_size: int, token: tensorflow.python.lib.io._pywrap_file_io.TransactionToken = None)

Invoked with: <tf.Tensor 'add_2:0' shape=() dtype=string>, 524288

the code works well in eager mode with load_file("file3") but in order to perform well, I need it to work even in the graph mode.

1 Answer 1

1

tf.io.read_file does the trick. Modifying the whole code to

@tf.function
def load_file(a):
    a = tf.convert_to_tensor(a)
    a_path = tf.strings.substr(a, 0, 2) + "/" + a
    return tf.io.read_file("gs://some_bucket" + a_path)

makes it work in both eager and graph environment, and makes the input variable consistently a tf tensor.

Sign up to request clarification or add additional context in comments.

Comments

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Start asking to get answers

Find the answer to your question by asking.

Ask question

Explore related questions

See similar questions with these tags.