使用torch.package将pytorch模型进行独立打包

点击上方“Deephub Imba”,关注公众号,好文章不错过 !

PyTorch 发布于 2016 并迅速成为深度学习研究人员的首选工具。随着PyTorch的逐步发展，它已经不仅仅是一个原型工具。现在PyTorch成为一个成熟的框架，并且逐渐成为学术界和工业界的标准。研究人员和机器学习工程师可以在本地 Jupyter 的服务器、云平台多节点 GPU 集群以及边缘智能设备高效运行 PyTorch。

但是在我看来，PyTorch 有一个明显的缺点：它存储模型的方式。官方推荐存储 PyTorch 模型的方法是什么？应该保存模型的参数，而不是模型本身。以下是官方文档的引用：

当你保存模型进行推理时，只需保存训练模型的学习参数即可。使用 torch.save() 函数保存模型的 state_dict 将为以后恢复模型提供最大的灵活性。

这就是官方推荐保存模型方法，但是他的缺点是什么？你需要在使用时再次定义模型，反序列化 state_dict 并将参数加载到模型中。

在最新的1.9版本中有了一个新的方法torch.package，可以帮我们简化上面的步骤。

torch.package

torch.package是一种将PyTorch模型打包成独立格式的新方法。打包后的文件包含模型参数和元数据及模型的结构，换句话说，我们使用时只要load就可以了。

此外，torch.package 添加了对任意 PyTorch 代码的独立打包的支持。这意味着可以使用它打包任何想要的东西（例如，DataLoaders、Datasets 等）。这是一个非常好的方法来重现的训练结果，并为模型服务提供了方便的迁移支持。

torch.package模块是PyTorch 1.9的一部分，所以我们首先需要检查并安装正确的PyTorch版本。使用下面的命令:

!pip uninstall -y torch
!pip install -f https://download.pytorch.org/whl/test/cpu/torch_test.html torch

打包模型

我们使用torch.package来打包和部署一个用于生成时装的DCGAN。这个预训练的模型可以在 Torch Hub 上找到。使用下面的代码进行下载：

import torch

use_gpu = True if torch.cuda.is_available() else False
model = torch.hub.load('facebookresearch/pytorch_GAN_zoo:hub', 'DCGAN', pretrained=True, useGPU=use_gpu

接下来我们可以测试一下使用这个DCGAN生成时装模特的随机图像 ：

def run_model(model):
    num_images = 64
    noise, _ = model.buildNoiseData(num_images)
    with torch.no_grad():
      generated_images = model.test(noise)

    # let's plot these images using torchvision and matplotlib
    import matplotlib.pyplot as plt
    import torchvision
    plt.imshow(torchvision.utils.make_grid(generated_images).permute(1, 2, 0).cpu().numpy())

run_model(model)

结果如下：

现在，让我们打包模型。我们将使用 PackageExporter 来创建一个存档文件，这个存档就包含了在另一台机器上运行模型所需的所有东西：

from torch import package

path = "/tmp/dcgan.pt"
package_name = "dcgan"
resource_name = "model.pkl"

with package.PackageExporter(path) as exp:
    exp.save_pickle(package_name, resource_name, model)

但是，运行上面的代码会产生一个错误：

PackagingError:
* Module did not match against any action pattern. Extern, mock, or intern it.
models.DCGAN
models.utils.config
models.loss_criterions.ac_criterion     models.loss_criterions.base_loss_criterions     models.networks.DCGAN_nets

这是什么意思？为了创建独立包，PackageExporter 将尝试访问将被打包模块和对象所需的所有依赖项的源代码。

在此过程中，它将跳过标记为 extern 或 mock 的依赖项，并将所有标记为 intern 的依赖项包含在存档中。所以我们需要为导入相应的models模块：

with package.PackageExporter(path) as exp:
    exp.intern("models.**")
    exp.save_pickle(package_name, resource_name, model)

现在，我们有另一个问题:

PackagingError:
* Module did not match against any action pattern. Extern, mock, or intern it.
numpy

我们继续将numpy也加入进去

with package.PackageExporter(path) as exp:
    exp.extern("numpy.**")
    exp.intern("models.**")
    exp.save_pickle(package_name, resource_name, model)

现在这个命令没有问题了。

加载模型

我们可以使用PackageImporter要将模型加载到内存中:

imp = package.PackageImporter(path)
loaded_model = imp.load_pickle(package_name, resource_name)

要测试一切是否正常，可以再次调用前面定义的run_model函数并检查结果。

torch.package 除了打包模型文件以外还提供了很多高级功能，如果你想试试torch.package 但是又找不到1.9的环境，看看这个colab，在这里你可以随意测试：

https://colab.research.google.com/drive/1lFZkLyViGfXxB-m3jqlyTQuYToo3XLo-

作者：Dimitris Poulopoulos

喜欢就关注一下吧：

点个 在看 你最好看！********** **********