tensorflow2.0 - 如何使用具有多输出和自定义损失函数的 tf.data.dataset 运行 keras.fit？

Question

我正在使用tf2.0并且想要使用tf.keras和tf.data.dataset训练我的网络。但是，我正在努力tf.keras.fit与tf.data.dataset多输出和自定义损失函数一起使用。

我的 tensorflow 版本是tf2.0，这是我尝试但失败的示例代码。

import tensorflow as tf
import numpy as np

# define model
inputs = tf.keras.Input((512,512,3), name='model_input')

x = tf.keras.layers.Conv2D(filters=256, kernel_size=3, padding='same',kernel_initializer=tf.random_normal_initializer(stddev=0.01), name='conv1')(inputs)
x = tf.keras.layers.Conv2D(filters=256, kernel_size=3, padding='same',kernel_initializer=tf.random_normal_initializer(stddev=0.01), name='conv2')(x)

output1 = tf.keras.layers.Conv2D(filters=256, kernel_size=3, padding='same',kernel_initializer=tf.random_normal_initializer(stddev=0.01), name='output1')(x)
output2 = tf.keras.layers.Conv2D(filters=256, kernel_size=3, padding='same',kernel_initializer=tf.random_normal_initializer(stddev=0.01), name='output2')(x)

model = tf.keras.Model(inputs, [output1, output2])

# define dataset
def parse_func(single_data): # just for example case
    input = single_data
    output1 = single_data
    output2 = single_data
    weight1 = output1
    weight2 = output2
    return input, output1, output2, weight1, weight2

def tf_parse_func(single_data):
    return tf.py_function(parse_func, [single_data], [tf.float32, tf.flaot32, tf.float32, tf.flaot32, tf.float32])

data = np.random.rand(10, 512, 512, 3)
dataset = tf.data.Dataset.from_tensor_slices(data)
dataset = dataset.map(tf_parse_func, num_parallel_calls=tf.data.experimental.AUTOTUNE)
dataset = dataset.batch(2, drop_remainder=True)

# def loss func
def loss_fn1(label, pred):
    return tf.reduce_mean(tf.keras.losses.MSE(label, pred))
def loss_fn2(label, pred):
    return tf.nn.l2_loss(label-pred)

# start training

model.compile(loss={'output1':loss_fn1, 'output2':loss_fn2},
              loss_weights={'output1':1, 'output2':2},
              optimizer=tf.keras.optimizers.Adam())

model.fit(dataset, epochs=5)

其实我想通过

loss_weights={'output1': 1, 'output2': 2}

像这样

loss_weights={'output1': weight1, 'output2': weight2},

但我不知道该怎么做。最好weight1/weight2作为损失函数参数传递，但我不知道该怎么做。我想loss_fn1使用output1, weight1fromdataset和loss_fn2使用output2, weight2.

当我运行上面的代码时，我收到如下错误：

2019-10-22 20:47:40.551618: W tensorflow/core/framework/cpu_allocator_impl.cc:81] Allocation of 62914560 exceeds 10% of system memory.
      1/Unknown - 0s 28ms/stepTraceback (most recent call last):
  File "tools/keras_train_test.py", line 65, in <module>
    model.fit(dataset, epochs=5)
  File "/usr/local/lib/python3.6/dist-packages/tensorflow_core/python/keras/engine/training.py", line 728, in fit
    use_multiprocessing=use_multiprocessing)
  File "/usr/local/lib/python3.6/dist-packages/tensorflow_core/python/keras/engine/training_v2.py", line 324, in fit
    total_epochs=epochs)
  File "/usr/local/lib/python3.6/dist-packages/tensorflow_core/python/keras/engine/training_v2.py", line 123, in run_one_epoch
    batch_outs = execution_function(iterator)
  File "/usr/local/lib/python3.6/dist-packages/tensorflow_core/python/keras/engine/training_v2_utils.py", line 86, in execution_function
    distributed_function(input_fn))
  File "/usr/local/lib/python3.6/dist-packages/tensorflow_core/python/eager/def_function.py", line 457, in __call__
    result = self._call(*args, **kwds)
  File "/usr/local/lib/python3.6/dist-packages/tensorflow_core/python/eager/def_function.py", line 503, in _call
    self._initialize(args, kwds, add_initializers_to=initializer_map)
  File "/usr/local/lib/python3.6/dist-packages/tensorflow_core/python/eager/def_function.py", line 408, in _initialize
    *args, **kwds))
  File "/usr/local/lib/python3.6/dist-packages/tensorflow_core/python/eager/function.py", line 1848, in _get_concrete_function_internal_garbage_collected
    graph_function, _, _ = self._maybe_define_function(args, kwargs)
  File "/usr/local/lib/python3.6/dist-packages/tensorflow_core/python/eager/function.py", line 2150, in _maybe_define_function
    graph_function = self._create_graph_function(args, kwargs)
  File "/usr/local/lib/python3.6/dist-packages/tensorflow_core/python/eager/function.py", line 2041, in _create_graph_function
    capture_by_value=self._capture_by_value),
  File "/usr/local/lib/python3.6/dist-packages/tensorflow_core/python/framework/func_graph.py", line 915, in func_graph_from_py_func
    func_outputs = python_func(*func_args, **func_kwargs)
  File "/usr/local/lib/python3.6/dist-packages/tensorflow_core/python/eager/def_function.py", line 358, in wrapped_fn
    return weak_wrapped_fn().__wrapped__(*args, **kwds)
  File "/usr/local/lib/python3.6/dist-packages/tensorflow_core/python/keras/engine/training_v2_utils.py", line 66, in distributed_function
    model, input_iterator, mode)
  File "/usr/local/lib/python3.6/dist-packages/tensorflow_core/python/keras/engine/training_v2_utils.py", line 112, in _prepare_feed_values
    inputs, targets, sample_weights = _get_input_from_iterator(inputs)
  File "/usr/local/lib/python3.6/dist-packages/tensorflow_core/python/keras/engine/training_v2_utils.py", line 145, in _get_input_from_iterator
    x, y, sample_weights = next_element
ValueError: too many values to unpack (expected 3)

我尝试了许多方法和其他方法，但我发现我无法完成这项工作。那么有人可以帮助我吗？非常感谢！

标签： tensorflow2.0tf.keras

Answer 1

我可以使用model.fit_generator而不是model.fit让它工作。下面是成功运行的源代码：

import tensorflow as tf
import numpy as np

# define model
inputs = tf.keras.Input((112, 112, 3), name='model_input')

x = tf.keras.layers.Conv2D(filters=256, kernel_size=3, padding='same',
                           kernel_initializer=tf.random_normal_initializer(stddev=0.01), name='conv1')(inputs)
x = tf.keras.layers.Conv2D(filters=256, kernel_size=3, padding='same',
                           kernel_initializer=tf.random_normal_initializer(stddev=0.01), name='conv2')(x)

output1 = tf.keras.layers.Conv2D(filters=3, kernel_size=3, padding='same',
                                 kernel_initializer=tf.random_normal_initializer(stddev=0.01), name='output1')(x)
output2 = tf.keras.layers.Conv2D(filters=3, kernel_size=3, padding='same',
                                 kernel_initializer=tf.random_normal_initializer(stddev=0.01), name='output2')(x)

model = tf.keras.Model(inputs, [output1, output2])


# define dataset
def parse_func(single_data):  # just for example case
    input = single_data
    output1 = single_data
    output2 = single_data
    weight1 = output1
    weight2 = output2
    return input, output1, output2, weight1, weight2


def tf_parse_func(single_data):
    input, output1, output2, weight1, weight2 = tf.py_function(parse_func, [single_data], [tf.float32, tf.float32, tf.float32, tf.float32, tf.float32])
    return input, output1, output2, weight1, weight2


data = np.random.rand(10, 112, 112, 3).astype(np.float32)
dataset = tf.data.Dataset.from_tensor_slices(data).repeat(-1)
dataset = dataset.map(tf_parse_func, num_parallel_calls=tf.data.experimental.AUTOTUNE)
dataset = dataset.batch(2, drop_remainder=True)

def generator():
    for input, output1, output2, weight1, weight2 in dataset:
        output1 = tf.concat([output1, weight1], axis=-1)
        output2 = tf.concat([output2, weight2], axis=-1)
        yield input, [output1, output2]

# def loss func
def loss_fn1(label, pred):
    weight = label[..., 3:]
    label = label[..., :3]
    return tf.reduce_mean(tf.keras.losses.MSE(label*weight, pred*weight))


def loss_fn2(label, pred):
    weight = label[..., 3:]
    label = label[..., :3]
    return tf.nn.l2_loss(label*weight - pred*weight)


# start training

model.compile(loss={'output1': loss_fn1, 'output2': loss_fn2},
              loss_weights={'output1': 1, 'output2': 2},
              optimizer=tf.keras.optimizers.Adam())

# model.fit(dataset, epochs=5)
model.fit_generator(generator(), steps_per_epoch=10, epochs=5)

输出是：

Epoch 1/5
10/10 [==============================] - 7s 661ms/step - loss: 6814.9424 - output1_loss: 0.0673 - output2_loss: 3407.4375
Epoch 2/5
10/10 [==============================] - 7s 656ms/step - loss: 1858.2006 - output1_loss: 0.0669 - output2_loss: 929.0669
Epoch 3/5
10/10 [==============================] - 7s 658ms/step - loss: 1141.7914 - output1_loss: 0.0403 - output2_loss: 570.8755
Epoch 4/5
10/10 [==============================] - 7s 656ms/step - loss: 854.0343 - output1_loss: 0.0341 - output2_loss: 427.0001
Epoch 5/5
10/10 [==============================] - 7s 656ms/step - loss: 708.3558 - output1_loss: 0.0179 - output2_loss: 354.1689

我认为这是一个噱头，仍然希望任何人的任何答案。