https://github.com/zasdfgbnm/autonvtx
autonvtx的Python项目详细描述
安装
pip install autonvtx
使用
像往常一样编写模型,autonvtx(model)
到模型中:
截图如下:
它也适用于现有模型:
importtorchimporttorchvisionimportautonvtxm=torchvision.models.resnet50()autonvtx(m)input_=torch.randn(10,3,224,224)torch.cuda.profiler.start()for_inrange(10):output=m(input_)torch.cuda.profiler.stop()
截图如下:
- 项目
标签: