constraint parameter for add_weight of layers.class does not accept functions.

[mannypaeza]

I am trying to convert a constraint function from Keras 2 to Keras 3. In Keras 2 API (with tensorflow), the code below works with constraint=masked_constraint(self.mask) in self.add_weight outputting a function.

import numpy as np
import os
import tensorflow as tf
from tensorflow.keras.layers import Input, Dense, Reshape, Layer, Activation
from tensorflow.keras import backend as K
from tensorflow.keras.initializers import Constant, RandomUniform

class Masked_Conv2D(Layer):
    """ Creates a trainable ring convolutional kernel with non zero entries between
    user specified radius_min and radius_max. Uses a random uniform non-negative
    initializer unless specified otherwise.

    Args:
        output_dim: int, default: 1
            number of output channels (number of kernels)

        kernel_size: (int, int), default: (5, 5)
            dimension of 2d bounding box

        strides: (int, int), default: (1, 1)
            stride for convolution (modifying that will downsample)

        radius_min: int, default: 2
            inner radius of kernel

        radius_max: int, default: 3
            outer radius of kernel (typically: 2*radius_max - 1 = kernel_size[0])

        initializer: 'uniform' or Keras initializer, default: 'uniform'
            initializer for ring weights. 'uniform' will choose from a non-negative
            random uniform distribution such that the expected value of the sum
            is 2.

        use_bias: bool, default: True
            add a bias term to each convolution kernel

    Returns:
        Masked_Conv2D: tensorflow.keras.layer
            A trainable layer implementing the convolution with a ring
    """
    def __init__(self, output_dim=1, kernel_size=(5,5), strides=(1,1),
               radius_min=2, radius_max=3, initializer='uniform',
               use_bias=True): #, output_dim):
        self.output_dim = output_dim
        self.kernel_size = kernel_size
        self.radius_min = radius_min
        self.radius_max = radius_max
        self.strides = strides
        self.use_bias = use_bias
        xx = np.arange(-(kernel_size[0]-1)//2, (kernel_size[0]+1)//2)
        yy = np.arange(-(kernel_size[0]-1)//2, (kernel_size[0]+1)//2)
        [XX, YY] = np.meshgrid(xx, yy)
        R = np.sqrt(XX**2 + YY**2)
        R[R<radius_min] = 0
        R[R>radius_max] = 0
        R[R>0] = 1
        self.mask = R  # the mask defines the non-zero pattern and will multiply the weights
        if initializer == 'uniform':
            self.initializer = RandomUniform(minval=0, maxval=2/R.sum())
        else:
            self.initializer = initializer
        super(Masked_Conv2D, self).__init__()

    def build(self, input_shape):
        # breakpoint()
        try:
            n_filters = input_shape[-1].value   # tensorflow < 2
        except:
            n_filters = input_shape[-1]         # tensorflow >= 2
        self.h = self.add_weight(name='h',
                                 shape= self.kernel_size + (n_filters, self.output_dim,),
                                 initializer=self.initializer,
                                 constraint=masked_constraint(self.mask), #look into
                                 trainable=True)

        self.b = self.add_weight(name='b',
                                 shape=(self.output_dim,),
                                 initializer=Constant(0),
                                 trainable=self.use_bias)
        super(Masked_Conv2D, self).build(input_shape)

    def call(self, x):
        #hm = tf.multiply(self.h, K.expand_dims(K.expand_dims(tf.cast(self.mask, float))))
        #hm = tf.multiply(hm, hm>0)
        #hm = tf.where(hm>0, hm, 0)
        # breakpoint() 
        y = K.conv2d(x, self.h, padding='same', strides=self.strides)
        if self.use_bias:
            y = y + tf.expand_dims(self.b, axis=0)
        return y

"Constraint Function I used"
def masked_constraint(R):
    """ Enforces constraint for kernel to be non-negative everywhere and zero outside the ring

    Args:
        R: np.array
            Binary mask that extracts 

    Returns:
        my_constraint: function
            Function that enforces the constraint
    """
    R = tf.cast(R, dtype=tf.float32)
    R_exp = tf.expand_dims(tf.expand_dims(R, -1), -1)
    def my_constraint(x):
        Rt = tf.tile(R_exp, [1, 1, 1, x.shape[-1]])
        Z = tf.zeros_like(x)
        return tf.where(Rt>0, x, Z)
    return my_constraint

( Photo of the model showing the function object in breakpoint()).

However, in Keras 3 (using pytorch for this), this function does not work. I tried to rewrite it as a Keras.constraints.Constraint class and it also fails to achieve my objective. Below is the following code (with both the function and class code written below).

class Masked_Conv2D(Layer): #Keras.layer
    """ Creates a trainable ring convolutional kernel with non zero entries between
    user specified radius_min and radius_max. Uses a random uniform non-negative
    initializer unless specified otherwise.

    Args:
        output_dim: int, default: 1
            number of output channels (number of kernels)

        kernel_size: (int, int), default: (5, 5)
            dimension of 2d boundaing box

        strides: (int, int), default: (1, 1)
            stride for convolution (modifying that will downsample)

        radius_min: int, default: 2
            inner radius of kernel

        radius_max: int, default: 3
            outer radius of kernel (typically: 2*radius_max - 1 = kernel_size[0])

        initializer: 'uniform' or Keras initializer, default: 'uniform'
            initializer for ring weights. 'uniform' will choose from a non-negative
            random uniform distribution such that the expected value of the sum
            is 2.

        use_bias: bool, default: True
            add a bias term to each convolution kernel

    Returns:
        Masked_Conv2D: keras.layer
            A trainable layer implementing the convolution with a ring
    """
    def __init__(self, output_dim=1, kernel_size=(5,5), strides=(1,1),
               radius_min=2, radius_max=3, initializer='uniform',
               use_bias=True): #, output_dim):
        self.output_dim = output_dim
        self.kernel_size = kernel_size
        self.radius_min = radius_min
        self.radius_max = radius_max
        self.strides = strides
        self.use_bias = use_bias
        xx = np.arange(-(kernel_size[0]-1)//2, (kernel_size[0]+1)//2)
        yy = np.arange(-(kernel_size[0]-1)//2, (kernel_size[0]+1)//2)
        [XX, YY] = np.meshgrid(xx, yy)
        R = np.sqrt(XX**2 + YY**2)
        R[R<radius_min] = 0
        R[R>radius_max] = 0
        R[R>0] = 1
        self.mask = R  # the mask defines the non-zero pattern and will multiply the weights
        if initializer == 'uniform':
            self.initializer = RandomUniform(minval=0, maxval=2/R.sum())
        else:
            self.initializer = initializer
        super(Masked_Conv2D, self).__init__()

    def build(self, input_shape):
        n_filters = input_shape[-1] 
        self.h = self.add_weight(name='h',
                                 shape= self.kernel_size + (n_filters, self.output_dim,),
                                 initializer=self.initializer,
                                 constraint=MaskedConstraint(self.mask), # MaskedConstraint()(self.mask)
                                 trainable=True)

        self.b = self.add_weight(name='b',
                                 shape=(self.output_dim,),
                                 initializer=Constant(0),
                                 trainable=self.use_bias) 
        super(Masked_Conv2D, self).build(input_shape)

    def call(self, x):
        y = ops.conv(x, self.h, strides=self.strides, padding='same')
        if self.use_bias:
            y = y + torch.unsqueeze(self.b, dim=0)
        return y

    def compute_output_shape(self, input_shape):
        return input_shape[:3] + (self.output_dim,)

# def masked_constraint(R):
    # Enforces constraint for kernel to be non-negative everywhere and zero outside the ring

    # Args:ad
    #    R: np.array
    #        Binary mask that extracts 

    # Returns:
    #     my_constraint: function
    #        Function that enforces the constraint
    # breakpoint()
#     R = torch.tensor(R).float()
#    R_exp = torch.unsqueeze(torch.unsqueeze(R, dim=-1), dim=-1)
#    def my_constraints(x):
#        Rt = torch.tile(R_exp, [1, 1, 1, x.shape[-1]])
#        Z = torch.zeros_like(x)
#        return torch.where(Rt > 0, x, Z)
#    return my_constraints

class MaskedConstraint(keras.constraints.Constraint):
    def __call__(self, R):
        # breakpoint() 
        R = torch.tensor(R).float()
        R_exp = torch.unsqueeze(torch.unsqueeze(R, dim=-1), dim=-1)
        def my_constraints(x):
            Rt = torch.tile(R_exp, [1, 1, 1, x.shape[-1]])
            Z = torch.zeros_like(x)
            return torch.where(Rt > 0, x, Z)
        return my_constraints

Here is the error produced by my terminal for both functions:

With this, what are the other work around for this. (Note, x in the code is defined as follows):

x_in = Input(shape=shape) radius_min = int(gSig*r_factor) radius_max = radius_min + width ks = 2*radius_max + 1 conv1 = Masked_Conv2D(output_dim=n_channels, kernel_size=(ks, ks), radius_min=radius_min, radius_max=radius_max, initializer=initializer, use_bias=use_bias) x = conv1(x_in)

[hertschuh]

Hi @mannypaeza ,

The issue with the MaskedConstraint class as it is defined is that __call__ returns a function. Instead, you should define an __init__ to save the mask, and then __call__ should take a value and return a value.

class MaskedConstraint(tf.keras.constraints.Constraint):
    def __init__(self, R):
        R = tf.cast(R, dtype=tf.float32)
        self.R_exp = tf.expand_dims(tf.expand_dims(R, -1), -1)

    def __call__(self, x):
        Rt = tf.tile(self.R_exp, [1, 1, 1, x.shape[-1]])
        Z = tf.zeros_like(x)
        return tf.where(Rt>0, x, Z)

Now, it appears that we could easily allow all callables, not just Constraint subclasses. I'll look into this.

[mannypaeza]

Hello! I ended up forgetting that I wrote this issue, and I was able to figure it out on my own (in the same way you had it, except in pytorch). thanks for the help, anyways!

[google-ml-butler]

Are you satisfied with the resolution of your issue?
Yes
No

[ghsanti]

@mannypaeza it's a beautiful code, are these things called ring convolutions? Can't find any paper or good source about it.

You are selecting the yellow area of the kernel right? Has this any specific use cases?
I'm only slightly suspicious of maxval=2/R.sum() since R can vary greatly..and the yy use kernel[0] not kernel[1] ?

also ig you are using torch but using keras.ops it'be multibackend

[mannypaeza]

@ghsanti yes, this is a ring convolutional model. For use case, it can be used for analyzing microendoscopic
one photon data (neuroscience related). As for the R, it is limited due to the radius of the average neuron (if that makes sense). For kernel_size, we kept it as the same (kernel_size=(a, a)) hence why we did kernel[0] for the yy as it would not matter.