使用 Tensorflow 2 输入管道将 trainX 和 trainY 分开进行多损失训练

如何解决使用 Tensorflow 2 输入管道将 trainX 和 trainY 分开进行多损失训练

我正在构建一个网络，将一张图片作为输入并输出两个预测，并根据这两个预测计算损失函数。所以：

输入（trainX）：RGBD 输出1（trainGT）：GT 输出 2（trainError）：错误

当我使用 TF 的标准输入管道时，它会压缩输入和目标。但是，当我定义损失函数时，我需要将两个目标分开。这是我的输入管道代码：

@tf.function
def load_image(point_file):
    name = tf.strings.split(point_file,'/')[-1]
    GT = tf.io.read_file(GT_path + name)
    RGBD = tf.io.read_file(RGBD_path + name)
    error = tf.io.read_file(error_path + name)
    GT = tf.image.decode_png(GT)
    RGBD = tf.image.decode_png(RGBD)
    error = tf.image.decode_png(error)[:,:,0]

    GT = tf.cast(GT,tf.float32)// 255.0
    RGBD = tf.cast(RGBD,tf.float32)/ 255.0
    error = tf.cast(error,tf.float32)
    error = tf.reshape(error,1024)

    return RGBD,GT,error

train_dataset = tf.data.Dataset.list_files(point_path + 'train/*.png')
train_dataset = train_dataset.map(load_image,num_parallel_calls=tf.data.experimental.AUTOTUNE)
train_dataset = train_dataset.shuffle(BUFFER_SIZE)
train_dataset = train_dataset.batch(BATCH_SIZE).repeat()

test_dataset = tf.data.Dataset.list_files(point_path+'test/*.png')
test_dataset = test_dataset.map(load_image)
test_dataset = test_dataset.batch(BATCH_SIZE)

稍后在 fit 函数中我想做的是：

losses = {
    "GT_output": loss_functions.weighted_dice_loss,"error_output": tf.keras.losses.MeanSquaredError(),}
lossWeights = {"GT_output": 0.9,"error_output": 0.1}

model.compile(optimizer='adam',loss=losses,loss_weights=lossWeights,)

model.fit(x=trainX,y={"GT_output": trainGT,"error_output": trainError},validation_data=(testX,{"GT_output": testGT,"error_output": testError}),epochs=EPOCHS,verbose=1)

但是有没有办法将 trainX、trainGT 和 trainError 与 train_dataset 分开？

谢谢...

解决方法

你可以这样做：

take_n = 1000
trainGT = train_dataset.take(take_n)
trainError = train_dataset.skip(take_n)