torch-liberator 0.2.1

Torch Liberator - 部署PyTorch模型

Torch Liberator是一个包含一组用于读取和写入深度网络相关部分的Python模块。

通常，当移动使用torch训练的深度网络时，除了包含学习权重的检查点文件外，还需要跟踪定义模型文件的整个代码库。Torch Liberator包含提取相关源代码并将其与权重捆绑在一起，然后将其序列化成一个单文件部署的工具。

注意：自torch 1.9版本起，torch附带了一个torch.package子模块，其中包含一个保存模型结构和模型权重的函数。我们建议使用torch.package来代替本包中提供的单文件部署。因此，load_partial_state逻辑是本模块中提供的核心代码。

安装

pip install torch_liberator

# OR with a specific branch

pip install git+https://gitlab.kitware.com/computer-vision/torch_liberator.git@main

部分状态加载

在0.1.0版本的torch liberator中，现在公开了一个load_partial_state函数，该函数尽可能地将一个模型的权重“推入”另一个模型。有几种方法可以计算一个模型中层名称与另一个模型中层名称之间的关联，最通用的是“嵌入”方法，而稍微更结构化的“同构”选项。

你有没有遇到过这种情况，即使用一个模型作为更大网络中的子模型？然后你必须将预训练的子网络状态加载到更大的模型中？

torch_liberator.load_patial_state的最新版本可以通过解决最大公共子树同构问题来处理这种情况。这计算了两个具有一致后缀的状态字典之间可能的最大映射。

>>> import torchvision
>>> import torch
>>> import torch_liberator
>>> resnet50 = torchvision.models.resnet50()
>>> class CustomModel(torch.nn.Module):
>>>     def __init__(self):
>>>         super().__init__()
>>>         self.module = resnet50
>>>         self.extra = torch.nn.Linear(1, 1)
>>> # Directly load resnet50 state into a model that has it as an embedded subnetwork
>>> model = CustomModel()
>>> model_state_dict = resnet50.state_dict()
>>> # load partial state returns information about what it did
>>> info = torch_liberator.load_partial_state(model, model_state_dict, association='isomorphism', verbose=1)
>>> print(len(info['seen']['full_add']))
>>> print(len(info['self_unset']))
>>> print(len(info['other_unused']))
320
2
0

它还可以处理加载具有某些底层结构共享的两个模型之间的公共状态。

>>> import torchvision
>>> import torch
>>> import torch_liberator
>>> resnet50 = torchvision.models.resnet50()
>>> class CustomModel1(torch.nn.Module):
>>>     def __init__(self):
>>>         super().__init__()
>>>         self.module = resnet50
>>>         self.custom_model1_layer = torch.nn.Linear(1, 1)
>>> class CustomModel2(torch.nn.Module):
>>>     def __init__(self):
>>>         super().__init__()
>>>         self.backbone = resnet50
>>>         self.custom_model2_layer = torch.nn.Linear(1, 1)
>>> # Load as much of model1 state into model2 as possible
>>> model1 = CustomModel1()
>>> model2 = CustomModel2()
>>> model2_state_dict = model2.state_dict()
>>> # load partial state returns information about what it did
>>> info = torch_liberator.load_partial_state(model1, model2_state_dict, association='isomorphism', verbose=1)
>>> print(len(info['seen']['full_add']))
>>> print(len(info['seen']['skipped']))
>>> print(len(info['self_unset']))
>>> print(len(info['other_unused']))
320
2
2
2

>>> import torchvision
>>> import torch_liberator
>>> #
>>> faster_rcnn = torchvision.models.detection.faster_rcnn.fasterrcnn_resnet50_fpn()
>>> resnet50 = torchvision.models.resnet50(pretrained=True)
>>> state_dict = resnet50.state_dict()
>>> # Load partial state return a dictionary that tells you how well it did
>>> info = torch_liberator.load_partial_state(faster_rcnn, state_dict, verbose=0, association='embedding')
>>> print(ub.map_vals(len, info['seen']))
>>> print(ub.map_vals(len, ub.dict_diff(info, ['seen'])))
{'full_add': 265, 'skipped': 55}
{'other_unused': 55, 'self_unset': 30}

>>> # Load partial state return a dictionary that tells you how well it did
>>> info = torch_liberator.load_partial_state(faster_rcnn, state_dict, verbose=0, association='isomorphism')
>>> print(ub.map_vals(len, info['seen']))
>>> print(ub.map_vals(len, ub.dict_diff(info, ['seen'])))
{'full_add': 265, 'skipped': 55}
{'other_unused': 55, 'self_unset': 30}

此外，如果张量的尺寸不完全匹配，它们将被扭曲，即尽可能“推入”。有关更多详细信息，请参阅文档字符串。

独立单文件模型部署

torch_liberator的原始目的是构建包含模型代码和模型权重的独立torch包。它仍然这样做，但torch 1.9中引入的torch.package可能是一个更好的解决方案。有关详细信息，请参阅torch.package。

Torch Liberator基于liberator库，静态提取定义模型拓扑的pytorch代码，并将其与预训练权重文件捆绑在一起。这结果是一个单文件部署包，并可能消除对用于训练模型的代码库的依赖。

Torch Liberator还可以读取这些部署文件，并创建一个使用正确的预训练权重初始化的模型实例。

API尚可，但确实需要改进。然而，当前版本处于工作状态。没有高级文档，但有大量的文档字符串和doctests。此处的示例通过从torchvision中提取AlexNet模型，很好地概述了代码。

>>> import torch_liberator
>>> from torch_liberator.deployer import DeployedModel
>>> from torchvision import models

>>> print('--- DEFINE A MODEL ---')
>>> model = models.alexnet(pretrained=False)  # false for test speed
>>> initkw = dict(num_classes=1000)  # not all models nicely supply this
>>> model._initkw = initkw
--- DEFINE A MODEL ---

>>> print('--- DEPLOY THE MODEL ---')
>>> zip_fpath = torch_liberator.deploy(model, 'test-deploy.zip')
--- DEPLOY THE MODEL ---
[DEPLOYER] Deployed zipfpath=/tmp/tmpeqd3y_rx/test-deploy.zip


>>> print('--- LOAD THE DEPLOYED MODEL ---')
>>> loader = DeployedModel(zip_fpath)
>>> model = loader.load_model()
--- LOAD THE DEPLOYED MODEL ---
Loading data onto None from <zopen(<_io.BufferedReader name='/tmp/tmpg1kln3kw/test-deploy/deploy_snapshot.pt'> mode=rb)>
Pretrained weights are a perfect fit

目前的主要问题是，你必须在部署时明确定义“initkw”（用于创建我们的模型实例的关键字参数），或者你可以将其设置为模型（或如果所有关键字参数都是类的成员变量，torch_liberator会尝试智能地推断initkw应该是什么）的_initkw属性。

还有一个torch-liberator CLI，可以用来打包权重文件、Python模型文件和可选的JSON元数据。

python -m torch_liberator \
    --model <path-to-the-liberated-py-file> \
    --weights <path-to-the-torch-pth-weight-file> \
    --info <path-to-train-info-json-file> \
    --dst my_custom_deployfile.zip