与tf.keras.predict和tf.data.Dataset保持示例索引对应

时间:2019-04-30 19:53:36

标签: python tensorflow keras

我正在TensorFlow2中使用tf.keras API。我有100,000张图像左右保存为TFRecords(每条记录128张图像)。每个记录都有一个输入图像,目标图像和帧索引。我找不到一种干净的方法来使帧索引与预测保持一致。

这里是一个例子,除了我用NumPy数组构建数据集,而不是从TFRecords中读取:

import tensorflow as tf
from tensorflow import keras
import numpy as np

# build dummy tf.data.Dataset
x = np.random.random(10000).astype(np.float32)
y = x + np.random.random(10000).astype(np.float32) * 0.1
idx = np.arange(10000, dtype=np.uint16)
np.random.shuffle(idx)  # frames are random in my TFRecord files
ds = tf.data.Dataset.from_tensor_slices((x, y, idx))
# pretend ds returned from TFRecord
ds = ds.map(lambda f0, f1, f2: (f0, f1))  # strip off idx
ds = ds.batch(32)

# build and train model
x = keras.Input(shape=(1,))
y_hat = keras.layers.Dense(1)(x)  # i.e. linear regression
model = keras.Model(x, y_hat)
model.compile('sgd', 'mse')
history = model.fit(ds, epochs=5)

# predict 1 batch
model.predict(ds, steps=1)

再次短暂地读取数据集以提取索引(容易出错),是否有一种干净的方法来使预测与图像索引保持一致?在TF1.x中很简单。但是我想利用TF2中干净的Keras compile(),fit(),predict()API。

1 个答案:

答案 0 :(得分:0)

好吧,我想得太刻苦,实际上很容易。进行预测时,只需将索引添加到数据集,然后在进行批处理时拉出索引:

rt tensorflow as tf
from tensorflow import keras
import numpy as np

def build_dataset(mode):
    np.random.seed(1)
    x = np.random.random(10000).astype(np.float32)
    y = x + np.random.random(10000).astype(np.float32) * 0.1
    idx = np.arange(10000, dtype=np.uint16)
    if mode == 'train':
        ds = tf.data.Dataset.from_tensor_slices((x, y))
        ds = ds.shuffle(128)
    else:
        ds = tf.data.Dataset.from_tensor_slices((x, idx))
    ds = ds.batch(32)

    return ds

# build and train simple linear regression model
x_tf = keras.Input(shape=(1,))
yhat_tf = keras.layers.Dense(1)(x_tf)
model = keras.Model(x_tf, yhat_tf)
model.compile(optimizer='sgd', loss='mse')
ds = build_dataset('train')
history = model.fit(ds, epochs=5)

# predict 1 batch
ds = build_dataset('predict')
for batch in ds:
    x_tf, indices_tf = batch 
    yhat_np = model.predict(x_tf)
    break