数据集返回多个参数，用这些参数调用拟合。使用整个 PrefetchDataset 调用 fit

Question

我有以下代码可以正常工作。

但问题是它只适用于一批，因为我正在使用

next(iter)

我创建了一个张量流数据集，它必须为我的问题返回 3 个值。

问题是，当我要调用

fit

时，我必须分别使用这 3 个值才能正确调用网络架构。

所以，我的问题是如何使用

fit

调用

PrefetchDataset

。

import tensorflow as tf
import numpy as np
from tensorflow.keras.layers import Input, Dense,  Activation, \
    Conv2DTranspose, Conv2D, Reshape
from tensorflow.keras.models import Model

AUTOTUNE = tf.data.experimental.AUTOTUNE

def scale(X, a=-1, b=1, dtype='float32'):
    if a > b:
        a, b = b, a
    xmin = tf.cast(tf.math.reduce_min(X), dtype=dtype)
    xmax = tf.cast(tf.math.reduce_max(X), dtype=dtype)
    X = (X - xmin) / (xmax - xmin)
    scaled = X * (b - a) + a
    return scaled, xmin, xmax

def set_shape_b(x, y, z):
    x = tf.reshape(x,  [16, 16, 2])
    y = tf.reshape(y, [1])
    z = tf.reshape(z, [1])
    return x, y, z

def set_shape_a(x, y, z):
    x = tf.reshape(x,  [4, 4, 2])
    y = tf.reshape(y, [1])
    z = tf.reshape(z, [1])
    return x, y, z

def First(lr):
    inp = Input(lr)
    x = Dense(16)(inp)
    x = Reshape((4, 4, 16))(x)
    x = Conv2DTranspose(2, kernel_size=3, strides=2, padding='same')(x)
    x = Conv2DTranspose(2, kernel_size=3, strides=2, padding='same')(x)
    output = Activation('tanh')(x)
    model = Model(inp, output, name='First')
    return model
    
def Second(hr):
    inp = Input(hr)
    x = Dense(16)(inp)
    x = Conv2D(2, kernel_size=3, strides=2, padding='same')(x)
    x = Conv2D(2, kernel_size=3, strides=2, padding='same')(x)
    output = Dense(1, activation='sigmoid')(x)
    model = Model(inputs=inp, outputs=output, name='Second')
    return model
    

def build_model(First, Second):
    inp = Input(shape=INP)
    gen = First(inp)
    output = Second(gen)
    model = Model(inputs=inp , outputs=[gen, output], name='model')
    return model

# Preproces --------------- #
a = np.random.random((20, 4, 4, 2)).astype('float32')
b = np.random.random((20, 16, 16, 2)).astype('float32')

dataset_a = tf.data.Dataset.from_tensor_slices(a)
dataset_b = tf.data.Dataset.from_tensor_slices(b)

dataset_b = dataset_b.map(lambda x: tf.py_function(scale,
                                                   [x], 
                                                   (tf.float32, tf.float32, tf.float32)))
dataset_b = dataset_b.map(set_shape_b)

dataset_a = dataset_a.map(lambda x: tf.py_function(scale,
                                                   [x], 
                                                   (tf.float32, tf.float32, tf.float32)))
dataset_a = dataset_a.map(set_shape_a)
 
dataset_ones = tf.data.Dataset.from_tensor_slices(tf.ones((len(b), 4, 4, 1)))   

dataset = tf.data.Dataset.zip((dataset_a, (dataset_b, dataset_ones)))

dataset = dataset.cache()
dataset = dataset.batch(2)
dataset = dataset.prefetch(buffer_size=AUTOTUNE)

# Prepare models -------------------- #
INP = (4, 4, 2)
OUT = (16, 16, 2)

first = First(INP)
second = Second(OUT)
model = build_model(first, second)

model.compile(loss=['mse', 'binary_crossentropy'],
              optimizer= tf.keras.optimizers.Adam(learning_rate=1e-4))


train_l, (train_h, train_ones) = next(iter(dataset))


# train ------------------
model.fit(train_l[0],
          [train_h[0], train_ones],
          epochs=2)
              

Answer 1

按照上面的评论，要解决您的问题，您可以应用自定义函数以仅返回目标值。另外，请查看 tf.data.Dataset.map 以获取参考。

def set_shape(x, y, z, dims):
    x = tf.reshape(x,  dims)
    y = tf.reshape(y, [1])
    z = tf.reshape(z, [1])
    return x, y, z

dataset_a = dataset_a.map(lambda x, y, z: set_shape(x, y, z, dims=[4, 4, 2]))
dataset_b = dataset_b.map(lambda x, y, z: set_shape(x, y, z, dims=[16, 16, 2]))

def only_scale(x, y, z):
    return x

dataset_a = dataset_a.map(only_scale)
dataset_b = dataset_b.map(only_scale)

压缩和批处理数据。

dataset = tf.data.Dataset.zip(
    (dataset_a, (dataset_b, dataset_ones))
)
dataset = dataset.cache()
dataset = dataset.batch(2)
dataset = dataset.prefetch(buffer_size=AUTOTUNE)

a, b = next(iter(dataset)) 
a.shape, b[0].shape, b[1].shape
(TensorShape([2, 4, 4, 2]),
 TensorShape([2, 16, 16, 2]),
 TensorShape([2, 4, 4, 1]))

现在，我们可以将它传递给 fit 方法。

# train ------------------
model.fit(
    dataset,
    epochs=2
)

Epoch 1/2
2s 6ms/step - loss: 1.0283 - First_loss: 0.3368 - Second_loss: 0.6914
Epoch 2/2
0s 4ms/step - loss: 1.0228 - First_loss: 0.3367 - Second_loss: 0.6860