Running the following on CPU is slower but not that bad:
In particular, the fold is “almost free” as it does not do any copy.
And the two other ops have a similar cost as they go through the big unfolded Tensor once each.
import torch
from torch import nn
from torch.nn import Parameter
import time

class MyConv2d1(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1, sratio=0):
        super(MyConv2d1, self).__init__()
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.kernel_size = kernel_size
        self.stride = stride
        self.padding = padding
        self.dilation = dilation
        self.sratio = sratio
        # print(self.sratio)
        # CHANGED: Remove cuda
        self.weight = Parameter(torch.rand(out_channels, in_channels, kernel_size, kernel_size))

    def forward(self, input):
        h_in, w_in = input.shape[2:]
        h_out = h_in
        w_out = w_in
        # x: [batchsize ksize num_sliding]
        start = time.time()
        x = torch.nn.functional.unfold(input, kernel_size=self.kernel_size, padding=self.padding)
        print(" Unfold: ", time.time() - start)

        batchsize = input.shape[0]
        ksize = self.in_channels*self.kernel_size*self.kernel_size
        num_sliding = x.shape[2]

        assert x.shape[1] == ksize


        w = self.weight
        start = time.time()
        out_unf = x.transpose(1, 2).matmul(w.view(w.size(0), -1).t()).transpose(1, 2)
        print(" Mm: ", time.time() - start)
        start = time.time()
        out = torch.nn.functional.fold(out_unf, output_size=[h_out, w_out], kernel_size=1, padding=0, dilation=self.dilation, stride=1)
        print(" Fold: ", time.time() - start)

        return out

in_c = 10
out_c = 15
kernel = 3
padding = 1 # Custom Module only works when output is same size as input
inp = torch.rand(2, in_c, 50, 50)

conv = torch.nn.Conv2d(in_c, out_c, kernel, padding=padding)
custom_conv = MyConv2d1(in_c, out_c, kernel, padding=padding)

start = time.time()
out = conv(inp)
print("Regular: ", time.time() - start)

start = time.time()
out = custom_conv(inp)
print("Custom: ", time.time() - start)