tensorflow 2：使用隐藏层输出的损失

我正在尝试实现https://arxiv.org/abs/1806.05372中描述的OSME MAMC模型。

我被困在必须增加不取决于y_true和y_pred而是取决于隐藏层和y_true的成本的地方。

tensorflow自定义损失是不对的，为此我们需要y_true和y_pred。

我将模型编写成类，然后尝试使用渐变磁带将NPairLoss添加到Softmax输出损耗，但是在训练过程中渐变为NaN。我认为我的方法不好，但是我不知道如何设计/编写它。

这是我的模型：

class OSME_network(tf.keras.Model):    
    def __init__(self, nbrclass=10, weight="imagenet",input_tensor=(32,32,3)):
        super(OSME_network, self).__init__()
        self.nbrclass = nbrclass
        self.weight = weight
        self.input_tensor=input_tensor
        self.Resnet_50=ResNet50(include_top=False, weights=self.weight, input_shape=self.input_tensor)
        self.Resnet_50.trainable=False
        self.split=Lambda(lambda x: tf.split(x,num_or_size_splits=2,axis=-1))
        self.s_1=OSME_Layer(ch=1024,ratio=16)
        self.s_2=OSME_Layer(ch=1024,ratio=16)
        self.fl1=tf.keras.layers.Flatten()
        self.fl2=tf.keras.layers.Flatten()
        self.d1=tf.keras.layers.Dense(1024, name='fc1')
        self.d2=tf.keras.layers.Dense(1024,name='fc2')
        self.fc=Concatenate()
        self.preds=tf.keras.layers.Dense(self.nbrclass,activation='softmax')

    @tf.function
    def call(self,x): #set à construire le model sequentiellement

        x=self.Resnet_50(x)     
        x_1,x_2=self.split(x)
        xx_1 = self.s_1(x_1)     
        xx_2 = self.s_2(x_2)
        xxx_1 = self.d1(xx_1)
        xxx_2 = self.d2(xx_2)
        xxxx_1 = self.fl1(xxx_1)
        xxxx_2 = self.fl2(xxx_2)
        fc = self.fc([xxxx_1,xxxx_2]) #fc1 + fc2
        ret=self.preds(fc)
        return xxxx_1,xxxx_2,ret
class OSME_Layer(tf.keras.layers.Layer):
    def __init__(self,ch,ratio):
        super(OSME_Layer,self).__init__()
        self.GloAvePool2D=GlobalAveragePooling2D()
        self.Dense1=Dense(ch//ratio,activation='relu')
        self.Dense2=Dense(ch,activation='sigmoid')
        self.Mult=Multiply()
        self.ch=ch
    def call(self,inputs):
        squeeze=self.GloAvePool2D(inputs)
        se_shape = (1, 1, self.ch)
        se = Reshape(se_shape)(squeeze)
        excitation=self.Dense1(se)
        excitation=self.Dense2(excitation)
        scale=self.Mult([inputs,excitation])
        return scale

class NPairLoss():
    def __init__(self):
        self._inputs = None
        self._y=None

    @tf.function
    def __call__(self,inputs,y):
        targets=tf.argmax(y, axis=1)
        b, p, _ = inputs.shape
        n = b * p

        inputs=tf.reshape(inputs, [n, -1])

        targets = tf.repeat(targets,repeats=p)

        parts = tf.tile(tf.range(p),[b])

        prod=tf.linalg.matmul(inputs,inputs,transpose_a=False,transpose_b=True)

        same_class_mask = tf.math.equal(tf.broadcast_to(targets,[n, n]),tf.transpose(tf.broadcast_to(targets,(n, n)))) 

        same_atten_mask = tf.math.equal(tf.broadcast_to(parts,[n, n]),tf.transpose(tf.broadcast_to(parts,(n, n))))

        s_sasc = same_class_mask & same_atten_mask
        s_sadc = (~same_class_mask) & same_atten_mask
        s_dasc = same_class_mask & (~same_atten_mask)
        s_dadc = (~same_class_mask) & (~same_atten_mask)

        loss_sasc = 0
        loss_sadc = 0
        loss_dasc = 0

        for i in range(n):
            #loss_sasc
            pos = prod[i][s_sasc[i]]
            neg = prod[i][s_sadc[i] | s_dasc[i] | s_dadc[i]]
            n_pos=tf.shape(pos)[0]
            n_neg=tf.shape(neg)[0]            
            pos = tf.transpose(tf.broadcast_to(pos,[n_neg,n_pos]))
            neg = tf.broadcast_to(neg,[n_pos,n_neg])
            exp=tf.clip_by_value(tf.math.exp(neg - pos),clip_value_min=0,clip_value_max=9e6) # need to clip value, else inf
            loss_sasc += tf.reduce_sum(tf.math.log(1 + tf.reduce_sum(exp,axis=1)))

            #loss_sadc
            pos = prod[i][s_sadc[i]]
            neg = prod[i][s_dadc[i]]
            n_pos = tf.shape(pos)[0]
            n_neg = tf.shape(neg)[0]
            pos = tf.transpose(tf.broadcast_to(pos,[n_neg,n_pos])) #np.transpose(np.tile(pos,[n_neg,1]))
            neg = tf.broadcast_to(neg,[n_pos,n_neg])#np.tile(neg,[n_pos,1])
            exp=tf.clip_by_value(tf.math.exp(neg - pos),clip_value_min=0,clip_value_max=9e6)            
            loss_sadc += tf.reduce_sum(tf.math.log(1 + tf.reduce_sum(exp,axis=1)))

            #loss_dasc
            pos = prod[i][s_dasc[i]]
            neg = prod[i][s_dadc[i]]
            n_pos = tf.shape(pos)[0]
            n_neg = tf.shape(neg)[0]
            pos = tf.transpose(tf.broadcast_to(pos,[n_neg,n_pos])) #np.transpose(np.tile(pos,[n_neg,1]))
            neg = tf.broadcast_to(neg,[n_pos,n_neg])#np.tile(neg,[n_pos,1])
            exp=tf.clip_by_value(tf.math.exp(neg - pos),clip_value_min=0,clip_value_max=9e6)
            loss_dasc += tf.reduce_sum(tf.math.log(1 + tf.reduce_sum(exp,axis=1)))
        return (loss_sasc + loss_sadc + loss_dasc) / n

然后，进行训练：

@tf.function
def train_step(x,y):
    with tf.GradientTape() as tape:
        fc1,fc2,y_pred=model(x,training=True)        
        stacked=tf.stack([fc1,fc2],axis=1)
        layerLoss=npair(stacked,y)
        loss=cce(y, y_pred) +0.001*layerLoss
    grads=tape.gradient(loss,model.trainable_variables)
    opt.apply_gradients(zip(grads,model.trainable_variables))
    return loss

model=OSME_network(weight="imagenet",nbrclass=10,input_tensor=(32, 32, 3))
model.compile(optimizer=opt, loss=categorical_crossentropy, metrics=["acc"])
model.build(input_shape=(None,32,32,3))

cce = tf.keras.losses.CategoricalCrossentropy(from_logits=True,name='categorical_crossentropy')
npair=NPairLoss()

for each batch :
    x=tf.Variable(x_train[start:end])
    y=tf.Variable(y_train[start:end])
    train_loss=train_step(x,y)

谢谢你的帮助 ：）