I have a Dataset that uses a keras model call within a map like this toy example:
ds=tf.data.Dataset.from_tensor_slices([tf.ones([1]) for i in range(10)])
model=tf.keras.models.Sequential([tf.keras.layers.Dense(1),tf.keras.layers.Dense(1)])
ds=ds.batch(4).map(lambda x:model(x))
I was wondering whether there is anyway to use the built-in model.predict(x) instead, since the map using the model call is quite slow (in my real project). I have tried
ds=tf.data.Dataset.from_tensor_slices([tf.ones([1]) for i in range(10)])
model=tf.keras.models.Sequential([tf.keras.layers.Dense(1),tf.keras.layers.Dense(1)])
ds=ds.batch(4).map(lambda x:model.predict(x))
and
ds=tf.data.Dataset.from_tensor_slices([tf.ones([1]) for i in range(10)])
model=tf.keras.models.Sequential([tf.keras.layers.Dense(1),tf.keras.layers.Dense(1)])
def predict(x):
return model.predict(x)
def predict_wrapper(x):
y=tf.py_function(predict,[x],tf.float32)
y.set_shape([None,None])
return y
ds=ds.batch(4).map(predict_wrapper)
for x in ds:
print(x)
Would it be possible? Would it make any difference in speed? My guess is that probably not since Dataset is already optimized for distributed strategy and it is like distributing ops within dristributed ops. But since I have no idea about the matter I thought I would ask.
Also I am working in Google Colab if that makes any difference.