autonvtx 0.1

https://github.com/zasdfgbnm/autonvtx

Install

pip install autonvtx

Usage

Write your model as usual and autonvtx(model) to your model:

import torch
import autonvtx

class Model(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.layer1 = torch.nn.Linear(5, 5)
        self.layer2 = torch.nn.Linear(5, 5)
    def forward(self, x):
        x = self.layer1(x)
        x = self.layer2(x)
        return x

m = Model().cuda()
autonvtx(m)
input_ = torch.randn(1024, 5, device='cuda')

torch.cuda.profiler.start()
for _ in range(10):
    output = m(input_)
torch.cuda.profiler.stop()

The screenshot for this would be:

It also works with existing models:

import torch
import torchvision
import autonvtx

m = torchvision.models.resnet50()
autonvtx(m)
input_ = torch.randn(10, 3, 224, 224)

torch.cuda.profiler.start()
for _ in range(10):
    output = m(input_)
torch.cuda.profiler.stop()

The screenshot for this would be: