训练一个LeNet模型

Linux Android 全流程 模型导出 模型转换 模型训练 初级 中级 高级

概述

本教程基于LeNet训练示例代码,演示MindSpore Lite训练功能的使用。

整个端侧训练流程分为以下三步:

  1. 基于MindSpore构建训练模型,并导出MindIR格式文件。

  2. 使用MindSpore Lite Converter工具,将MindIR模型转为端侧MS模型。

  3. 调用MindSpore Lite训练API,加载端侧MS模型,执行训练。

下面章节首先通过示例代码中集成好的脚本,帮你快速部署并执行示例,再详细讲解实现细节。

准备

以下操作均在PC上完成,推荐使用x86平台的Ubuntu 18.04 64位操作系统。

下载数据集

我们示例中用到的MNIST数据集是由10类28*28的灰度图片组成,训练数据集包含60000张图片,测试数据集包含10000张图片。

MNIST数据集下载页面:http://yann.lecun.com/exdb/mnist/。页面提供4个数据集下载链接,其中前2个文件是训练数据需要,后2个文件是测试结果需要。

将数据集下载并解压到本地路径下,这里将数据集解压分别存放到/PATH/MNIST_Data/train/PATH/MNIST_Data/test路径下。

目录结构如下:

MNIST_Data/
├── test
│   ├── t10k-images-idx3-ubyte
│   └── t10k-labels-idx1-ubyte
└── train
    ├── train-images-idx3-ubyte
    └── train-labels-idx1-ubyte

安装MindSpore

安装MindSpore CPU环境,具体请参考MindSpore安装

获取Converter和Runtime

可以通过MindSpore Lite源码编译生成模型训练所需的train-linux-x64以及train-android-aarch64包。编译命令如下:

# 生成converter工具以及x86平台的runtime包
bash build.sh -I x86_64 -T on -e cpu -j8

# 生成arm64平台的runtime包
bash build.sh -I arm64 -T on -e cpu -j8

你也可以下载MindSpore Lite直接下载所需要的转换工具以及模型训练框架,并将它们放在MindSpore源码下的output目录(如果没有output目录,请创建它)。

连接安卓设备

准备好一台Android设备,并通过USB与工作电脑正确连接。手机需开启“USB调试模式”,华为手机一般在设置->系统和更新->开发人员选项->USB调试中打开“USB调试模式”。

本示例使用adb工具与Android设备进行通信,在工作电脑上远程执行各类设备操作;如果没有安装adb工具,可以执行apt install adb安装。

模型训练和验证

示例代码在MindSpore源码下的mindspore/lite/examples/train_lenet目录。本地克隆MindSpore源码后,进入mindspore/lite/examples/train_lenet目录,执行如下命令后,脚本会导出lenet_tod.mindir模型,然后利用converter工具将MindIR模型转换为MindSpore Lite可以识别的lenet_tod.ms模型;最后,将lenet_tod.ms模型文件、MNIST数据集以及MindSpore Lite训练runtime包推送到Andorid设备上,执行训练。

bash prepare_and_run.sh -D /PATH/MNIST_Data -t arm64

其中,/PATH/MNIST_Data是你工作电脑上存放MNIST数据集的绝对路径,-t arm64表示我们将在Android设备上执行训练和推理。

在Android设备上训练LeNet模型每100轮会输出损失值和准确率;最后选择训练完成的模型执行推理,验证MNIST手写字识别精度。在端侧训练的LeNet模型能够达到97%的识别率,结果如下所示(测试准确率会受设备差异的影响):

Training on Device
100:    Loss is 0.853509 [min=0.581739]  max_acc=0.674079
200:    Loss is 0.729228 [min=0.350235]  max_acc=0.753305
300:    Loss is 0.379949 [min=0.284498]  max_acc=0.847957
400:    Loss is 0.773617 [min=0.186403]  max_acc=0.867788
500:    Loss is 0.477829 [min=0.0688716]  max_acc=0.907051
600:    Loss is 0.333066 [min=0.0688716]  max_acc=0.93099
700:    Loss is 0.197988 [min=0.0549653]  max_acc=0.940905
800:    Loss is 0.128299 [min=0.048147]  max_acc=0.946314
900:    Loss is 0.43212 [min=0.0427626]  max_acc=0.955729
1000:   Loss is 0.446575 [min=0.033213]  max_acc=0.95643
1100:   Loss is 0.162593 [min=0.025461]  max_acc=0.95643
1200:   Loss is 0.177662 [min=0.0180249]  max_acc=0.95643
1300:   Loss is 0.0425688 [min=0.00832943]  max_acc=0.95643
1400:   Loss is 0.270186 [min=0.00832943]  max_acc=0.963041
1500:   Loss is 0.0340949 [min=0.00832943]  max_acc=0.963041
1600:   Loss is 0.205415 [min=0.00832943]  max_acc=0.969551
1700:   Loss is 0.0269625 [min=0.00810314]  max_acc=0.970152
1800:   Loss is 0.197761 [min=0.00680999]  max_acc=0.970152
1900:   Loss is 0.19131 [min=0.00680999]  max_acc=0.970152
2000:   Loss is 0.182704 [min=0.00680999]  max_acc=0.970453
2100:   Loss is 0.375163 [min=0.00313038]  max_acc=0.970453
2200:   Loss is 0.296488 [min=0.00313038]  max_acc=0.970453
2300:   Loss is 0.0556241 [min=0.00313038]  max_acc=0.970453
2400:   Loss is 0.0753383 [min=0.00313038]  max_acc=0.973057
2500:   Loss is 0.0732852 [min=0.00313038]  max_acc=0.973057
2600:   Loss is 0.220644 [min=0.00313038]  max_acc=0.973057
2700:   Loss is 0.0159947 [min=0.00313038]  max_acc=0.973257
2800:   Loss is 0.0800904 [min=0.00168969]  max_acc=0.973257
2900:   Loss is 0.0210299 [min=0.00168969]  max_acc=0.97476
3000:   Loss is 0.256663 [min=0.00168969]  max_acc=0.97476
accuracy = 0.970553

Load trained model and evaluate accuracy
accuracy = 0.970553

如果你没有Android设备,也可以执行bash prepare_and_run.sh -D /PATH/MNIST_Data -t x86直接在PC上运行本示例。

示例程序详解

示例程序结构

  train_lenet/
  ├── model
  │   ├── lenet_export.py
  │   ├── prepare_model.sh
  │   └── train_utils.py
  ├── scripts
  │   ├── eval.sh
  │   ├── run_eval.sh
  │   ├── train.sh
  │   └── run_train.sh
  │
  ├── src
  │   ├── dataset.cc
  │   ├── dataset.h
  │   ├── net_runner.cc
  │   └── net_runner.h
  │
  ├── README.md
  └── prepare_and_run.sh

定义并导出模型

首先我们基于MindSpore框架创建一个LeNet5模型,你也可以直接用MindSpore model_zoo的现有LeNet5模型

本小节完全使用MindSpore云侧功能,进一步了解MindSpore请参考MindSpore教程

import sys
from mindspore import context, Tensor, export
from mindspore import dtype as mstype
from lenet import LeNet5
import numpy as np
from train_utils import TrainWrap

sys.path.append('./mindspore/model_zoo/official/cv/lenet/src/')

n = LeNet5()
n.set_train()
context.set_context(mode=context.PYNATIVE_MODE, device_target="CPU", save_graphs=False)

然后定义输入和标签张量大小:

batch_size = 32
x = Tensor(np.ones((batch_size, 1, 32, 32)), mstype.float32)
label = Tensor(np.zeros([batch_size, 10]).astype(np.float32))
net = TrainWrap(n)

定义损失函数、网络可训练参数、优化器,并启用单步训练,由TrainWrap函数实现。

import mindspore.nn as nn
from mindspore import ParameterTuple
def TrainWrap(net, loss_fn=None, optimizer=None, weights=None):
    if loss_fn == None:
        loss_fn = nn.SoftMaxCrossEntropyWithLogits()
    loss_net = nn.WithLossCell(net, loss_fn)
    loss_net.set_train()
    if weights == None:
        weights = ParameterTuple(net.trainable_params())
    if optimizer == None:
        optimizer = nn.Adam(weights, learning_rate=1e-3, beta1=0.9 beta2=0.999, eps=1e-8, use_locking=False, use_nesterov=False, weight_decay=0.0, loss_scale=1.0)
    train_net = nn.TrainOneStepCell(loss_net, optimizer)

最后调用export接口将模型导出为MindIR文件保存(目前端侧训练仅支持MindIR格式)。

export(net, x, label, file_name="lenet_tod", file_format='MINDIR')
print("finished exporting")

如果输出finished exporting表示导出成功,生成的lenet_tod.mindir文件在当前目录下。完整代码参见lenet_export.pytrain_utils.py

转换模型

得到lenet_tod.mindir文件后,使用MindSpore Lite converter工具将其转还为可用于端侧训练的模型文件,执行指令如下:

./converter_lite --fmk=MINDIR --trainModel=true --modelFile=lenet_tod.mindir --outputFile=lenet_tod

转换成功后,当前目录下会生成lenet_tod.ms模型。

详细的converter工具使用,可以参考训练模型转换

训练模型

源码src/net_runner.cc使用MindSpore Lite训练API完成模型训练,主函数如下:

int NetRunner::Main() {
  InitAndFigureInputs();

  InitDB();

  TrainLoop();

  float acc = CalculateAccuracy();
  std::cout << "accuracy = " << acc << std::endl;

  if (cycles_ > 0) {
    auto trained_fn = ms_file_.substr(0, ms_file_.find_last_of('.')) + "_trained_" + std::to_string(cycles_) + ".ms";
    session_->SaveToFile(trained_fn);
  }
  return 0;
}
  1. 加载模型

    InitAndFigureInputs函数加载转换后的MS模型文件,调用CreateSession接口创建TrainSession实例(下述代码中的ms_file_就是转换模型阶段生成的lenet_tod.ms模型)。同时根据模型的输入Tensor设置batch_sziedata_size

    void NetRunner::InitAndFigureInputs() {
      mindspore::lite::Context context;
      context.device_list_[0].device_info_.cpu_device_info_.cpu_bind_mode_ = mindspore::lite::NO_BIND;
      context.thread_num_ = 1;
    
      session_ = mindspore::session::TrainSession::CreateSession(ms_file_, &context);
      assert(nullptr != session_);
    
      auto inputs = session_->GetInputs();
      assert(inputs.size() > 1);
      this->data_index_  = 0;
      this->label_index_ = 1;
      this->batch_size_ = inputs[data_index_]->shape()[0];
      this->data_size_  = inputs[data_index_]->Size() / batch_size_;  // in bytes
      if (verbose_) {
        std::cout << "data size: " << data_size_ << "\nbatch size: " << batch_size_ << std::endl;
      }
    }
    
  2. 数据集处理

    InitDB函数初始化MNIST数据集,调用DataSet加载训练数据以及相应标签。

    int NetRunner::InitDB() {
      if (data_size_ != 0) ds_.set_expected_data_size(data_size_);
      int ret = ds_.Init(data_dir_, DS_MNIST_BINARY);
      num_of_classes_ = ds_.num_of_classes();
      if (ds_.test_data().size() == 0) {
        std::cout << "No relevant data was found in " << data_dir_ << std::endl;
        assert(ds_.test_data().size() != 0);
      }
    
      return ret;
    }
    
  3. 执行训练

    TrainSessionDataSet创建完毕后,就可以开始训练了。首先调用TrainSessionTrain方法,将模型设置为训练模式;然后循环调用RunGraph函数,通过DataSet读取训练数据,执行训练;在一轮训练完成后,可以调用SaveToFile方法保存成CheckPoint模型,CheckPoint模型包含已更新的权重,在应用崩溃或设备出现故障时可以直接加载CheckPoint模型,继续开始训练。

    int NetRunner::TrainLoop() {
        session_->Train();
        float  min_loss = 1000.;
        float max_acc = 0;
        for (int i = 0; i < cycles_; i++) {
            FillInputData(ds_.train_data());
            session_->RunGraph(nullptr, verbose_? after_callback : nullptr);
            float loss = GetLoss();
            if (min_loss > loss) {
                min_loss = loss;
            }
            if (save_checkpoint_ != 0 && (i+1)%save_checkpoint_ == 0) {
                auto cpkt_file = ms_file_.substr(0, ms_file_.find_last_of('.')) + "_trained_" + std::to_string(i+1) + ".ms";
                session_->SaveToFile(cpkt_file);
            }
    
            std::cout << i + 1 << ": Loss is " << loss[0] << " [min=" << min_loss << "]";
            if (max_acc > 0) {
                std::cout << "max_acc=" << max_acc;
            }
            std::cout << std::endl;
            if ((i+1)%100 == 0) {
                float acc = CalculateAccuracy(10);
                if (max_acc < acc) {
                    max_acc = acc;
                }
                std::cout << "accuracy (on " << batch_size_ * 10 << " samples) = " << acc;
                std::cout << "max accuracy= " << max_acc << std::endl;
            }
        }
        return 0;
    }
    
  4. 验证精度

    在每一轮训练结束后,都会调用CalculateAccuracy评估模型精度。验证精度是执行推理流程,需要首先调用TrainSessionEval方法,将当前的模型设置为推理模式;然后通过读取/PATH/MNIST_Data/test测试数据集,调用RunGraph执行推理,获取模型输出并比对结果得到识别率。为了下一轮的训练能够继续执行,在获取当前阶段的精度后,调用Train方法,将TrainSession重新设置为训练模式。

    float NetRunner::CalculateAccuracy(int max_tests) const {
      float accuracy = 0.0;
      const std::vector<DataLabelTuple> test_set = ds_.test_data();
      int tests = test_set.size() / batch_size_;
      if (max_tests != -1 && tests < max_tests) tests = max_tests;
    
      session_->Eval();
      for (int i = 0; i < tests; i++) {
        auto labels = FillInputData(test_set, (max_tests == -1));
        session_->RunGraph();
        auto outputsv = SearchOutputsForSize(batch_size_ * num_of_classes_);
        assert(outputsv != nullptr);
        auto scores = reinterpret_cast<float *>(outputsv->MutableData());
        for (int b = 0; b < batch_size_; b++) {
          int max_idx = 0;
          float max_score = scores[num_of_classes_ * b];
          for (int c = 0; c < num_of_classes_; c++) {
            if (scores[num_of_classes_ * b + c] > max_score) {
              max_score = scores[num_of_classes_ * b + c];
              max_idx = c;
            }
          }
          if (labels[b] == max_idx) accuracy += 1.0;
        }
      }
      session_->Train();
      accuracy /= static_cast<float>(batch_size_ * tests);
      return accuracy;
    }