https://github.com/zasdfgbnm/autonvtx
Project description
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:
Project details
Release history Release notifications | RSS feed
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.
Source Distribution
Built Distribution
autonvtx-0.1-py3-none-any.whl
(2.7 kB
view hashes)
Close
Hashes for autonvtx-0.1.linux-x86_64.tar.gz
Algorithm | Hash digest | |
---|---|---|
SHA256 | 25ae971a23af60f75a3eaea056434c4b7c5733124a1b2f7c6704a8dd12547e9d |
|
MD5 | 01b8c8b24c93e1fbd282606ae5bbf51b |
|
BLAKE2b-256 | 5db4b176868ee80b2c47c9e4f58ffa26b764d33d74755ffebb55f237c987e23a |