手把手教你从0到1学会tensorflow进行模型训练，并能在网页轻松运行

项目体验地址：http://at.iunitv.cn/

效果预览：

花絮：

很多小伙伴嘴上说着学不动了，其实身体还是很诚实的。

毕竟读书还是有很多好处的：比如让你的脑门散发智慧的光芒，再或者让你有理由说因为读书太忙了所以没有女朋友等等。所以在这个特殊的日子里，你这一年的图书我们承包了。不为别的，只为帮助在座的各位在2020年能够遇见更好的自己!

今天，我们也想要借助这个特殊的机会，普及一下Tensorflow相关的知识，我们会用TensorFlow.js做一个图书识别的模型，并在Vue Application中运行，赋予网页识别图书的能力。

本文讲述了AI相关的概念知识和如何运用SSD Mobile Net V1模型进行迁移学习的方法，从而帮助大家完成一个可以在网页上运行的图书识别模型。

【文末有活动哦】

正文：

什么是迁移学习

迁移学习和域适应指的是在一种环境中学到的知识被用在另一个领域中来提高它的泛化性能。——《深度学习》，第 526 页

再简单一点理解，以今天图书识别模型训练为例，我们利用前人训练好的具备图片识别能力的AI模型，保留AI模型中对图片特征提取的能力的基础上再训练，使AI模型具备识别图书的能力。

迁移学习能够大大提高模型训练的速度，并达到相对不错的正确率。

而我们今天所要迁移学习的对象就是SSD Mobile Net V1模型，初次接触神经网络的同学可以将其理解为一种具备图片识别的轻便小巧的AI模型，它能够在移动设备上高效地运行。对这个模型具体的神经网络设计结构感兴趣的同学可以自行搜索。

了解了基本的概念之后，我们便开始动手吧！我们可以基于SSD Mobile Net模型去设计一个属于自己的AI模型，并让它在Vue Application中运行。

Object Detection（目标识别）

本次项目是为了训练一个Object Detection的模型，即目标识别的模型，该模型能够识别并圈选出图片中相应的目标对象。

准备工作

同步开发环境

为了避免小伙伴因为环境问题遇到各种各样的坑，在工作开展之前，我们先跟大家同步一下运行的环境。大家如果要动手去做，也尽量跟我们的运行环境保持一致，这样可以有效避免踩坑，规避“从入门到放弃”的现象。

开发环境：

• 系统Mac OS系统
• Python版本：3.7.3
• TensorFlow版本：1.15.2
• TensorFlowJS版本：1.7.2
• 开发工具：Pycharm和Webstorm

下载项目

同步完开发环境后，终于要开始动工了。首先我们需要在Github上下载几个项目：

• 迁移模型训练项目
• 图片格式转换项目
• 图片圈选工具LabelImg

准备图片素材

我们可以通过搜索引擎收集有关图书的图片素材：

其次，我们可以在Github上克隆LabelImg项目，并根据Github的使用说明，按照不同的环境安装运行LabelImg项目，运行后的页面如下：

然后我们按照以下步骤，将图片格式转换为圈选区域后的XML文件：

1. 打开图片存放的目录
2. 选择圈选后的存放目录
3. 圈选图片目标区域
4. 设置圈选区域的标签
5. 保存成XML格式

存放完后我们在存放的目录下会看到许多XML格式的文件，这个文件记录了图片的位置信息、圈选信息和标签信息等，用于后续的模型训练。

配置安装Object Detection的环境

从Github克隆迁移模型训练的项目迁移模型训练项目，注意要在r1.5分支运行，并用PyCharm打开项目。

项目的目录环境为上图，首先我们需要下载TensorFlow1.15.2版本：

pip install tensorflow==1.15.2

其次安装依赖包：

sudo pip install pillow
sudo pip install lxml
sudo pip install jupyter
sudo pip install matplotlib

然后通过终端切换到research目录，并执行几行配置命令，具体请参考Github的使用说明：

cd ./research
protoc object_detection/protos/*.proto --python_out=.
export PYTHONPATH=$PYTHONPATH:`pwd`:`pwd`/slim

最后我们运行model_builder_test.py文件，如果在终端中看到OK字样，表示配置成功。

python object_detection/builders/model_builder_test.py

将XML格式转换为TensorFlow需要的TFRecord格式

克隆并打开图片格式转换项目，然后我们对该项目加以小改造：

改造文件目录:

1. 删除annotations、data、training目录中的内容
2. 增加一个xmls目录，用以存放xml文件

改造文件：
接着，我们再改造以下2个文件并新增一个文件，方便我们转换图片格式

1. 改造xml_to_csv.py为：

   import os
   import glob
   import pandas as pd
   import xml.etree.ElementTree as ET
   import random
   import time
   import shutil
   
   class Xml2Cvs:
       def __init__(self):
           self.xml_filepath = r'./xmls'
           self.save_basepath = r"./annotations"
           self.trainval_percent = 0.9
           self.train_percent = 0.85
   
       def xml_split_train(self):
   
           total_xml = os.listdir(self.xml_filepath)
           num = len(total_xml)
           list = range(num)
           tv = int(num * self.trainval_percent)
           tr = int(tv * self.train_percent)
           trainval = random.sample(list, tv)
           train = random.sample(trainval, tr)
           print("train and val size", tv)
           print("train size", tr)
           start = time.time()
           test_num = 0
           val_num = 0
           train_num = 0
           for i in list:
               name = total_xml[i]
               if i in trainval:
                   if i in train:
                       directory = "train"
                       train_num += 1
                       xml_path = os.path.join(os.getcwd(), 'annotations/{}'.format(directory))
                       if (not os.path.exists(xml_path)):
                           os.mkdir(xml_path)
                       filePath = os.path.join(self.xml_filepath, name)
                       newfile = os.path.join(self.save_basepath, os.path.join(directory, name))
                       shutil.copyfile(filePath, newfile)
                   else:
                       directory = "validation"
                       xml_path = os.path.join(os.getcwd(), 'annotations/{}'.format(directory))
                       if (not os.path.exists(xml_path)):
                           os.mkdir(xml_path)
                       val_num += 1
                       filePath = os.path.join(self.xml_filepath, name)
                       newfile = os.path.join(self.save_basepath, os.path.join(directory, name))
                       shutil.copyfile(filePath, newfile)
               else:
                   directory = "test"
                   xml_path = os.path.join(os.getcwd(), 'annotations/{}'.format(directory))
                   if (not os.path.exists(xml_path)):
                       os.mkdir(xml_path)
                   test_num += 1
                   filePath = os.path.join(self.xml_filepath, name)
                   newfile = os.path.join(self.save_basepath, os.path.join(directory, name))
                   shutil.copyfile(filePath, newfile)
   
           end = time.time()
           seconds = end - start
           print("train total : " + str(train_num))
           print("validation total : " + str(val_num))
           print("test total : " + str(test_num))
           total_num = train_num + val_num + test_num
           print("total number : " + str(total_num))
           print("Time taken : {0} seconds".format(seconds))
   
       def xml_to_csv(self, path):
           xml_list = []
           for xml_file in glob.glob(path + '/*.xml'):
               tree = ET.parse(xml_file)
               root = tree.getroot()
               print(root.find('filename').text)
               for object in root.findall('object'):
                   value = (root.find('filename').text,
                            int(root.find('size').find('width').text),
                            int(root.find('size').find('height').text),
                            object.find('name').text,
                            int(object.find('bndbox').find('xmin').text),
                            int(object.find('bndbox').find('ymin').text),
                            int(object.find('bndbox').find('xmax').text),
                            int(object.find('bndbox').find('ymax').text)
                            )
                   xml_list.append(value)
           column_name = ['filename', 'width', 'height', 'class', 'xmin', 'ymin', 'xmax', 'ymax']
           xml_df = pd.DataFrame(xml_list, columns=column_name)
           return xml_df
   
       def main(self):
           for directory in ['train', 'test', 'validation']:
               xml_path = os.path.join(os.getcwd(), 'annotations/{}'.format(directory))
               xml_df = self.xml_to_csv(xml_path)
               xml_df.to_csv('data/mask_{}_labels.csv'.format(directory), index=None)
               print('Successfully converted xml to csv.')
   
   
   if __name__ == '__main__':
       Xml2Cvs().xml_split_train()
       Xml2Cvs().main()
   

2. 改造generate_tfrecord.py文件，将csv格式转换为TensorFlow需要的record格式：

将该区域的row_label改成我们LabelImg中的标签名，因为我们只有一个标签，所以直接修改成book即可。

3. 新增一个generate_tfrecord.sh脚本，方便执行generate_tfrecord.py文件

   #!/usr/bin/env bash
   python generate_tfrecord.py --csv_input=data/mask_train_labels.csv  --output_path=data/mask_train.record --image_dir=images
   python generate_tfrecord.py --csv_input=data/mask_test_labels.csv  --output_path=data/mask_test.record --image_dir=images
   python generate_tfrecord.py --csv_input=data/mask_validation_labels.csv  --output_path=data/mask_validation.record --image_dir=images
   
   

配置Object Decation的环境：

export PYTHONPATH=$PYTHONPATH:你的models/research/slim所在的全目录路径

最后我们将图片文件复制到images目录，将xml文件复制到xmls目录下，再执行xml_to_csv.py文件，我们会看到data目录下产生了几个csv格式结尾的文件；这时，我们在终端执行generate_tfrecord.sh文件，TensorFlow所需要的数据格式就大功告成啦。

迁移训练模型：

在这个环节我们要做以下几件事：

• 将刚刚生成好的record文件放到对应目录下
• 下载SSD Mobile Net V1模型文件
• 配置book.pbtxt文件和book.config文件

放置record文件和SSD Mobile Net V1模型

为了方便我直接将models/research/object_detection/test_data下的目录清空，放置迁移训练的文件。

首先我们下载SSD Mobile Net V1模型文件：

我们下载第一个ssd_mobilenet_v1_coco模型即可，下载完毕后，我们解压下载的模型压缩包文件，并将模型相关的文件放在test_data的model目录下。并将我们刚刚生成的record文件放置在test_data目录下。

完成pbtxt和config配置文件

我们在test_data目录下，新建一个book.pbtxt文件，并完成配置内容:

item {
  id: 1
  name: 'book'
}

由于我们只有一个标签，我们就直接配置一个id值为1，name为book的item对象。

由于我们使用SSD Mobile Net V1模型进行迁移学习，因此我们到sampleconfigs目录下复制一份ssd_mobilenet_v1_coco.config文件并重命名为book.config文件。

接着我们修改book.config中的配置文件：

将num_classes修改为当前的标签数量：

由于我们只有一个book标签，因此修改成1即可。

修改所有PATH_TO_BE_CONFIGURED的路径：

我们将此处的模型文件地址设置成testdata/model/model.ckpt的全路径地址。

我们将train_input_reader的input_path设置成mask_train.record的全路径地址；将label_map_path设置成book.pbtxt的全路径地址；将eval_input_reader的input_path设置成mask_test.record的全路径地址。

到目前为止我们所有配置都已经完成啦。接下来就是激动人心的训练模型的时刻。

运行train.py文件训练模型

我们在终端中运行train.py文件，开始迁移学习、训练模型。

python3 train.py --logtostderr --train_dir=./test_data/training/ --pipeline_config_path=./test_data/book.config

其中train_dir为我们训练后的模型存放的目录，pipeline_config_path为我们book.config文件所在的相对路径。

运行命令后，我们可以看到模型在进行一步一步的训练：

并在/test_data/training目录下存放训练后的模型文件:

将ckpt文件转换为pb文件

我们通过export_inference_graph.py文件，将训练好的模型转换为pb格式的文件，这个文件格式在后面我们要用来转换为TensorFlow.js能够识别的文件格式。终于我们见到TensorFlow.js的影子啦。

我们执行命令，运行export_inference_graph.py文件：

python export_inference_graph.py --input_type image_tensor --pipeline_config_path ./test_data/book.config --trained_checkpoint_prefix ./test_data/training/model.ckpt-1989 --output_directory ./test_data/training/book_model_test

其中pipeline_config_path为book.config的相对文件路径，trained_checkpoint_prefix为模型文件的路径，例如我们选择训练了1989步的模型文件，output_directory为我们输出pb文件的目标目录。

运行完后，我们可以看到一个生成了book_model_test目录：

将pb文件转换为TensorFlowJs模型

首先我们需要依赖TensorFlowjs的依赖包

pip install tensorflowjs

然后通过命令行转换刚刚生成的pb文件

tensorflowjs_converter --input_format=tf_saved_model --output_node_names='detection_boxes,detection_classes,detection_features,detection_multiclass_scores,detection_scores,num_detections,raw_detection_boxes,raw_detection_scores' --saved_model_tags=serve --output_format=tfjs_graph_model ./saved_model ./web_model

其中我们设置最后两个参数，即saved_model的目录与TensorFlow.js识别模型的输出目录。

运行结束后，我们可以看到一个新生成的web_model目录，其中包括了我们迁移学习训练后的模型。

到这里，模型训练的阶段终于结束了。

在Vue中运行模型

准备工作

新建Vue项目，在Vue项目的public目录下放入我们训练好的模型，即web_model目录。

接着我们借助Tensorflow.js的依赖包，在package.json的dependencies中加入：

"@tensorflow/tfjs": "^1.7.2",
"@tensorflow/tfjs-core": "^1.7.2",
"@tensorflow/tfjs-node": "^1.7.2",

然后通过npm命令安装依赖包。

加载模型

在我们的JS代码部分引入TensorFlow的依赖包：

import * as tf from '@tensorflow/tfjs';
import {loadGraphModel} from '@tensorflow/tfjs-converter';

接着第一步，我们先加载模型文件中的model.json文件：

const MODEL_URL = process.env.BASE_URL+"web_model/model.json";
this.model = await loadGraphModel(MODEL_URL);

通过loadGraphModel方法，我们加载好训练的模型，再将模型对象打印出来：

随后，我们可以看到模型会输出一个长度为4的数组:

• detection_scores：表示识别对象模型的置信度，置信度越高，则代表模型认为对应区域识别为书本的可能性越高
• detection_classes：表示模型识别的区域对应的标签，例如在本案例中，识别出来的是book
• num_detections：表示模型识别出目标对象的个数
• detection_boxes：表示模型识别出来目标对象的区域，为一个长度为4的数组，分别是:[x_pos,y_pos,x_width,y_height] 。第一个位代表圈选区域左上角的x坐标，第二位代表圈选左上角的y坐标，第三位代表圈选区域的宽度，第四位代表圈选区域的长度。

模型识别

知道了输出值，我们就可以开始将图片输入到模型中，从而得到模型预测的结果：

const img = document.getElementById('img');
let modelres =await this.model.executeAsync(tf.browser.fromPixels(img).expandDims(0));

我们通过model.executeAsync方法，将图片输入到模型中，从而得到模型的输出值。

结果是我们前文提到的一个长度为4的数组。接着我们通过自定义方法，将得到的结果进行整理，从而输出一个想要的结果格式：

buildDetectedObjects:function(scores, threshold, imageWidth, imageHeight, boxes, classes, classesDir) {
          const detectionObjects = [];
          scores.forEach((score, i) => {
              if (score > threshold) {
                  const bbox = [];
                  const minY = boxes[i * 4] * imageHeight;
                  const minX = boxes[i * 4 + 1] * imageWidth;
                  const maxY = boxes[i * 4 + 2] * imageHeight;
                  const maxX = boxes[i * 4 + 3] * imageWidth;
                  bbox[0] = minX;
                  bbox[1] = minY;
                  bbox[2] = maxX - minX;
                  bbox[3] = maxY - minY;
                  detectionObjects.push({
                      class: classes[i],
                      label: classesDir[classes[i]].name,
                      score: score.toFixed(4),
                      bbox: bbox
                  });
              }
          });

          return detectionObjects
}

我们通过调用buildDetectedObjects来整理和返回最后的结果。

• scores：输入模型的detection_scores数组
• threshold：阈值，即结果score>threshold我们才会将对应的结果放入结果对象detectionObjects中
• imageWidth：图片的宽度
• imageHeight：图片的长度
• boxes：输入模型的detection_boxes数组
• classes：输入模型的detection_classes数组
• classesDir：即模型标签对象

调用buildDetectedObjects方法示例：

let classesDir = {
	1: {
		name: 'book',
		id: 1,
		}
	};
let res=this.buildDetectedObjects(modelres[0].dataSync(),0.20,img.width,img.height,modelres[3].dataSync(),modelres[1].dataSync(),classesDir);

我们通过modelres[0].dataSync()，来获取对应结果的数组对象，再输入到方法中，从而最终获得res结果对象。

最后我们通过Canvas的API，将图片根据bbox返回的数组对象，画出对应的区域即可。由于篇幅原因，就不赘述了，最终效果如下：

最后

本案例的模型存在一定的不足，由于训练时间较短，图书的封面类型众多，存在人像、风景图等等的样式，导致模型在识别过程中可能会将少部分的人脸、风景照等图片错误地识别成图书封面。各位小伙伴在训练自己模型的过程中可以考虑优化此问题。

当然，本案例的模型在识别非图书的场景会存在识别不准确的情况，一方面这是因为本案例从网络收集的图书样本具有一定局限性，而且图书的封面类型千差万别，存在人像、风景图等等的样式；另一方面因为本文在仅为起到抛砖引玉的作用，为各位前端小伙伴普及TensorFlow.js相关的知识，并提供训练自己的模型的解决方案，所以在收集样本和模型训练时间较短。感兴趣的小伙伴可以自己琢磨琢磨如何优化样本和在避免过拟合的情况下提高训练时长，从而提高模型对被识别物体的准确性。

我们写下本文仅为起到抛砖引玉的作用，为各位前端小伙伴普及TensorFlow.js相关知识并提供一种AI的解决方案。

愿每位热爱学习的开发者都能畅游书海，遇见更好的自己！