批评你不可怕，对你失望才可怕。 —— 《看见》，柴静

LeZero是Le（法语，代词）+Zéro（法语，零）的合成词，意为“零的起点”。零是一切数字的原点，这个项目也是深度学习的入门的第一步。

项目源码：github.com/LeiLiGithub…

通过阅读本文，你可以了解到：

1. 一个极简的深度学习框架，应该包含哪些模块
2. 如何实现可以灵活配置的数据源、模型、优化算法
3. 怎样把训练好的模型集成入Android项目，并且实现Python与Android互相调用
4. 一个完整的基于MNIST手写数字识别Demo

背景

一方面，大语言模型（LLMs）是当下最炙手可热的话题之一，虽然目前大部分模型都是以云端服务的方式提供（ChatGPT、文心一言、星火、豆包、通义千问），也有很多公司团队在探索边缘计算的课题。由于隐私性、数据安全、响应速度需求的增加，端侧AI势必将在未来的科技浪潮中占有一席之地。

另一方面，作为拥有十余年经验的移动端研发人员，自己也努力在当下寻求一些突破，端侧智能是很好的一个点，能够利用起自己积累的终端开发丰富经验。

因此，未来我将长期投入到端侧AI的知识图谱建立与技能积累，系统化提升自己在这方面的竞争力。

框架设计目标

结合背景，以及自己查阅到的深度学习相关资料，为这个框架列出如下设计目标：

1. 反映深度学习/神经网络核心思想。作为一个原教旨主义者，该框架应当反映出深度学习最本质的理论形态，从感知器（神经元）开始，扩展为不同的激活层，多个激活层交织构成最终模型
2. 项目代码精简结构清晰。采用易于学习、编写、调试的Python语言进行开发，最大程度地减少逻辑无关的模板代码，聚焦核心功能。
3. 对扩展开放。模型本身虽然只提供了基础的感知器、激活层、自动求导、数据接口等实现，但应当便于扩展实现更复杂算法。
4. 与Android无缝集成。支持与Android互相调用、传递数据，不需要像JNI那样写语法陌生的接口文件。

框架具体实现

目录结构

1. LeZero：项目根目录
2. framework：深度学习框架，包含已经训练好的模型，可以独立运行，也可以作为Android项目的依赖，打包时会将内部除.py外的全部文件放入assets目录
3. lezero：框架代码，全部由Python实现，可独立运行其内部的test_*.py文件
4. mnist_dataset：下载好的MNIST原始数据集，由于网络不稳定，直接用本地数据进行训练和推理
5. mpl_v2.npz：已经训练好的SGD模型，包含一个维度为1000的隐藏层，激活函数为Sigmoid，具体参数可见于lezero/train_model.py
6. lezero-android：Android项目，可用Android Studio打开运行

深度学习框架设计

整个框架的数据流如上图所示，包含数据的加载、模型初始化、正向传播、反向传播、参数的优化过程，支持训练+推理和仅推理两种模式。

可灵活替换的数据源

数据仓库Dataset

该类用于描述数据从哪里加载，可以是本地或者网络；用train=True来指定加载训练数据，如果仅用于验证则可以指明train=False，transform和target_transform分别是对data以及label的变换函数。

class Dataset:
    # train 是否为训练集
    # transform, target_transform 数据、标签预处理
    def __init__(self, train=True, transform=None, target_transform=None):
        self.train = train
        self.transform = transform
        self.target_transform = target_transform
        if self.transform is None:
            self.transform = lambda x: x
        if self.target_transform is None:
            self.target_transform = lambda x: x
        self.data = None
        self.label = None
        self.prepare()
        
    # 支持[]随机取数、长度函数len
    def __getitem__(self, index):
        assert np.isscalar(index) # 只支持标量
        if self.label is None:
            return self.transform(self.data[index]), None
        else:
            return self.transform(self.data[index]), self.target_transform(self.label[index])

    def __len__(self):
        return len(self.data)

    # 准备数据，子类需要实现该方法
    def prepare(self):
        pass

数据传输通道DataLoader

它连接Dataset与Model，提供对数据集的打乱、遍历操作，支持设置batch_size来小批量读取数据。

class DataLoader:
    # shuffle - 每个epoch当中是否打乱
    def __init__(self, dataset, batch_size, shuffle=True):
        self.dataset = dataset
        self.batch_size = batch_size
        self.shuffle = shuffle
        self.data_size = len(dataset)
        self.max_iter = math.ceil(self.data_size / batch_size)
        self.reset()

由点及面构建模型

该部分代码主要位于core.py，是模型的核心数据结构。

神经元Variable

神经网络的最小单元是Variable，作为神经元/感知器，它封装了ndarray作为data，同时记录自身的梯度用于误差反向传播（仅推理时为了增加运算速度、减少内存占用，不需要记录梯度）。同时还对Variable重载了+ - * / ** T 等操作符

# 最基本的运算单元
class Variable:
    def __init__(self, data, name=None):
        if data is not None:
            if not isinstance(data, np.ndarray):
                raise TypeError('{} is not supported'.format(type(data)))

        self.data = data
        self.grad = None
        self.creator = None
        self.generation = 0
        self.name = name

正向、反向传播函数Function

Function是对Variable进行运算的所有函数的基类，这里的函数不仅仅包含Add Sub等运算符重载函数，也包括Linear（线性关系）、MeanSquareError（均方差）、SoftmaxWithLoss（Softmax及损失函数）、ReLU（激活函数）等。

后续可以通过对Function进行扩展，实现更加复杂的训练及推理算法。

# 所有函数的基类，核心方法forward、backward
class Function:
    def __call__(self, *inputs): # *将变长参数转换为列表
        inputs = [as_variable(x) for x in inputs] # 支持直接输入ndarray
        xs = [x.data for x in inputs]
        ys = self.forward(*xs) # 使用*解包，ys是ndarray类型
        if not isinstance(ys, tuple):
            ys = (ys, )
        outputs = [Variable(as_array(y)) for y in ys]

        if Config.enable_backprop: # 训练时开启反向传播，推理时关闭
            self.generation = max([x.generation for x in inputs])
            for output in outputs:
                output.set_creator(self)
            self.inputs = inputs
            self.outputs = [weakref.ref(output) for output in outputs] # 弱引用防止循环引用
        
        return outputs if len(outputs) > 1 else outputs[0]

    def forward(self, xs):
        raise NotImplementedError()

    def backward(self, gys):
        raise NotImplementedError()

激活层Layer

Layer则是神经网络中“层”的概念，它保存了一层中的权重信息。模型的保存与加载也是以Layer为单位来实现的。

代码里实现了一个基础的线性关系层Linear，包含权重与偏置。

class Layer:
    def __init__(self):
        self._params = set() # 保存所有的params

    def __setattr__(self, name, value):
        if isinstance(value, (Parameter, Layer)): # Parameter、Layer 进行记录
            self._params.add(name)
        super().__setattr__(name, value)

    def __call__(self, *inputs):
        outputs = self.forward(*inputs)
        if not isinstance(outputs, tuple):
            outputs = (outputs,) # 保证outputs是tuple
        self.inputs = [weakref.ref(x) for x in inputs]
        self.outputs = [weakref.ref(y) for y in outputs]
        return outputs if len(outputs) > 1 else outputs[0] # 返回tuple或者单个元素

    def forward(self, inputs):
        raise NotImplementedError()

    def params(self): 
        for name in self._params: # yield是生成器，可用于for循环、next方法，只能被遍历一次
            obj = self.__dict__[name] 

            if isinstance(obj, Layer):
                yield from obj.params() #yield from获取生成器中的值
            else:
                yield obj

    def cleargrads(self):
        for param in self.params():
            param.cleargrad()
        
    # 生成 String-Parameter摊平集合
    def _flatten_params(self, params_dict, parent_key=""):
        for name in self._params:
            obj = self.__dict__[name]
            key = (parent_key + '/' + name) if parent_key else name

            if isinstance(obj, Layer):
                obj._flatten_params(params_dict, key)
            else:
                params_dict[key] = obj

    def save_weights(self, path):
        params_dict = {}
        self._flatten_params(params_dict)
        array_dict = {key: param.data for key, param in params_dict.items()
                      if param is not None}
        try:
            np.savez_compressed(path, **array_dict)
        except (Exception, KeyboardInterrupt) as e: # 保存失败（如Ctrl+C）则删除
            if os.path.exists(path):
                os.remove(path)
            raise

    def load_weights(self, path):
        npz = np.load(path)
        params_dict = {}
        self._flatten_params(params_dict)
        for key, param in params_dict.items(): # key-String, param-Variable
            param.data = npz[key]
            
# 线性关系
class Linear(Layer):
    
    def __init__(self, out_size, nobias=False, dtype=np.float32, in_size=None):
        super().__init__()
        self.in_size = in_size
        self.out_size = out_size
        self.dtype = dtype

        self.W = Parameter(None, name='W')
        if self.in_size is not None: # 未指定in_size则延后处理
            self._init_W()

        if nobias: # 无偏置
            self.b = None
        else:
            self.b = Parameter(np.zeros(out_size, dtype=dtype), name='b') # size=输出层

    def _init_W(self):
        I, O = self.in_size, self.out_size
        W_data = np.random.randn(I, O).astype(self.dtype) * np.sqrt(1/I)
        self.W.data = W_data

    def forward(self, x):
        # 在传播时根据输入矩阵shape初始化权重
        if self.W.data is None:
            self.in_size = x.shape[1]
            self._init_W()

        y = F.linear(x, self.W, self.b) # 调用functions.linear
        return y

模型Model

Model是隐藏层与激活函数的集合，可以在Model中对这两者进行数量、类型的配置。

class Model(L.Layer):
    def plot(self, *inputs, to_file='model.png'):
        y = self.forward(*inputs)
        return utils.plot_dot_graph(y, verbose=True, to_file=to_file)

# Multi-Layer Perceptron 多层感知器
class MLP(Model):
    def __init__(self, fc_output_sizes, activation=F.sigmoid):
        super().__init__()
        self.activation = activation
        self.layers = []

        for i, out_size in enumerate(fc_output_sizes):
            layer = L.Linear(out_size)
            setattr(self, 'l'+str(i), layer)
            self.layers.append(layer)

    def forward(self, x):
        for l in self.layers[:-1]:
            x = self.activation(l(x)) # 依次执行Linear+激活函数
        return self.layers[-1](x) # 最后一层

区分训练和推理

权重优化器Optimizer

目前实现了随机梯度下降（SGD）和动量（Momentum）两个方法，后续计划加上Adam、AdaGrad等。所有Optimizer应当继承自基类模板。

class Optimizer:
    def __init__(self):
        self.target = None
        self.hooks = []

    def setup(self, target):
        self.target = target
        return self

    def update(self):
        params = [p for p in self.target.params() if p.grad is not None]

        # 预处理
        for f in self.hooks:
            f(params)

        # 更新参数
        for param in params:
            self.update_one(param)

    # 更新单个参数
    def update_one(self, param):
        raise NotImplementedError()

    def add_hook(self, f):
        self.hooks.append(f)

# 随机梯度下降法
class SGD(Optimizer):
    def __init__(self, lr=0.01):
        super().__init__()
        self.lr = lr

    def update_one(self, param):
        param.data -= self.lr * param.grad.data

训练&模型保存

代码位于train_model.py。

1. 在创建Dataset时声明train=True
2. 在每个batch训练完成后进行反向传播loss.backward()，同时更新参数optimizer.update()
3. 最终通过model.save_weights(path)完成模型的保存

模型加载&推理

代码位于train_model.py。

1. 创建model与optimizer
2. 直接用model.save_weights()加载已经训练完成的权重给optimizer
3. 使用model(input_data)进行推理

Android Demo

手写输入DoodleView

onTouchEvent记录手指移动路径，利用Paint画笔，在onDraw时绘制涂鸦轨迹。

private val paint = Paint()
private val path = Path()

init {
    // 背景色
    setBackgroundColor(Color.BLACK)

    // 画笔色
    paint.color = Color.WHITE
    paint.strokeWidth = 50f
    paint.style = Paint.Style.STROKE
    paint.strokeCap = Paint.Cap.ROUND
}

override fun onDraw(canvas: Canvas) {
    super.onDraw(canvas)
    canvas.drawPath(path, paint)
}

override fun onTouchEvent(event: MotionEvent): Boolean {
    val x = event.x
    val y = event.y

    when (event.action) {
        MotionEvent.ACTION_DOWN -> {
            path.moveTo(x, y)
        }
        MotionEvent.ACTION_MOVE -> {
            path.lineTo(x, y)
        }
    }
    invalidate()
    return true
}

在生成Bitmap时，注意利用Matrix.postScale()对原图进行缩放，这里我偷懒直接缩放为28*28的尺寸，这也是MNIST数据集中的图片大小。在将Bitmap保存为IntArray时，应当留意Bitmap.getPixel(col, row)才是正确的像素读取顺序。

// 从Bitmap生成IntArray, 写入到参数array中
private fun fillBitmapData(rawBitmap: Bitmap, array: IntArray, smallWidth: Int, smallHeight: Int) {
    // 缩放比例
    val scaleWidth = smallWidth.toFloat() / rawBitmap.width
    val scaleHeight = smallHeight.toFloat() / rawBitmap.height
    // 缩放matrix
    val matrix = Matrix().apply {
        postScale(scaleWidth, scaleHeight)
    }
    val smallBitmap = Bitmap.createBitmap(rawBitmap, 0, 0, rawBitmap.width, rawBitmap.height, matrix, true)
    // 将bitmap生成像素图
    for (i in 0 until smallWidth) {
        for (j in 0 until smallHeight) { // 注意getPixe入参为横纵坐标，对应的是col、row
            val isDot = (
                    if (smallBitmap.getPixel(j, i) == Color.BLACK)
                        0
                    else
                        255
            )
            array[smallWidth * i + j] = isDot
        }
    }
}

集成模型

这里采用了chaquopy框架（github.com/chaquo/chaq… ，需要在模块的build.gradle文件里进行如下配置（本项目使用的是Groovy DSL写法，如果用Kotlin DSL的话，语法有所不同）。

plugins {
    alias(libs.plugins.androidApplication)
    alias(libs.plugins.jetbrainsKotlinAndroid)
    id 'com.chaquo.python'
}

android {

    defaultConfig {
        ...
        // 因为chaqopy是基于JNI实现，必须声明ndk
        ndk {
            abiFilters "arm64-v8a"
        }

        // 需要安装的三方库
        python {
            pip {
                install("numpy")
                install("matplotlib")
            }
            // 本机python可执行路径
            buildPython("D:\Software\anaconda3\python.exe")
        }
    }
    ...
}


// Python根目录，其下的全部.py文件会被作为JNI进行编译，其它类型文件则在安装后保存在应用files目录下。注意这个目录不可以是本Android项目的父目录，会导致gradle构建失败
chaquopy {
    sourceSets.getByName("main") {
        srcDirs = ["../../framework"]
    }
}

Python-Java互调

在APP层生成Bitmap的数组（类型为IntArray）后，调用Python模块inference_demo.py中定义的infer_user_input函数，并获取其输出：

val result = mPyModule.callAttr("infer_user_input", inputArray).toInt()

在Python函数中把入参转换为ndarray类型，直接加载已经训练好的模型（注意路径），进行推理。

# 推理用户输入
def infer_user_input(user_input):
    # 用np.array将java.jarray('I')转为nparray
    input_data = np.array(user_input)

    model_file_name = 'mlp_v2.npz'
    model_file_path = os.path.join(os.path.dirname(__file__), '..', model_file_name)
    
    hidden_size = 1000
    model = MLP((hidden_size, 10))

    if os.path.exists(model_file_path):
        model.load_weights(model_file_path)
        print('load finish:', model_file_path)
    else:
        raise ValueError(model_file_path, 'not exist!')

    print(model(input_data).data)
    infer = model(input_data).data.argmax(axis=0)
    print("infer=", infer)
    return infer

Github项目地址

github.com/LeiLiGithub…