Невозможно передать метки VGG в пользовательской функции потери Keras

Question

При использовании последовательной модели с последовательностями (5D-тензоры в качестве выходов модели и таргет-модели) в сочетании с экземпляром модели VGG16 в пользовательской функции потери восприятия, keras выдает assert key in layers_to_output_shapes Ошибка подтверждения в "/ usr / local /lib/python3.6/dist-packages/keras/engine/network.py ", строка 676, в compute_output_shape. Это происходит при компиляции последовательной модели. Работает пользовательская реализация потерь L1 без VGG-модели, но с теми же тензорами и формами. Работает также последовательная модель без последовательностей (4D Tensors) и пользовательские потери с VGG-моделью.

Код для воспроизведения проблемы

def build_vgg_model(self, weights="/mnt/workspace/lm78463/diforem/data/weights/vgg16_weights_tf_dim_ordering_tf_kernels_notop.h5"):
    # Input image to extract features from
    img = Input(batch_shape=(self.seq_len, self.img_rows, self.img_cols, 3))

    # Mean center and rescale by variance as in PyTorch
    mean = [0.485*255, 0.456*255, 0.406*255]
    std = [0.229*255, 0.224*255, 0.225*255]
    processed = Lambda(lambda x: (x - mean) / std)(img)

    # Get the vgg network from Keras applications
    if weights in ['imagenet', None]:
        vgg = VGG16(weights=weights, include_top=False)
        # model = Model(inputs=base_model.input, outputs=base_model.get_layer('block4_pool').output)
    else:
        vgg = VGG16(weights=None, include_top=False)
        vgg.load_weights(weights, by_name=True)

    # Output the first three pooling layers
    vgg.outputs = [vgg.layers[i].output for i in self.vgg_layers]
    # Process the input images with mean center and rescale
    vgg_processed = vgg(processed)
    # Create model and compile
    model = Model(inputs=img, outputs=vgg_processed)
    model.trainable = False

    return model

def loss_total(self):
    def loss(y_trues, y_preds):
        o = y_preds[:, -1, :, :, :]
        t = y_trues[:, -1, :, :, :]

        vgg_outs = self.vgg(o)

        # assertion error
        vgg_gts = self.vgg(t)

        # Compute loss components
        l1 = self.l1(y_trues[:, -1, :, :, :], y_preds[:, -1, :, :, :])
        l2 = self.loss_perceptual(vgg_outs, vgg_gts)

        # Return loss function
        return l1 + 2 * l2
    return loss

def loss_perceptual(self, vgg_out, vgg_gt):
    loss = 0
    for o, g in zip(vgg_out, vgg_gt):
        loss += self.l1(o, g)
    return loss

@staticmethod
def l1(y_true, y_pred):
    if K.ndim(y_true) == 4:
        return K.mean(K.abs(y_pred - y_true), axis=[1, 2, 3])
    elif K.ndim(y_true) == 3:
        return K.mean(K.abs(y_pred - y_true), axis=[1, 2])
    else:
        raise NotImplementedError("Calculating L1 loss on 1D tensors? should not occur for this network")



seq = Sequential()
seq.add(ConvLSTM2D(filters=40, kernel_size=(3, 3),
                   input_shape=(None, 40, 40, 3),
                   padding='same', return_sequences=True))
seq.add(BatchNormalization())

seq.add(ConvLSTM2D(filters=40, kernel_size=(3, 3),
                   padding='same', return_sequences=True))
seq.add(BatchNormalization())

seq.add(ConvLSTM2D(filters=40, kernel_size=(3, 3),
                   padding='same', return_sequences=True))
seq.add(BatchNormalization())

seq.add(ConvLSTM2D(filters=40, kernel_size=(3, 3),
                   padding='same', return_sequences=True))
seq.add(BatchNormalization())

seq.add(Conv3D(filters=3, kernel_size=(3, 3, 3),
               activation='sigmoid',
               padding='same', data_format='channels_last'))
seq.compile(loss=self.loss_total(), optimizer='adadelta')

Другая информация / журналы

Traceback (most recent call last):
  File "convlstm_main.py", line 51, in <module>
    model = ConvLSTM(64, 64, 3, 10)
  File "/mnt/workspace/lm78463/RecurrentDiFoRem/src/model/conv_lstm.py", line 40, in __init__
    self.model = self.compile()
  File "/mnt/workspace/lm78463/RecurrentDiFoRem/src/model/conv_lstm.py", line 117, in compile
    optimizer='adadelta')
  File "/usr/local/lib/python3.6/dist-packages/keras/engine/training.py", line 229, in compile
    self.total_loss = self._prepare_total_loss(masks)
  File "/usr/local/lib/python3.6/dist-packages/keras/engine/training.py", line 692, in _prepare_total_loss
    y_true, y_pred, sample_weight=sample_weight)
  File "/usr/local/lib/python3.6/dist-packages/keras/losses.py", line 71, in __call__
    losses = self.call(y_true, y_pred)
  File "/usr/local/lib/python3.6/dist-packages/keras/losses.py", line 132, in call
    return self.fn(y_true, y_pred, **self._fn_kwargs)
  File "/mnt/workspace/lm78463/RecurrentDiFoRem/src/model/conv_lstm.py", line 192, in loss
    vgg_gts = self.vgg(t)
  File "/usr/local/lib/python3.6/dist-packages/keras/engine/base_layer.py", line 506, in __call__
    output_shape = self.compute_output_shape(input_shape)
  File "/usr/local/lib/python3.6/dist-packages/keras/engine/network.py", line 676, in compute_output_shape
    assert key in layers_to_output_shapes
AssertionError

Информация о системе

• Написал ли я собственный код (в отличие от использования примера каталога): Да
• Платформа и распространение ОС (например, Linux Ubuntu 16.04): Ubuntu 18.04 / nvcr.io/nvidia/tensorflow:19.12-tf1-py3 docker image
• Серверная часть TensorFlow (да / нет): да
• Версия TensorFlow: 1.14.0 + nv
• * Версия 1042 *: Python 3.6.9
• Версия CUDA / cuDNN: 10.1 / 7.6.4
• Модель GPU и память: 2x NVIDIA RTX2080Ti 11 ГБ

Ссылка на мою проблему с keras на Github