Hello
I’ve been trying to build a custom channel-wise convolution operation with multiplication and sum operation so that I can alter the bit-precision during the convolution process.
(I’ve checked several forum articles and all of them told me I should build it with unfold and fold functions.)
But when I replace the torch.nn.conv2d function with my code
There is no error but accuracy seems to degrade a lot
I’m completely new to python and pytorch, and only know basics about neural network so I can’t find what mistakes I made. 
Can anyone please help? I would really appreciate any help.
Also, if there is any tips to save gpu usage I would also appreciate it very much, since my code seems to consume tons of memory so I had to cut the batch size and it’s taking forever to see the results. 
class Conv2DFunctionCustom(Function):
@staticmethod
def forward(ctx, input, weight, bias=None, stride=1, padding=0, dilation=1, groups=1):
ctx.save_for_backward(input, weight, bias)
ctx.stride, ctx.padding, ctx.dilation, ctx.groups = stride, padding, dilation, groups
zeropad = nn.ZeroPad2d(ctx.padding[0])
batch_size = len(input)
input_channel = len(input[0])
input_size = len(input[0][0])
ochannel = len(weight)
ichannel = len(weight[0])
kernel_size = len(weight[0][0])
if padding:
inp_pad = zeropad(input)
else:
inp_pad = input
input_unfold = torch.nn.functional.unfold(inp_pad, kernel_size).transpose(1, 2)
if input_channel < 256:
input_unfold = input_unfold.unfold(2,input_channel,input_channel)[None,:]
else:
input_unfold = input_unfold.unfold(2,256,256)[None,:]
weight_unfold = weight.view(weight.size(0), -1)
if input_channel < 256:
weight_unfold = weight_unfold.unfold(1,ichannel,ichannel)[:,None,None]
else:
weight_unfold = weight_unfold.unfold(1,256,256)[:,None,None]
mac = input_unfold * weight_unfold
mac = torch.sum(mac, dim=4)
mac = torch.sum(mac, dim=3)
out = torch.nn.functional.fold(mac, (input_size, input_size), (1, 1)).transpose(0,1)
return out
def backward(ctx, grad_output):
input, weight, bias = ctx.saved_tensors
stride, padding, dilation, groups = ctx.stride, ctx.padding, ctx.dilation, ctx.groups
grad_input = grad_weight = grad_weight = grad_input = grad_bias = grad_stride = grad_padding = grad_dilation = grad_groups = None
if ctx.needs_input_grad[1]:
grad_input = torch.nn.grad.conv2d_input(input.shape, weight, grad_output, stride, padding, dilation, groups)
if ctx.needs_input_grad[2]:
grad_weight = torch.nn.grad.conv2d_weight(input, weight.shape, grad_output, stride, padding, dilation, groups)
return grad_input, grad_weight, grad_input, grad_weight, grad_bias, grad_stride, grad_padding, grad_dilation, grad_groups