Hi! I am struggling to understand why the code below introduces any overhead (especially noticeable in backward) compared to torch.nn.Linear. What is suboptimal in my implementation and how to improve it?
(for some reasons I really need to wrap torch.nn.functional.linear in MyLinearFunction and torch.nn.Linear in MyLinear)
You can run the code in Colab:
https://colab.research.google.com/drive/1_ab1XLMfH3Zkwgoas4CChjnZYUp6LAy2 3
If you have any problems with the link above, here is the same code:
import time

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader
from torchvision.datasets import MNIST
from torchvision.transforms import Compose, ToTensor


def bench(module, n=1000):
    forward_total = 0
    backward_total = 0
    for _ in range(n):
        x, y = next(iter(dataloader))
        x = x.view(-1, 784)
        now = time.perf_counter()
        out = module(x)
        forward_total += time.perf_counter() - now
        loss = loss_fn(out, y)
        now = time.perf_counter()
        loss.backward()
        backward_total += time.perf_counter() - now
    return forward_total / n, backward_total / n


class MyLinearFunction(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x, weight, bias=None):
        ctx.x = x
        ctx.weight = weight
        ctx.bias = bias
        return F.linear(x, weight, bias)

    @staticmethod
    def backward(ctx, grad_out):
        x, weight, bias = ctx.x, ctx.weight, ctx.bias
        return grad_out @ weight, grad_out.t() @ x, grad_out.sum(dim=0)


class MyLinear(nn.Linear):
    def forward(self, x):
        return MyLinearFunction.apply(x, self.weight, self.bias)


mnist_path = 'mnist_data'
mnist = MNIST(
    mnist_path,
    train=True,
    transform=Compose([
        ToTensor(),
        lambda x: x[0]  # mnist images contain one channel
    ]),
    download=True
)
dataloader = DataLoader(mnist, batch_size=128)
loss_fn = nn.CrossEntropyLoss()

m1 = nn.Linear(784, 10)
m2 = MyLinear(784, 10)
m1.weight.data = m2.weight.data.clone()
m1.bias.data = m2.bias.data.clone()

vf, vb = bench(m1)
cf, cb = bench(m2)
print(f'vanilla forward: {vf:.6f}, vanilla backward: {vb:.6f}')
print(f'custom  forward: {cf:.6f}, custom  backward: {cb:.6f}')