【硬核】自己动手实现轻量级神经网络推理框架

点击上方“视学算法”，选择加"星标"或“置顶”

重磅干货，第一时间送达

作者丨东林钟声@知乎（已授权）

来源丨https://zhuanlan.zhihu.com/p/347752954

编辑丨极市平台

导读

 

本文介绍了作者编写的 Planer (Powerful Light Artificial NEuRon)框架，该框架仅依赖 NumPy 作为其矩阵计算库，通过正则表达式对 PyTorch 模型 TorchScript 进行解析实现模型自动转换到 Planer 框架进行推理。 作者已经转换了多种 CNN 模型在 Planer 实现成功推理。

介绍：

前段时间投入了大量的时间与精力编写了 Planer (Powerful Light Artificial NEuRon)框架，该框架仅依赖 NumPy 作为其矩阵计算库，并设计了 JSON 格式的及其精简的中间表达格式。最终通过正则表达式对 PyTorch 模型 TorchScript 进行解析实现模型自动转换到 Planer 框架进行推理。笔者已经转换了多种 CNN 模型在 Planer 实现成功推理。同时，笔者希望 Planer 能够成为在对部署难度有要求的计算环境中成为有竞争力的一个框架。Planer 在设计之初就考虑到 了可扩展性以及可移植性，笔者编写的另一个目的是希望大家能够加入完善 Planer 的队伍当中，实现更多的 Layer 并支持更多更新的模型，让 Planer 的生态壮大起来。

https://github.com/Image-Py/planergithub.com

下图是已经成功转换的PyTorch训练好的模型在Planer上实现推理（HED边缘检测、CRAFT场景文字检测、ResNet18、ESRGAN超分辨率、UNet以及YOLO-v3）。Image-Py/planer（https://github.com/Image-Py/planer）下图是已经成功转换的PyTorch训练好的模型在Planer上实现推理（HED边缘检测、CRAFT场景文字检测、ResNet18、ESRGAN超分辨率、UNet以及YOLO-v3）。

特点：

• 纯NumPy实现，依赖简单、部署快捷

• 非常精简的IR实现，基于JSON

• 自带模型可视化（基于networkx）

• 支持模型从PyTorch自动转换

• 比较丰富的示例CNN模型

顶层设计：

我们要设计一个神经网络推理框架，首先要先把框架的顶层设计想好。我们的目的是实现一个部署友好、自主可控且轻量级的推理框架。部署友好如何实现?因为推理框架的本质还是涉及到数据的计算，我们采用NumPy来支撑我们所有的计算。同时为了做到轻量级，我们将所有神经网络中涉及到比较单独的部分统称为Layer，比如卷积层、全连接层以及非线性激活函数等。同时这种设计方式使得以后对框架进行扩充只需要再实现新的Layer功能就好。这里我们给出我们Layer这个基本类的Python抽象：

class Layer:    name = 'layer'
    def __init__(self, name):        self.name = name
    def forward(self, x): pass
    def backward(self, grad_y): pass
    def para(self): return None
    def load(self, buf): return 0
    def __call__(self, x):        return self.forward(x)

Layer中最重要的两个api就是前向计算 forward 以及参数加载load。这个设计基本上可以覆盖大部分的模型中所涉及到的层与操作。同时我们这种统一的Layer 设计，可以统一有训练参数层操作(全连接、卷积层等)与无训练参数操作(激活函数、池化等)。无训练参数操作只需要实现forward即可。为了保证模型在部署加载中的简便性，我们将所有的权重都拉直后拼接成一个一维npy文件保存起来，通过NumPy 的io功能实现模型的保存于加载。同时我们将模型的计算表示，也就是IR用json文件保存。这样在实际部署推理模型的时候，只需要NumPy模块与一个npy和一个json文件。这样的设计极大程度的较小了部署难度，将所有的业务实现交给NumPy来完成。同时由于PyTorch更加直观且高效，在研究中广泛使用。Planer框架主要将PyTorch作为对照，计算风格与api设计成与之类似。

构建方式：

Planer支持两种构建方式，一种是手动构建，基本和PyTorch一样的实现，另外一种是基于json文件自动生成模型。

手动构建：

from planer import *# ========== write a net manually ========== class CustomNet(Net):    def __init__(self):        self.conv = Conv2d(3, 64, 3, 1)        self.relu = ReLU()        self.pool = Maxpool(2)        self.upsample = UpSample(2)        self.concatenate = Concatenate()        self.sigmoid = Sigmoid()
    def forward(self, x):        x = self.conv(x)        x = self.relu(x)        y = self.pool(x)        y = self.upsample(y)        z = self.concatenate([x, y])        return self.sigmoid(z)

JSON构建：

# ========== load net from json ========== layer = [('conv', 'conv', (3, 64, 3, 1)),        ('relu', 'relu', None),        ('pool', 'maxpool', (2,)),        ('up', 'upsample', (2,)),        ('concat', 'concat', None),        ('sigmoid', 'sigmoid', None)]
flow = [('x', ['conv', 'relu'], 'x'),        ('x', ['pool', 'up'], 'y'),        (['x','y'], ['concat', 'sigmoid'], 'z')]
net = Net()net.load_json(layer, flow)

模型自动转换（ResNet18、pytorch 1.1.0）：

from torchvision.models import resnet18import torchfrom planer import torch2planer
net = resnet18(pretrained=True)x = torch.randn(1, 3, 224, 224, device='cpu')torch2planer(net, 'resnet18', x)
# then you will get a resnet18.json and resnet18.npy in current folder.
from planer import read_netimport planerimport numpy as np
# get the planer array libpal = planer.core(np)x = pal.random.randn(1, 3, 224, 224).astype('float32')net = read_net('resnet18')net(x) # use the net to predict youre data

GPU加速：

这里可以直接使用CuPy替换掉NumPy作为backend来进行gpu计算加速

import planer, cupyplaner.core(cupy) # use cupy as backend
import planer, clpyplaner.core(clpy) # use clpy as backend

模型可视化（UNet）：

已经支持的Layer：

FC、Conv、Flatten、Upsample、MaxPool、BachNorm、ReLU、Sigmoid、LeakyReLU、Softmax等。其实再有了这些基本的操作后，可以成功推理复杂的CNN，比如YOLO-v3：

YOLO-v3的成功转换代表了一个里程碑，表示一些基本的模型都可以通过Planer来进行推理实现，而且Planer推理只依赖NumPy，这使得使用Planer可以大幅度降低部署难度与门槛。同时我正在写一本详细讲解Planer的电子书《动手编写深度学习推理框架 Planer》，第一个版本的电子书已经编写完毕，详情可以私聊。

点个在看 paper不断！