I have a math function that generates a custom kernel for use in the convolution layer.
This function has 3 parameters ‘a’, ‘b’ and ‘c’, as follow:
kernel = generator(a, b, c)
I would like to build a CNN model that, instead trainingg the kernel matrix values, train just the parameters ‘a’, ‘b’ and ‘c’ that generates the kernel.
How can I do that in PyTorch?
Example:
class MyModel(nn.Module):
    def __init__(self, kernel):
        super(MyModel, self).__init__()
        self.conv1 = nn.Conv2d(1, 1, kernel_size=7, bias=False)
        # Other layers here

        # Generate the custom kernel
        kernel = kernel_generator(a, b, c)

        # Question: how train a, b, c
        # in the gradient descent???

        # Initialize conv1 with custom kernel
        self.conv1.weight = nn.Parameter(kernel)

    def forward(self, x):
        x = self.conv1(x)
        # pass x to other modules
        return x