Пользовательская функция потерь в керасе / тензор потока

Question

Предположим, что я хочу построить custom loss function. Эта функция потерь должна измерять difference между two sets of weigths. Предположим, у нас есть binary classificator. И мы хотим вычесть весовые коэффициенты для предсказания class0 из весовых коэффициентов для class1.

Вопрос: Как я могу извлечь весовые коэффициенты в модель керас и чем вычитать эти веса? Например:

# define a composite model for updating generators by adversarial and cycle loss
def Model(generator,aux1,aux2,aux3,aux4,aux5 segmentation_model, image_shape):

  generator.trainable = True

  aux1.trainable=False
  aux2.trainable=False
  aux3.trainable=False
  aux4.trainable=False
  aux5.trainable=False

  segmentation_model.trainable=False

  # generated element
  input_B = Input(shape=image_shape)
  gen1_out = generator(input_B)

  # segmentation element, GAP = Global Activation Map (shape=(batch_size,channels))
  seg_out,GAP5,GAP4,GAP3,GAP2,GAP1=segmentation_model(gen1_out)

  # binary Aux-Classificator, output=[0,1]
  a1=aux1(GAP1)
  a2=aux2(GAP2)
  a3=aux3(GAP3)
  a4=aux4(GAP4)
  a5=aux5(GAP5)

  # ---> Now extract the weigths for class 0 and for class 1 and subtract them

  ...
  ...
  ...
  diff_a1= ....



  # define model
  model = Model([input_B] ,[seg_out,diff_a1,diff_a2,diff_a3,diff_a4,diff_a5])

  opt = Adam(lr=0.0002, beta_1=0.5)
  model.summary()
  return model

Является ли здесь уместным lambda layer, или мне лучше написать custom layer с использованием только методов call и output_shape?