Hello,
I tried to visualize the runtime performance improvement made by convolution layer which I implemented from scratch Vs torchscript version of convolution layer Vs torch.nn.conv2d() module for 100 iterations with input (128,3,28,28), out_channel =64, kernel size=3.
Convolution layer from scratch in CUDA -> 9.366 seconds
torchscript convolution layer from scratch in CUDA -> 6.636 seconds
torch.nn.conv2d() -> 475.614 milliseconds.
Is there any problem in my approach? and how to optimize even more? . I request you to help me with this problem.
My code
class conv2D(nn.Module):
  def __init__(self, in_channel, out_channel, kernel_size):
    super(conv2D,self).__init__()
    self.weight = torch.nn.Parameter(torch.ones(out_channel,in_channel,kernel_size, kernel_size))
    self.bias = torch.nn.Parameter(torch.zeros(out_channel))
    self.kernel_size = kernel_size
    self.in_channel = in_channel
    self.out_channel = out_channel

  def forward(self, image):
    img_height = image.shape[3]
    img_width = image.shape[2]
    batch_size = image.shape[0]
    out_height = img_height-self.kernel_size+1
    out_width = img_width-self.kernel_size+1

    output = torch.zeros(batch_size,self.out_channel,out_width,out_height)
    for k in range(batch_size):
      for i in range(out_height):
        for j in range(out_width):
          temp = torch.sum(image[k,:,j:j+self.kernel_size,i:i+self.kernel_size]*self.weight,dim=(1,2,3))
          output[k,:,i,j]=torch.add(temp,self.bias)
      return output

x = torch.ones(128,3,28,28).to("cuda")
c = conv2D(3,64,3).to("cuda")
c_s = torch.jit.script(c).to("cuda")
c_s(x)

Scripting the model and running with a sample input to get an optimized graph
with torch.autograd.profiler.profile(use_cuda=True) as prof:
  with torch.no_grad():
    for i in range(100):
      c(x)
print(prof.table())

Profiling both the scripted and normal method.
with torch.autograd.profiler.profile(use_cuda=True) as prof:
  with torch.no_grad():
    for i in range(100):
      c_s(x)
print(prof.table())