『深度应用』OneFlow快速上手指南

 

现实中没有复杂的系统
--《极简主义》范式一：事情其实很简单

0.引子

为什么会有写这篇文章？

1.这不是一篇广告文，笔者不是大V，没人联系写稿，请放心食用。

2.这不是一片解析文，笔者水平有限，还无法做出深度解读。zhuanlan.zhihu.com/p/85111240 此篇对TF2.0的解析就差很大火候，让大家见笑了。

3.这可以算是一篇技术文章，通过对OneFlow安装，简单上手等操作来对比OneFlow较主流深度学习框架的难易程度。

简单分析下深度学习框架现状：

虽然不是解析文章，但针对深度学习框架分析还是要有一些的。现在主流的深度学习框架（不包括推理框架）可以说是TensorFlow，PyTorch，MXNet三足鼎力，其实这也是ta们背后资本与资源的角逐。TensorFlow凭借其完备性以及在部署方面的积累依旧处于优势地位，逼得PyTorch，MXNet等其他小众框架祭出了ONNX来打破TensorFlow的部署的垄断地位。PyTorch能成功从TensorFlow抢下学术界一块大蛋糕，很大程度上归功于其方向的正确性，没有在TensorFlow优势方面拼刺刀，而是在TensorFlow痛点处发力，通过易用性来争取用户。而MXNet就很难与TensorFlow与PyTorch有一战之力了，MXNet本身的一些优势也不是TensorFlow与PyTorch的使用痛点很难让用户有迁移的动力。

反观国内的赛场，百度的飞桨PaddlePaddle，虽然开源很早，但一直处于不温不火的状态。国内的框架的确还无法与国际主流框架展开竞争，差距还是巨大。但是今年还是有四家国内机构或公司选择了开源，华为的MineSpore，清华的Jittor，旷世的MegEngine，以及今天主角一流科技的OneFlow。虽竞争力不大，但在中美贸易摩擦下，自主产权越来越受到重视。

其实以笔者愚见，开发者并不会在意使用的是那个框架，更在意的是这个框架的完备性与便易性。当时大规模开发者从TensorFlow1.x转到PyTorch就可以说明，品牌对于开发者并不是很重要，重要的还是框架好不好用。所以笔者对国内开源框架还是积极持乐观态度的，做的好就会有用户支持。

下面，我们来看看OneFlow究竟做的如何。

1.上手

以下部分内容来自OneFlow开发文档： docs.oneflow.org/index.html

1.0简单介绍

OneFlow 是什么？

OneFlow 是开源的、采用全新架构设计，世界领先的工业级通用深度学习框架。

OneFlow优势是什么？

• 分布式训练全新体验，多机多卡如单机单卡一样简单
• 完美契合一站式平台(k8s + docker)
• 原生支持超大模型
• 近零运行时开销、线性加速比
• 灵活支持多种深度学习编译器
• 自动混合精度
• 中立开放，合作面广
• 持续完善的算子集、模型库

1.1安装指南

安装 OneFlow 稳定发布版

系统要求：前提Nvidia driver 与 CUDA安装完成

• Python >= 3.5
• Nvidia Linux x86_64 driver version >= 440.33

安装 OneFlow with legacy CUDA 指令如下：

pip install --find-links https://oneflow-inc.github.io/nightly oneflow_cu102 --user
pip install --find-links https://oneflow-inc.github.io/nightly oneflow_cu101 --user
pip install --find-links https://oneflow-inc.github.io/nightly oneflow_cu100 --user
pip install --find-links https://oneflow-inc.github.io/nightly oneflow_cu92 --user
pip install --find-links https://oneflow-inc.github.io/nightly oneflow_cu91 --user
pip install --find-links https://oneflow-inc.github.io/nightly oneflow_cu90 --user
复制代码

如何查看本机Nvidia Linux x86_64 driver version？可以看到笔者的Driver Version: 450.57

(base) song@songxpc:~$ nvidia-smi
Thu Aug 20 12:10:44 2020       
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 450.57       Driver Version: 450.57       CUDA Version: 11.0     |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|                               |                      |               MIG M. |
|===============================+======================+======================|
|   0  GeForce RTX 208...  Off  | 00000000:01:00.0 Off |                  N/A |
| 47%   54C    P2    80W / 257W |   5136MiB / 11019MiB |     38%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+
                                                                               
+-----------------------------------------------------------------------------+
| Processes:                                                                  |
|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |
|        ID   ID                                                   Usage      |
|=============================================================================|
|    0   N/A  N/A      1668      G   /usr/bin/totem                     10MiB |
|    0   N/A  N/A      1974      G   /usr/lib/xorg/Xorg                857MiB |
|    0   N/A  N/A      2128      G   /usr/bin/gnome-shell              331MiB |
|    0   N/A  N/A      2563      G   ...token=2514374358980620094      376MiB |
|    0   N/A  N/A     14664      C   python                           2531MiB |
|    0   N/A  N/A     27295      G   ...AAAAAAAAA= --shared-files      872MiB |
|    0   N/A  N/A     31690      G   ...token=3577040725527546973      149MiB |
+-----------------------------------------------------------------------------+
(base) song@songxpc:~$ 
复制代码

如何查看本机cuda版本？可以看到笔者的Cuda为10.2

(base) song@songxpc:~$ nvcc -V
nvcc: NVIDIA (R) Cuda compiler driver
Copyright (c) 2005-2019 NVIDIA Corporation
Built on Wed_Oct_23_19:24:38_PDT_2019
Cuda compilation tools, release 10.2, V10.2.89
复制代码

下面通过Conda来安装OneFlow框架，Conda推荐安装miniconda

依次输入：

conda create --name OF python=3.7 -y
conda activate OF
conda install cudatoolkit=10.2 cudnn=7 -y
pip install --find-links https://oneflow-inc.github.io/nightly oneflow_cu102 --user
复制代码

使用以下指令测试安装完成:

(OF) song@songxpc:~$ python
Python 3.7.8 | packaged by conda-forge | (default, Jul 31 2020, 02:25:08) 
[GCC 7.5.0] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import oneflow as of
>>> of.__version__
'0.1.8'
>>> 
复制代码

3分钟快速上手

这篇文章介绍了如何快速上手 OneFlow ，我们可以在3分钟内完成一个完整的神经网络训练过程。

运行一个例子

如果已经安装好了 OneFlow ，可以使用以下命令下载文档仓库中的mlp_mnist.py脚本，并运行。

wget https://docs.oneflow.org/code/quick_start/mlp_mnist.py #下载脚本
python mlp_mnist.py #运行脚本
复制代码

将得到类似以下输出：

2.7290366
0.81281316
0.50629824
0.35949975
0.35245502
...
复制代码

输出的是一串数字，每个数字代表了每一轮训练后的损失值，训练的目标是损失值越小越好。到此您已经用 OneFlow 完成了一个完整的神经网络的训练。

代码解读

下面是完整代码。

# mlp_mnist.py
import oneflow as flow
import oneflow.typing as tp

BATCH_SIZE = 100


@flow.global_function(type="train")
def train_job(
    images: tp.Numpy.Placeholder((BATCH_SIZE, 1, 28, 28), dtype=flow.float),
    labels: tp.Numpy.Placeholder((BATCH_SIZE,), dtype=flow.int32),
) -> tp.Numpy:
    with flow.scope.placement("cpu", "0:0"):
        initializer = flow.truncated_normal(0.1)
        reshape = flow.reshape(images, [images.shape[0], -1])
        hidden = flow.layers.dense(
            reshape,
            512,
            activation=flow.nn.relu,
            kernel_initializer=initializer,
            name="dense1",
        )
        logits = flow.layers.dense(
            hidden, 10, kernel_initializer=initializer, name="dense2"
        )
        loss = flow.nn.sparse_softmax_cross_entropy_with_logits(labels, logits)

    lr_scheduler = flow.optimizer.PiecewiseConstantScheduler([], [0.1])
    flow.optimizer.SGD(lr_scheduler, momentum=0).minimize(loss)

    return loss


if __name__ == "__main__":
    check_point = flow.train.CheckPoint()
    check_point.init()

    (train_images, train_labels), (test_images, test_labels) = flow.data.load_mnist(
        BATCH_SIZE, BATCH_SIZE
    )
    for i, (images, labels) in enumerate(zip(train_images, train_labels)):
        loss = train_job(images, labels)
        if i % 20 == 0:
            print(loss.mean())
复制代码

接下来让我们简单介绍下这段代码。

OneFlow 相对其他深度学习框架较特殊的地方是这里：

@flow.global_function(type="train")
def train_job(
    images: tp.Numpy.Placeholder((BATCH_SIZE, 1, 28, 28), dtype=flow.float),
    labels: tp.Numpy.Placeholder((BATCH_SIZE,), dtype=flow.int32),
) -> tp.Numpy:
复制代码

train_job是一个被@flow.global_function修饰的函数，通常被称为作业函数(job function)。只有被@flow.global_function修饰的作业函数才能够被 OneFlow 识别，通过type来指定job的类型：type="train"为训练作业；type="predict"为验证或预测作业。

在 OneFlow 中一个神经网络的训练或者预测作业需要两部分信息：

• 一部分是这个神经网络本身的结构和相关参数，这些在上文提到的作业函数里定义；
• 另外一部分是使用什么样的配置去训练这个网络，比如 learning rate 、模型优化更新的方法。这些job function里配置如下：

lr_scheduler = flow.optimizer.PiecewiseConstantScheduler([], [0.1]) flow.optimizer.SGD(lr_scheduler, momentum=0).minimize(loss)

这段代码里包含了训练一个神经网络的所有元素，除了上面说的作业函数及其配置之外：

• check_point.init(): 初始化网络模型参数；
• flow.data.load_mnist(BATCH_SIZE,BATCH_SIZE): 准备并加载训练数据；
• train_job(images, labels): 返回每一次训练的损失值；
• print(loss.mean()): 每训练20次，打印一次损失值。

以上只是一个简单网络的示例，在使用卷积神经网络进行手写体识别中，我们对使用OneFlow的流程进行了更加全面和具体的介绍。 另外，还可参考 OneFlow 基础专题中对于训练中各类问题的详细介绍。同时还提供了一些经典网络的样例代码及数据供参考。

2.解析

识别 MNIST 手写体数字

在这篇文章中，我们将学习：

• 使用 oneflow 接口配置软硬件环境
• 使用 oneflow 的接口定义训练模型
• 实现 oneflow 的训练作业函数
• 模型的保存和加载
• 实现 oneflow 的校验作业函数

本文通过使用 LeNet 模型，训练 MNIST 数据集向大家介绍使用 OneFlow 的各个核心环节，文末附有完整的示例代码。

在学习之前，也可以通过以下命令查看各脚本功能( 脚本运行依赖GPU )。

首先，同步本文档仓库并切换到对应路径：

git clone https://github.com/Oneflow-Inc/oneflow-documentation.git
cd oneflow-documentation/cn/docs/code/quick_start/
复制代码

• 模型训练python lenet_train.py
  以上命令将对 MNIST 数据集进行训练，并保存模型。

输出：

File mnist.npz already exist, path: ./mnist.npz
5.9947124
1.0865117
0.5317516
0.20937675
0.26428983
0.21764673
0.23443426
...
复制代码

训练模型是以下  lenet_eval.py 与  lenet_test.py 的前提条件，你也可以直接下载使用我们已经训练好的模型，略过训练步骤

#在仓库docs/code/quick_start/目录下
wget https://oneflow-public.oss-cn-beijing.aliyuncs.com/online_document/docs/quick_start/lenet_models_1.zip
unzip lenet_models_1.zip
复制代码

• 模型校验python lenet_eval.py
  以上命令，使用 MNIST 测试集对刚刚生成的模型进行校验，并给出准确率。

输出：

File mnist.npz already exist, path: ./mnist.npz
accuracy: 99.4%
复制代码

• 图像识别

python lenet_test.py ./9.png
# 输出：prediction: 9
复制代码

以上命令将使用之前训练的模型对我们准备好的 "9.png" 图片中的内容进行预测。 你也可以下载我们提取好的mnist图片，自行验证自己训练模型的预测效果。

MNIST数据集介绍

MNIST 是一个手写数字的数据库。包括了训练集与测试集；训练集包含了60000张图片以及图片对应的标签，测试集包含了60000张图片以及图片测试的标签。Yann LeCun 等已经将图片进行了大小归一化及居中处理，并且打包为二进制文件供下载。yann.lecun.com/exdb/mnist/

定义训练模型

在 oneflow.nn 及 oneflow.layers 提供了部分用于构建模型的算子。

def lenet(data, train=False):
    initializer = flow.truncated_normal(0.1)
    conv1 = flow.layers.conv2d(
        data,
        32,
        5,
        padding="SAME",
        activation=flow.nn.relu,
        name="conv1",
        kernel_initializer=initializer,
    )
    pool1 = flow.nn.max_pool2d(
        conv1, ksize=2, strides=2, padding="SAME", name="pool1", data_format="NCHW"
    )
    conv2 = flow.layers.conv2d(
        pool1,
        64,
        5,
        padding="SAME",
        activation=flow.nn.relu,
        name="conv2",
        kernel_initializer=initializer,
    )
    pool2 = flow.nn.max_pool2d(
        conv2, ksize=2, strides=2, padding="SAME", name="pool2", data_format="NCHW"
    )
    reshape = flow.reshape(pool2, [pool2.shape[0], -1])
    hidden = flow.layers.dense(
        reshape,
        512,
        activation=flow.nn.relu,
        kernel_initializer=initializer,
        name="dense1",
    )
    if train:
        hidden = flow.nn.dropout(hidden, rate=0.5, name="dropout")
    return flow.layers.dense(hidden, 10, kernel_initializer=initializer, name="dense2")
复制代码

以上代码中，我们搭建了一个 LeNet 网络模型。

实现训练作业函数

OneFlow 中提供了 oneflow.global_function 装饰器，通过它，可以将一个 Python 函数转变为训练作业函数（job function）。

global_function装饰器

oneflow.global_function 装饰器接收一个type作为参数用于指定作业类型，type="train"指定作业类型为训练，type="predict"指定作业类型为预测或验证。装饰器内部还有一个 function_config 对象参数用于具体类型的作业配置。

@flow.global_function(type="train")
def train_job(
    images: tp.Numpy.Placeholder((BATCH_SIZE, 1, 28, 28), dtype=flow.float),
    labels: tp.Numpy.Placeholder((BATCH_SIZE,), dtype=flow.int32),
    #作业函数实现 ...
复制代码

其中的 tp.Numpy.Placeholder 是数据占位符， tp.Numpy 指定这个作业函数在调用时，将返回一个 numpy 对象。

指定优化特征

我们可以通过 flow.optimizer 接口指定待优化参数和优化器。这样，OneFlow在每次迭代训练作业的过程中，将以指定的方式作为优化目标。

@flow.global_function(type="train")
def train_job(
    images: tp.Numpy.Placeholder((BATCH_SIZE, 1, 28, 28), dtype=flow.float),
    labels: tp.Numpy.Placeholder((BATCH_SIZE,), dtype=flow.int32),
) -> tp.Numpy:
    with flow.scope.placement("gpu", "0:0"):
        logits = lenet(images, train=True)
        loss = flow.nn.sparse_softmax_cross_entropy_with_logits(
            labels, logits, name="softmax_loss"
        )

    lr_scheduler = flow.optimizer.PiecewiseConstantScheduler([], [0.1])
    flow.optimizer.SGD(lr_scheduler, momentum=0).minimize(loss)
    return loss
复制代码

以上，我们通过 flow.nn.sparse_softmax_cross_entropy_with_logits 求得 loss ，并且将优化 loss 作为目标参数。

lr_scheduler 设定了学习率计划，[0.1]表明初始学习率为0.1；

flow.optimizer.SGD 则指定了优化器为sgd；loss作为参数传递给minimize表明优化器将以最小化loss为目标。

调用作业函数并交互

调用作业函数就可以开始训练。

调用作业函数时的返回结果，由定义作业函数时指定的返回值类型决定，可以返回一个，也可以返回多个结果。

返回一个结果的例子

在lenet_train.py中定义作业函数：

@flow.global_function(type="train")
def train_job(
    images: tp.Numpy.Placeholder((BATCH_SIZE, 1, 28, 28), dtype=flow.float),
    labels: tp.Numpy.Placeholder((BATCH_SIZE,), dtype=flow.int32),
) -> tp.Numpy:
    with flow.scope.placement("gpu", "0:0"):
        logits = lenet(images, train=True)
        loss = flow.nn.sparse_softmax_cross_entropy_with_logits(
            labels, logits, name="softmax_loss"
        )

    lr_scheduler = flow.optimizer.PiecewiseConstantScheduler([], [0.1])
    flow.optimizer.SGD(lr_scheduler, momentum=0).minimize(loss)
    return loss
复制代码

该作业函数的返回值类型为tp.Numpy，则当调用时，会返回一个numpy对象：

for epoch in range(20):
        for i, (images, labels) in enumerate(zip(train_images, train_labels)):
            loss = train_job(images, labels)
            if i % 20 == 0:
                print(loss.mean())
复制代码

我们调用了 train_job 并每循环20次打印1次 loss。

返回多个结果的例子

在校验模型的代码lenet_eval.py中定义的作业函数：

@flow.global_function(type="predict")
def eval_job(
    images: tp.Numpy.Placeholder((BATCH_SIZE, 1, 28, 28), dtype=flow.float),
    labels: tp.Numpy.Placeholder((BATCH_SIZE,), dtype=flow.int32),
) -> Tuple[tp.Numpy, tp.Numpy]:
    with flow.scope.placement("gpu", "0:0"):
        logits = lenet(images, train=False)
        loss = flow.nn.sparse_softmax_cross_entropy_with_logits(
            labels, logits, name="softmax_loss"
        )

    return (labels, logits)
复制代码

该作业函数的返回值类型为 Tuple[tp.Numpy, tp.Numpy]，则当调用时，会返回一个 tuple 容器，里面有2个元素，每个元素都是一个 numpy 对象：

for i, (images, labels) in enumerate(zip(test_images, test_labels)):
            labels, logits = eval_job(images, labels)
            acc(labels, logits)
复制代码

我们调用作业函数返回了labels与logits，并用它们评估模型准确率。

同步与异步调用

本文所有代码都是同步方式调用作业函数，实际上 OneFlow 还支持异步方式调用作业函数，具体内容在获取作业函数的结果一文中详细介绍。

模型的初始化、保存与加载

模型的初始化与保存

oneflow.train.CheckPoint 类构造的对象，可以用于模型的初始化、保存与加载。 在训练过程中，我们可以通过 init 方法初始化模型，通过 save 方法保存模型。如下例：

if __name__ == '__main__':
  check_point = flow.train.CheckPoint()
  check_point.init()
  #加载数据及训练 ...  
  check_point.save('./lenet_models_1')
复制代码

保存成功后，我们将得到名为 "lenet_models_1" 的 目录 ，该目录中包含了与模型参数对应的子目录及文件。

模型的加载

在校验或者预测过程中，我们可以通过 oneflow.train.CheckPoint.load 方法加载现有的模型参数。如下例：

if __name__ == '__main__':
  check_point = flow.train.CheckPoint()
  check_point.load("./lenet_models_1")
  #校验过程 ...
复制代码

load 自动读取之前保存的模型，并加载。

模型的校验

校验作业函数与训练作业函数 几乎没有区别 ，不同之处在于校验过程中的模型参数来自于已经保存好的模型，因此不需要初始化，也不需要在迭代过程中更新模型参数。

校验作业函数的编写

@flow.global_function(type="predict")
def eval_job(
    images: tp.Numpy.Placeholder((BATCH_SIZE, 1, 28, 28), dtype=flow.float),
    labels: tp.Numpy.Placeholder((BATCH_SIZE,), dtype=flow.int32),
) -> Tuple[tp.Numpy, tp.Numpy]:
    with flow.scope.placement("gpu", "0:0"):
        logits = lenet(images, train=False)
        loss = flow.nn.sparse_softmax_cross_entropy_with_logits(
            labels, logits, name="softmax_loss"
        )

    return (labels, logits)
复制代码

以上是校验训练作业函数的编写，声明了返回值类型是 Tuple[tp.Numpy, tp.Numpy]， 因此返回一个 tuple， tuple 中有2个元素，每个元素都是1个 numpy 对象。我们将调用训练作业函数，并根据返回结果计算准确率。

迭代校验

以下 acc 函数中统计样本的总数目，以及校验正确的总数目，我们将调用作业函数，得到 labels 与 logits：

g_total = 0
g_correct = 0


def acc(labels, logits):
    global g_total
    global g_correct

    predictions = np.argmax(logits, 1)
    right_count = np.sum(predictions == labels)
    g_total += labels.shape[0]
    g_correct += right_count
复制代码

调用校验作业函数：

if __name__ == "__main__":

    check_point = flow.train.CheckPoint()
    check_point.load("./lenet_models_1")
    (train_images, train_labels), (test_images, test_labels) = flow.data.load_mnist(
        BATCH_SIZE, BATCH_SIZE
    )

    for epoch in range(1):
        for i, (images, labels) in enumerate(zip(test_images, test_labels)):
            labels, logits = eval_job(images, labels)
            acc(labels, logits)

    print("accuracy: {0:.1f}%".format(g_correct * 100 / g_total))
复制代码

以上，循环调用校验函数，并且最终输出对于测试集的判断准确率。

预测图片

将以上校验代码修改，使得校验数据来自于原始的图片而不是现成的数据集，我们就可以使用模型进行图片预测。

def load_image(file):
    im = Image.open(file).convert("L")
    im = im.resize((28, 28), Image.ANTIALIAS)
    im = np.array(im).reshape(1, 1, 28, 28).astype(np.float32)
    im = (im - 128.0) / 255.0
    im.reshape((-1, 1, 1, im.shape[1], im.shape[2]))
    return im


def main():
    if len(sys.argv) != 2:
        usage()
        return

    check_point = flow.train.CheckPoint()
    check_point.load("./lenet_models_1")

    image = load_image(sys.argv[1])
    logits = eval_job(image, np.zeros((1,)).astype(np.int32))

    prediction = np.argmax(logits, 1)
    print("prediction: {}".format(prediction[0]))


if __name__ == "__main__":
    main()
复制代码

完整代码

训练模型

代码：lenet_train.py

#lenet_train.py
import oneflow as flow
import oneflow.typing as tp

BATCH_SIZE = 100


def lenet(data, train=False):
    initializer = flow.truncated_normal(0.1)
    conv1 = flow.layers.conv2d(
        data,
        32,
        5,
        padding="SAME",
        activation=flow.nn.relu,
        name="conv1",
        kernel_initializer=initializer,
    )
    pool1 = flow.nn.max_pool2d(
        conv1, ksize=2, strides=2, padding="SAME", name="pool1", data_format="NCHW"
    )
    conv2 = flow.layers.conv2d(
        pool1,
        64,
        5,
        padding="SAME",
        activation=flow.nn.relu,
        name="conv2",
        kernel_initializer=initializer,
    )
    pool2 = flow.nn.max_pool2d(
        conv2, ksize=2, strides=2, padding="SAME", name="pool2", data_format="NCHW"
    )
    reshape = flow.reshape(pool2, [pool2.shape[0], -1])
    hidden = flow.layers.dense(
        reshape,
        512,
        activation=flow.nn.relu,
        kernel_initializer=initializer,
        name="dense1",
    )
    if train:
        hidden = flow.nn.dropout(hidden, rate=0.5, name="dropout")
    return flow.layers.dense(hidden, 10, kernel_initializer=initializer, name="dense2")


@flow.global_function(type="train")
def train_job(
    images: tp.Numpy.Placeholder((BATCH_SIZE, 1, 28, 28), dtype=flow.float),
    labels: tp.Numpy.Placeholder((BATCH_SIZE,), dtype=flow.int32),
) -> tp.Numpy:
    with flow.scope.placement("gpu", "0:0"):
        logits = lenet(images, train=True)
        loss = flow.nn.sparse_softmax_cross_entropy_with_logits(
            labels, logits, name="softmax_loss"
        )

    lr_scheduler = flow.optimizer.PiecewiseConstantScheduler([], [0.1])
    flow.optimizer.SGD(lr_scheduler, momentum=0).minimize(loss)
    return loss


if __name__ == "__main__":
    flow.config.gpu_device_num(1)
    check_point = flow.train.CheckPoint()
    check_point.init()

    (train_images, train_labels), (test_images, test_labels) = flow.data.load_mnist(
        BATCH_SIZE, BATCH_SIZE
    )

    for epoch in range(20):
        for i, (images, labels) in enumerate(zip(train_images, train_labels)):
            loss = train_job(images, labels)
            if i % 20 == 0:
                print(loss.mean())
    check_point.save("./lenet_models_1")  # need remove the existed folder
    print("model saved")
复制代码

校验模型

代码：lenet_eval.py

预训练模型：lenet_models_1.zip

#lenet_eval.py
import numpy as np
import oneflow as flow
from typing import Tuple
import oneflow.typing as tp

BATCH_SIZE = 100


def lenet(data, train=False):
    initializer = flow.truncated_normal(0.1)
    conv1 = flow.layers.conv2d(
        data,
        32,
        5,
        padding="SAME",
        activation=flow.nn.relu,
        name="conv1",
        kernel_initializer=initializer,
    )
    pool1 = flow.nn.max_pool2d(
        conv1, ksize=2, strides=2, padding="SAME", name="pool1", data_format="NCHW"
    )
    conv2 = flow.layers.conv2d(
        pool1,
        64,
        5,
        padding="SAME",
        activation=flow.nn.relu,
        name="conv2",
        kernel_initializer=initializer,
    )
    pool2 = flow.nn.max_pool2d(
        conv2, ksize=2, strides=2, padding="SAME", name="pool2", data_format="NCHW"
    )
    reshape = flow.reshape(pool2, [pool2.shape[0], -1])
    hidden = flow.layers.dense(
        reshape,
        512,
        activation=flow.nn.relu,
        kernel_initializer=initializer,
        name="dense1",
    )
    if train:
        hidden = flow.nn.dropout(hidden, rate=0.5, name="dropout")
    return flow.layers.dense(hidden, 10, kernel_initializer=initializer, name="dense2")


@flow.global_function(type="predict")
def eval_job(
    images: tp.Numpy.Placeholder((BATCH_SIZE, 1, 28, 28), dtype=flow.float),
    labels: tp.Numpy.Placeholder((BATCH_SIZE,), dtype=flow.int32),
) -> Tuple[tp.Numpy, tp.Numpy]:
    with flow.scope.placement("gpu", "0:0"):
        logits = lenet(images, train=False)
        loss = flow.nn.sparse_softmax_cross_entropy_with_logits(
            labels, logits, name="softmax_loss"
        )

    return (labels, logits)


g_total = 0
g_correct = 0


def acc(labels, logits):
    global g_total
    global g_correct

    predictions = np.argmax(logits, 1)
    right_count = np.sum(predictions == labels)
    g_total += labels.shape[0]
    g_correct += right_count


if __name__ == "__main__":

    check_point = flow.train.CheckPoint()
    check_point.load("./lenet_models_1")
    (train_images, train_labels), (test_images, test_labels) = flow.data.load_mnist(
        BATCH_SIZE, BATCH_SIZE
    )

    for epoch in range(1):
        for i, (images, labels) in enumerate(zip(test_images, test_labels)):
            labels, logits = eval_job(images, labels)
            acc(labels, logits)

    print("accuracy: {0:.1f}%".format(g_correct * 100 / g_total))
复制代码

数字预测

代码：lenet_test.py

预训练模型：lenet_models_1.zip

MNIST 数据集图片mnist_raw_images.zip

import numpy as np
import oneflow as flow
from PIL import Image
import sys
import os
import oneflow.typing as tp

BATCH_SIZE = 1


def usage():
    usageHint = """
usage:
        python {0} <image file>
    eg: 
        python {0} {1}
    """.format(
        os.path.basename(sys.argv[0]), os.path.join(".", "9.png")
    )
    print(usageHint)


def lenet(data, train=False):
    initializer = flow.truncated_normal(0.1)
    conv1 = flow.layers.conv2d(
        data,
        32,
        5,
        padding="SAME",
        activation=flow.nn.relu,
        name="conv1",
        kernel_initializer=initializer,
    )
    pool1 = flow.nn.max_pool2d(
        conv1, ksize=2, strides=2, padding="SAME", name="pool1", data_format="NCHW"
    )
    conv2 = flow.layers.conv2d(
        pool1,
        64,
        5,
        padding="SAME",
        activation=flow.nn.relu,
        name="conv2",
        kernel_initializer=initializer,
    )
    pool2 = flow.nn.max_pool2d(
        conv2, ksize=2, strides=2, padding="SAME", name="pool2", data_format="NCHW"
    )
    reshape = flow.reshape(pool2, [pool2.shape[0], -1])
    hidden = flow.layers.dense(
        reshape,
        512,
        activation=flow.nn.relu,
        kernel_initializer=initializer,
        name="dense1",
    )
    if train:
        hidden = flow.nn.dropout(hidden, rate=0.5, name="dropout")
    return flow.layers.dense(hidden, 10, kernel_initializer=initializer, name="dense2")


@flow.global_function(type="predict")
def eval_job(
    images: tp.Numpy.Placeholder((BATCH_SIZE, 1, 28, 28), dtype=flow.float),
    labels: tp.Numpy.Placeholder((BATCH_SIZE,), dtype=flow.int32),
) -> tp.Numpy:
    with flow.scope.placement("gpu", "0:0"):
        logits = lenet(images, train=False)
    return logits


def load_image(file):
    im = Image.open(file).convert("L")
    im = im.resize((28, 28), Image.ANTIALIAS)
    im = np.array(im).reshape(1, 1, 28, 28).astype(np.float32)
    im = (im - 128.0) / 255.0
    im.reshape((-1, 1, 1, im.shape[1], im.shape[2]))
    return im


def main():
    if len(sys.argv) != 2:
        usage()
        return

    check_point = flow.train.CheckPoint()
    check_point.load("./lenet_models_1")

    image = load_image(sys.argv[1])
    logits = eval_job(image, np.zeros((1,)).astype(np.int32))

    prediction = np.argmax(logits, 1)
    print("prediction: {}".format(prediction[0]))


if __name__ == "__main__":
    main()
复制代码

3.总结

可以看出OneFlow模型搭建和训练还是很简单易用的，迁移到此框架成本不大。笔者要息笔锻炼一下了，希望大家都保持个好身体（~~）。。

-1.参考

-10.『王霸之路』从0.1到2.0一文看尽TensorFlow奋斗史 - 极简AI·小宋是呢的文章 - 知乎 zhuanlan.zhihu.com/p/85111240

-11.docs.oneflow.org/build_ship/…