How to perform channel wise convolution

Question

I have an input tensor of shape (10,3,4,6). The 6 is the channel dimension. If I need to perform convolution (1D and 2D both) channel-wise ( each channel should have different weights and bias) using Pytorch, how can I implement it?

Please give me an example code

Answer 1

depth-wise convolution is what you need:

import torch
from torch import nn
input = torch.randn(10,3,4,6)
# change the channel dimension as conv in torch need BxCxHxW
input = input.permute(0, 3, 1, 2).contiguous() 
# set the group as the number of channels
conv = nn.Conv2d(6, 6, kernel_size=3, stride=1, padding=1, groups=6)
print(conv(input))