PyTorch element-wise filter layer

Hi, I want to add element-wise multiplication layer to duplicate the input to multi-channels like this figure. (So, the input size M x N and multiplication filter size M x N is same), as illustrated in this figure

I want to add custom initialization value to filter, and also want them to get gradient while training. However, I can't find element-wise filter layer in PyTorch. Can I make it? Or is it just impossible in PyTorch?

Solution 1:

In pytorch you can always implement your own layers, by making them subclasses of nn.Module. You can also have trainable parameters in your layer, by using nn.Parameter.
Possible implementation of such layer might look like

import torch
from torch import nn

class TrainableEltwiseLayer(nn.Module)
  def __init__(self, n, h, w):
    super(TrainableEltwiseLayer, self).__init__()
    self.weights = nn.Parameter(torch.Tensor(1, n, h, w))  # define the trainable parameter

  def forward(self, x):
    # assuming x is of size b-1-h-w
    return x * self.weights  # element-wise multiplication

You still need to worry about initializing the weights. look into nn.init on ways to init weights. Usually one init the weights of all the net prior to training and prior to loading any stored model (so partially trained models can override random init). Something like

model = mymodel(*args, **kwargs)  # instantiate a model
for m in model.modules():
  if isinstance(m, nn.Conv2d):
     nn.init.normal_(m.weights.data)  # init for conv layers
  if isinstance(m, TrainableEltwiseLayer):
     nn.init.constant_(m.weights.data, 1)  # init your weights here...