自作loss関数でエラーがでます[Python, TensorFlow, Keras]

Question

実現したいこと

自作のloss関数を走らせたい。
（可能ならばgeneraterを使用しない手法が知りたいです。）

前提

現在医療系大学にて超解像を勉強中です。
今年度から本格的に取り組み始めた初学者です。
SSIM lossなどの自作loss関数などの作成はできましたが、
トレーニングの低画質画像をlossに含もうとするとエラーが出ます。
Python, TensorFlow, Kerasを使用。

発生している問題・エラーメッセージ

超解像の学習時に自作のloss関数でlow resolution画像を計算に含みたいができません。
本来def loss(y_true, y_pred):より高解像度画像と推論画像とでlossを計算しますが、
ここに低解像度画像を含みたいです。

    tensors = pywrap_tfe.TFE_Py_Execute(ctx._handle, device_name, op_name,
    tensorflow.python.framework.errors_impl.InvalidArgumentError:  Incompatible shapes: [16,48,48,1] vs. [41,48,48,1]
	 [[node loss/sub_1 (defined at \Users\PycharmProjects\DL\loss_function.py:105) ]] 
    [Op:__inference_train_function_17330]

Errors may have originated from an input operation.
Input Source operations connected to node loss/sub_1:
  IteratorGetNext (defined at /Users/PycharmProjects/DL/RUN.py:153)

Function call stack:
train_function

該当のソースコード

# 前略

def custom_loss(x_true):
    def loss(y_true, y_pred):
        return K.mean(K.square(y_pred - y_true) + K.square(y_true - x_true))
    return loss

model.compile(optimizer=Adam(lr=lr, beta_1=0.9, beta_2=0.999, epsilon=1e-8),
              metrics=[psnr], loss=[custom_loss(x_true=imagelow)])

training = model.fit(imagelow, imagehigh, epochs=ep, batch_size=bs,
                     validation_split=vs, shuffle=True, verbose=1)

# 後略

試したこと

ネットでかなり調べたつもりですが、x_trueを使ってるサイトが見つかりません。
custom lossなどの作成についても調べましたが、（おそらく）batchサイズが異なることしかわかりませんでした。

補足情報（FW/ツールのバージョンなど）

Python, TensorFlow, Kerasを使用。

Comments

[meg_]

> 自作loss関数でエラーがでます
自作でなければエラーなく実行できる、と理解して良いでしょうか？

[ps_aux_grep]

バッチサイズが違うのでエラーになっていますが，このままでは一生解決ができないスタイルのKerasの使い方をしているように見えます．Subclassing APIを使ってtrain_stepを実装するか，add_lossを利用する
https://www.tensorflow.org/guide/keras/custom_layers_and_models#the_add_metric_method
などして解決してください．

[nishimizuumi]

自作でなければ（MAE,MSEなど）実行できています

Answer 1

x_trueを使うにはSubclassing APIを利用して，keras.Modelのcompute_lossをカスタマイズしてください．質問にあるcustom_lossを9行目で記述しました．

import tensorflow as tf
from keras.layers import Input, Activation, Add, BatchNormalization, Conv2D
from keras.models import Model

class OriginalModel(Model):
    def compute_loss(self, x, y, y_pred, sample_weight):
        loss = super().compute_loss(x, y, y_pred, sample_weight)
        # append your loss
        loss += tf.math.reduce_mean(tf.math.squared_difference(y, y_pred)) + tf.math.reduce_mean(tf.math.squared_difference(y, x))
        return loss

def ResBlock(ch, k):
    def apply(inputs):
        x = Conv2D(ch, k, padding = "same", kernel_initializer = "he_uniform")(inputs)
        x = BatchNormalization()(x)
        x = Activation("swish")(x)
        x = Conv2D(ch, k, padding = "same")(x)
        return Add()([x, inputs])
    return apply
        
def build_model():
    inputs = Input(shape = (32, 32, 3))
    x = Conv2D(64, 7, padding = "same")(inputs)
    for _ in range(8):
        x = ResBlock(64, 3)(x)
    x = Conv2D(3, 3, padding = "same")(x)
    return OriginalModel(inputs, x) # Use your model included original loss function

model = build_model()

model.compile(
    optimizer = "adam",
    loss = "huber",
    metrics = ["mse", "mae"]
)

もちろんコメントでも申し上げた通り，keras.Modelのtrain_stepをカスタマイズして損失を新たに与えるのも可能です．

Comments

[nishimizuumi]

ご回答くださりありがとうございます。またmodelの記述も含めご丁寧に回答くださりありがとうございます。現在所用でご回答の内容を自身で実行できておりませんが、早急にご回答くださったためベストアンサーにさせていただきました。自身で実行し改めて追加の質問をさせていただくことがあるかもしれません。よろしくお願いいたします。