I am doing an experiment on images of small size as follows:
The input image has only one channel, then I apply 2 customed filters (denoted by A and B) to do the convolution with the image. The filter has the same size as the image and each of A and B have 10 channels respectively. So after the convolution we have two 1 * 10 tensors. Denote them as (a1, a2, … a10) and (b1, b2, … b10). Let w = (w1, w2, …, w10) be a trainable 1 *10 tensor. Then I want to compute the following fraction: (the numerator is the norm of weighted sum of vectors (ai, bi) and the denominator is the sum of weighted norm. The norm here is the ordinary Euclidean norm.)
Screen Shot 2020-09-07 at 6.46.40 PM1060×138 4.88 KB
I am wondering how to do this computation in the forward function and where I should claim w as trainable tensor? The template of my code is as follows:
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.filter_A = ...
        self.filter_B = ...

    def forward(self, x):
        x_A = F.conv2d(x, self.filter_A)
        x_B = F.conv2d(x, self.filter_B)
        ...