条件计算的中间层损失计算

如何解决条件计算的中间层损失计算

我想创建一个基于 MLP 的自定义 CNN 模型（多尺度），由几个并行的小网络（胶囊）组成。这些简单的小型网络被实例化为每个卷积尺度（即 3x3、5x5）的自定义层（conv2d->Flatten->Dense）。这些胶囊网络的目的是产生中间损失意识，以使用 CNN 模型减少整体全局损失。我写了一些粗略的代码，但我无法编写正确的代码来使用这些胶囊计算局部损失。代码如下：

from tensorflow.keras import layers
import tensorflow as tf
from tensorflow.keras.datasets import mnist
from tensorflow.keras.layers import Layer

class capsule(tf.keras.layers.Layer):
  def __init__(self):
     super(capsule,self).__init__()
     self.loss_fn = tf.keras.losses.CategoricalCrossentropy(from_logits=True)
     self.Flatten = tf.keras.layers.Flatten()
     self.conv2D = tf.keras.layers.Conv2D(3,3,(1,1),padding='same',activation='relu',name="LocalLoss3x3")
     self.classifier = tf.keras.layers.Dense(10,activation='softmax',name='capsule3Output')

 def call(self,inputs):
    x=self.conv2D(inputs)
    x=self.Flatten(x)
    x=self.classifier(x)
    pred=self(x_train)
    loss=self.loss_fn(pred,y_train)
    #self.add_loss(self.rate * tf.reduce_sum(tf.square(inputs)))
    return loss,x

(x_train,y_train),(x_test,y_test)=  mnist.load_data()
from tensorflow.keras import layers
class SparseMLP(tf.keras.models.Model):

def __init__(self,output_dim):
  super(SparseMLP,self).__init__()
  self.dense_1 = layers.Dense(1,activation=tf.nn.relu)
  self.capsule = capsule()
  self.dense_2 = layers.Dense(output_dim)

def call(self,inputs):
  x = self.dense_1(inputs)
  loss,x = self.capsule(inputs)
  return self.dense_2(x)


mlp = SparseMLP(10)
#x_train=x_train.reshape(-1,28,1)
y = mlp(x_train)