本文共 4031 字，大约阅读时间需要 13 分钟。

LeNet-5网络在昇腾训练芯片上的实现与训练

摘要

LeNet-5是Yann LeCun于1998年设计的用于手写数字识别的经典卷积神经网络，被誉为深度学习领域的“Hello world！”。它曾被广泛应用于支票识别领域。近年来，我将LeNet-5迁移到昇腾训练芯片上，并分享了完整的训练实现和代码。以下是实现过程和详细说明。

环境约束

昇腾910芯片目前支持TensorFlow 1.15版本，开发环境需按照官方文档进行配置。请确保已完成基础镜像的上传和启动。

启动训练镜像

以下命令用于挂载昇腾910的所有GPU资源，并准备好训练环境：

docker run -it --net=host --device=/dev/davinci0 --device=/dev/davinci1 --device=/dev/davinci2 --device=/dev/davinci3 --device=/dev/davinci4 --device=/dev/davinci5 --device=/dev/davinci6 --device=/dev/davinci7 --device=/dev/davinci_manager --device=/dev/devmm_svm --device=/dev/hisi_hdc -v /var/log/npu/slog/container/docker:/var/log/npu/slog -v /var/log/npu/conf/slog/slog.conf:/var/log/npu/conf/slog/slog.conf -v /usr/local/Ascend/driver/lib64/:/usr/local/Ascend/driver/lib64/ -v /usr/local/Ascend/driver/tools/:/usr/local/Ascend/driver/tools/ -v /data/:/data/ -v /home/code:/home/local/code -v ~/context:/cache ubuntu_18.04-docker.arm64v8:v2 /bin/bash

LeNet-5网络简介

LeNet-5由7层组成（不计输入层），每层都包含可训练的参数。其网络结构如下：

• C1：5x5卷积核，输出大小为28x28。
• S2：2x2子采样，输出16个14x14特征图。
• C3：6个5x5卷积核，输出16个特征图。
• S4：2x2子采样，输出16个5x5特征图。
• C5：120个5x5卷积核，输出1x1特征图。
• F6：完全连接层，输出84个特征图。

LeNet-5的设计相对简单，但在1990年代奠定了深度学习的基础。

训练案例

本文使用LeNet-5作为训练案例，基于MNIST手写数字识别数据集。以下是训练脚本的代码：

import os
import numpy as np
import tensorflow as tf
from tensorflow.examples.tutorials.mnist import input_data
from npu_bridge.estimator import npu_ops
# 导入NPU算子库
from tensorflow.core.protobuf.rewriter_config_pb2 import RewriterConfig
batch_size = 100
learning_rate = 0.1
training_step = 10000
model_save_path = "./model/"
model_name = "model.ckpt"
def train(mnist):
    x = tf.placeholder(tf.float32, [batch_size, mnist_inference.image_size, mnist_inference.image_size, mnist_inference.num_channels], name='x-input')
    y_ = tf.placeholder(tf.float32, [batch_size, mnist_inference.num_labels], name='y-input')
    
    regularizer = tf.contrib.layers.l2_regularizer(0.001)
    y = mnist_inference.inference(x, train=True, regularizer=regularizer)
    
    global_step = tf.Variable(0, trainable=False)
    cross_entropy = tf.nn.sparse_softmax_cross_entropy_with_logits(logits=y, labels=tf.argmax(y_, 1))
    cross_entropy_mean = tf.reduce_mean(cross_entropy)
    loss = cross_entropy_mean + tf.add_n(tf.get_collection("loss"))
    
    train_step = tf.train.GradientDescentOptimizer(learning_rate).minimize(loss, global_step=global_step)
    saver = tf.train.Saver()
    config = tf.ConfigProto(
        allow_soft_placement=True,
        log_device_placement=False
    )
    custom_op = config.graph_options.rewrite_options.custom_optimizers.add()
    custom_op.name = "NpuOptimizer"
    custom_op.parameter_map["use_off_line"].b = True
    
    writer = tf.summary.FileWriter("./log_dir", tf.get_default_graph())
    writer.close()
    with tf.Session(config=config) as sess:
        tf.global_variables_initializer().run()
        start_time = time.time()
        
        for i in range(training_step):
            xs, ys = mnist.train.next_batch(batch_size)
            reshaped_xs = np.reshape(xs, (batch_size, mnist_inference.image_size, mnist_inference.image_size, mnist_inference.num_channels))
            
            _, loss_value, step = sess.run([train_step, loss, global_step], feed_dict={x: reshaped_xs, y_: ys})
            
            if i % 10 == 0:
                print("After %d training steps, loss on training batch is %g, total time in this 1000 steps is %s." % (step, loss_value, time.time() - start_time))
                saver.save(sess, os.path.join(model_save_path, model_name), global_step=global_step)
                start_time = time.time()
def main():
    mnist = input_data.read_data_sets('MNIST_DATA/', one_hot=True)
    train(mnist)
if __name__ == "__main__":
    main()

代码解释

该代码实现了LeNet-5网络的训练过程。以下是主要部分的说明：

模型训练

运行上述代码时，模型将在./model/目录下生成训练结果。训练完成后，可以加载预训练模型进行推理任务。

总结

LeNet-5是深度学习的经典网络之一，其实现和迁移至昇腾训练芯片的过程展示了从经典算法到现代硬件实现的演变。通过以上代码，读者可以轻松在昇腾平台上完成LeNet-5的训练和部署。

转载地址：http://zrmuz.baihongyu.com/