Poetry上传一个属于自己的库

前言

其实这是一个拖了很久很久的坑,不知道多少人看过我之前的一篇博客关于torch.fx的使用,在这里面我用torch.fx实现了一些很有趣的功能比如模型可视化.所以当时就有一个想法,把代码封装一下写成一个属于自己的三方库,正好今天有点时间就把这个坑给填上.

这个工具的主要功能很简单,直接指定某个py文件工具会自动寻找文件中所有的nn.Module并进行解析可视化.

开始

关于模型的trace以及算子的解析在之前的博客中已经写的比较清楚了,这里就不过多赘述.今天的主要内容是封装+Poetry上传.

1 封装

将上次的代码直接拿来用,然后写个函数调用一下

def draw(model: torch.nn.Module, inputs: torch.Tensor, save_dir: str = './Save', save_name: str = 'model'):
    graph = model_graph(model, inputs)
    graph.render(outfile=save_dir + '/' + save_name, view=False)

然后解析py文件中所有的nn.Module

def parse_py(py_path: str) -> list[nn.Module]:
    spec = importlib.util.spec_from_file_location('module_name', py_path)
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    module_classes = [m for _, m in module.__dict__.items() if
                      isinstance(m, (nn.Module, type)) or issubclass(type(m),nn.Module)]
    if len(module_classes) == 0:
        raise ValueError('No nn.Module class found in the file.')
    return module_classes
    
for index, model_name in enumerate(model_list):
    try:
        if not isinstance(model_name, nn.Module):
            model = model_name()
        else:
            model = model_name
        args.name=type(model).__name__+'.svg'
        draw(model, inputs, save_dir=args.dir, save_name=args.name)
    except:
        print(f"{model_name} draw failed")
        pass

这里为了防止某些子类算子不能实现所有输入需求,其实也是目前设计的不太灵活,所以直接用try简化这部分处理.

2 测试一下

输入的测试文件如下test.py

import torch
from torch import nn
from torchvision.models import resnet18,regnet_x_8gf
​
​
class TestConv2d(nn.Module):
    def __init__(self, in_channels, out_channels, **kwargs):
        super(TestConv2d, self).__init__()
        self.conv = nn.Conv2d(in_channels, out_channels, **kwargs)
        self.bn = nn.BatchNorm2d(out_channels)
        self.relu = nn.ReLU(True)
​
    def forward(self, x):
        x = self.conv(x)
        x = self.bn(x)
        x = self.relu(x)
        return x
​
​
class TestModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = TestConv2d(3, 32, kernel_size=3)
        self.conv2 = TestConv2d(32, 64, kernel_size=3)
        self.dropout = nn.Dropout(0.3)
​
    def forward(self, x):
        x = self.conv1(x)
        x = self.conv2(x)
        x = self.dropout(x)
        return x
​
​
model= resnet18()
model2=regnet_x_8gf()

貌似还可以,接下来就是Poetry打包上传了

3 打包&&上传

首先下载Poetry,根据官网推荐利用curl -sSL https://install.python-poetry.org | python3 -进行安装,命令细节根据自己情况更改

使用poetry new xxx新建一个项目,这里项目的结构已经生成好了,只需要把之前的文件复制到对应位置.接下来就是比较重要的一步了,修改.toml文件

这里对这本地环境pip list把包的相关依赖加进去,关于^与~的区别可以自行google一下版本命名规则相关的.接下来就去pypi上注册一个帐号,得到用户名和密码

最后进行上传

poetry publish --build -u username -p password

这里我没有去细讲关于poetry lock/poetry build/poetry show --tree等相关的指令,大家有兴趣可以去看相关内容.

经过一系列操作终于在pypi上看到了我们自己的库,仅仅写了简单的readme没有去写license.

4 下载测试

我们从pypi上把这个库给install下来,并且在本地重新写个测试脚本调用库进行测试

from pytorch_show import main
​
if __name__ == '__main__':
    main()

残暴如此,直接调用main就好,再来写个model.py

from torchvision.models import vgg16
​
model= vgg16()

直接调用测试python test.py -f model.py

直筒形VGG一切正常

总结

今天也算是小小填了一下去年留下的坑,不过还是有很多遗憾的地方.本来这个工具的初衷除了能解析py文件之外还能对onnx进行绘制,但是今天实践下来有很多坑.一开始的想法是将onnx反转回pytorch模型,但是torch.fx的trace并不能支持动态流,因此很多算子包含if或者for loop的地方都会报错.然后尝试了直接利用onnx.tools.net_drawer进行绘制,成功得到了图片但是太过于复杂,很多多余的输入都被展示出来严重干扰了核心算子的展示,同时必须借助运行时才能得到每一步的shape,所以这部分想优化还是要好好想想办法,这也算是为下一阶段改进再留一个坑吧.

另外留个彩蛋,下一期会讲讲最近看到Rust与DL相关的内容.