欢迎大家来到IT世界,在知识的湖畔探索吧!

您可以用相同的方式转换各种模型。

要将机器学习模型转换为 CoreML 以在 iOS 上使用，您之前必须从模型存储库初始化模型架构并加载权重。

Spandrel 是一个可以从 pytorch 检查点文件（pth）恢复模型结构的库，并支持包括超分辨率和修复在内的多种模型。所有的你
需要的是权重（检查点文件）。

将它们转换为 CoreML 并在 iOS 和 macOS 设备上使用它们会很方便，因此我们将转换它们。

这次，我们尝试使用 RetinexFormer 来增强低光图像。

示例：RetinexFormer 转换程序

安装

使用 pip 安装。

pip install spandrel spandrel_extra_arches # git clone https://github.com/chaiNNer-org/spandrel.git # import shutil # shutil.copytree("spandrel/libs/spandrel/spandrel/","/usr/local/lib/python3.10/dist-packages/spandrel")

欢迎大家来到IT世界,在知识的湖畔探索吧!

基本上，您可以只使用 pip，但由于某种原因，当您使用 pip 安装时，某些软件包不包含模型架构，因此如果稍后初始化模型时遇到问题，请克隆存储库并将其替换为软件包。

初始化模型

将 pth 文件加载到 ModelLoader 中并初始化模型。
通常，这就是您需要做的全部事情，但是如果您没有商业许可证，则需要从 spandrel_extra_arches 注册模型。 Codeformer 也是如此。

欢迎大家来到IT世界,在知识的湖畔探索吧!from spandrel import MAIN_REGISTRY, ModelLoader, ImageModelDescriptor from spandrel_extra_arches import EXTRA_REGISTRY # add extra architectures before `ModelLoader` is used MAIN_REGISTRY.add(*EXTRA_REGISTRY) # load a model from disk model = ModelLoader().load_from_file(r"LOL_v2_real.pth") # make sure it's an image to image model assert isinstance(model, ImageModelDescriptor) # send it to the GPU and put it in inference mode model.cuda().eval()

我们用火炬来检查一下模型的运行情况。
推理方法见下文。

制作一个包裹模型

创建一个 nn.module，向模型添加后处理。

import torch torch_model = model.model.eval().cpu() class CoreMLModel(torch.nn.Module): def __init__(self, m): super(CoreMLModel, self).__init__() self.m = m def forward(self, image): pred = self.m(image) output = torch.clamp(pred * 255, min=0, max=255) return output coremlmodel = CoreMLModel(torch_model)

转换您的包装模型。

欢迎大家来到IT世界,在知识的湖畔探索吧!import coremltools as ct import torch ex = torch.randn(1,3,512,512).cpu() ts = torch.jit.trace(coremlmodel, ex) mlmodel = ct.convert(ts, inputs=[ct.ImageType(shape=ex.shape,scale=1/255)],outputs=[ct.ImageType(name="output")]) mlmodel.save("retiinexformerNTIRE.mlpackage")