tensorflow 2 : loss using hidden layers output

user_11786464 0
Viewed 786

I am trying to implement the OSME MAMC model describe in article https://arxiv.org/abs/1806.05372.

I'm stuck where I have to add a cost that doesn't depend on y_true and y_pred but on hidden layers and y_true.

It can't be right as tensorflow custom loss, for which we need y_true and y_pred.

I wrote the model into class, then tried to use gradient tape to add NPairLoss to Softmax output loss, but gradient is NaN during training. I think my approach isn't good, but I have no idea how to design / write it.

Here my model :

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

then, for training :

@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)

Thanks for any help :)

conv-neural-network tensorflow python
1 Answers

You can use tensorflow's add_loss.

model.compile() loss functions in Tensorflow always take two parameters y_true and y_pred. Using model.add_loss() has no such restriction and allows you to write much more complex losses that depend on many other tensors, but it has the inconvenience of being more dependent on the model, whereas the standard loss functions work with just any model.

You can find the official documentation of add_loss here. Add loss tensor(s), potentially dependent on layer inputs. This method can be used inside a subclassed layer or model's call function, in which case losses should be a Tensor or list of Tensors. There are few example in the documentation to explain the add_loss.

This method can also be called directly on a Functional Model during construction. In this case, any loss Tensors passed to this Model must be symbolic and be able to be traced back to the model's Inputs. These losses become part of the model's topology and are tracked in get_config.

Example :

inputs = tf.keras.Input(shape=(10,))
x = tf.keras.layers.Dense(10)(inputs)
outputs = tf.keras.layers.Dense(1)(x)
model = tf.keras.Model(inputs, outputs)
# Activity regularization.
model.add_loss(tf.abs(tf.reduce_mean(x)))

You can call self.add_loss(loss_value) from inside the call method of a custom layer. Here's a simple example that adds activity regularization.

Example:

class ActivityRegularizationLayer(layers.Layer):

  def call(self, inputs):
    self.add_loss(tf.reduce_sum(inputs) * 0.1)
    return inputs  # Pass-through layer.

inputs = keras.Input(shape=(784,), name='digits')
x = layers.Dense(64, activation='relu', name='dense_1')(inputs)

# Insert activity regularization as a layer
x = ActivityRegularizationLayer()(x)

x = layers.Dense(64, activation='relu', name='dense_2')(x)
outputs = layers.Dense(10, name='predictions')(x)

model = keras.Model(inputs=inputs, outputs=outputs)
model.compile(optimizer=keras.optimizers.RMSprop(learning_rate=1e-3),
              loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True))

# The displayed loss will be much higher than before
# due to the regularization component.
model.fit(x_train, y_train,
          batch_size=64,
          epochs=1)

You can find good example using add_loss here and here with explanations.

Hope this answers your question. Happy Learning.

user_134652580
Related