Hello everyone! I’ve been working on a customized Conv2d layer which allows me to do operations on the data immediately after every multiplication and sum (I’m trying to prove that we can quantize a CNN to use less bits while keeping the accuracy, it’s for my masters). My starting point is the code available here, where unfold and fold is used for that.
The code is:
# Vanilla Convolution
def myconv2d(input, weight, bias=None, stride=1, padding=0, dilation=1, groups=1):
    batch_size, in_channels, in_h, in_w = input.shape
    out_channels, in_channels, kh, kw =  weight.shape

    unfold = torch.nn.Unfold(kernel_size=(kh, kw), dilation=dilation, padding=padding, stride=stride)
    inp_unf = unfold(input)

    if bias is None:
        out_unf = inp_unf.transpose(1, 2).matmul(weight.view(weight.size(0), -1).t()).transpose(1, 2)
    else:
        out_unf = (inp_unf.transpose(1, 2).matmul(w_) + bias).transpose(1, 2)
    out = out_unf.view(batch_size, out_channels, out_h, out_w)
    return out

The problem is that I need to do call a function immediately after every multiplication and sum, so I could not use the .matmul() method since it does both for me. My first attempt was to do the operations independently, calling myfunc() after each one of them:
# Vanilla Convolution
def myconv2d(input, weight, bias=None, stride=1, padding=0, dilation=1, groups=1):
    batch_size, in_channels, in_h, in_w = input.shape
    out_channels, in_channels, kh, kw =  weight.shape

    unfold = torch.nn.Unfold(kernel_size=(kh, kw), dilation=dilation, padding=padding, stride=stride)
    inp_unf = unfold(input)

    weight_t = weight.view(weight.size(0), -1).t()
    inp_unf_t = inp_unf.transpose(1, 2)
    inp_unf_t = inp_unf_t[:,:,:, None]
    inp_unf_t_exp = inp_unf_t.transpose(2,3).repeat((1,1,weight_t.t().size(0),1))
    weight_t_exp = weight_t.t().repeat((1,inp_unf_t.size(1),1,1))
    out_unf = myfunc(inp_unf_t_exp.mul(weight_t_exp))
    out_unf = myfunc(torch.sum(out_unf, dim=3).transpose(1, 2))
    if bias is not None:
        out_unf = out_unf + myfunc(bias.view(-1, 1))

    out = out_unf.view(batch_size, out_channels, out_h, out_w)
    return out

This approach however is terribly slower and (B) consumes a  lot of memory, to the point that a “RuntimeError: CUDA out of memory” error is thrown. Is there an alternative to that approch which would be less memory greedy?
Any help is much appreciated!