使用TensorFlow一步步进行目标检测(4)

本文翻译自Medium上的文章：Step by Step TensorFlow Object Detection API Tutorial — Part 4: Training the Model，原文地址：https://medium.com/@WuStangDan/step-by-step-tensorflow-object-detection-api-tutorial-part-4-training-the-model-68a9e5d5a333

到目前为止，我们已选择了预训练模型，找到了现成数据集，也创建了自己的数据集，并将其转换为TFRecord文件。接下来，我们开始训练模型。

模型配置文件

如果您之前有转移学习的经历，那么自本教程的第2部分以来，您可能会遇到一个问题，如何修改设计用于90个COCO数据集类别的预训练模型，以处理我的新数据集的X个类别？要达到这一目的，在调用目标检测API之前，您必须删除网络的最后90个神经元分类层并将其替换为新层。TensorFlow中有一个类似的示例，如下所示：

# Assume fc_2nd_last is the 2nd_last fully connected layer in your network and nb_classes is the number of classes in your new dataset.
shape = (fc_2nd_last.get_shape().as_list()[-1], nb_classes)
fc_last_W = tf.Variable(tf.truncated_normal(shape, stddev=1e-2))
fc_last_b = tf.Variable(tf.zeros(nb_classes))
logits = tf.nn.xw_plus_b(fc_2nd_last, fc_last_W, fc_last_b)

要使用目标检测API实现此目的，您需要做的就是修改模型配置文件中的一行。在所克隆的TensorFlow模型库的位置，导航到object_detection/samples/configs文件夹，在此文件夹中，您可以找到所有预训练模型的配置文件。

复制您所选模型的配置文件，并将其移动到新文件夹，您将在该文件夹中执行所有训练。在这个新文件夹中，创建一个名为data的文件夹，并将TFRecord文件移动到其中。创建另一个名为models的文件夹，并将您选择的预训练模型的.ckpt(检查点)文件(其中3个)移动到此文件夹中。回想一下，model_detection_zoo.md包含每个预训练模型的下载链接，每个模型的下载内容不仅包含.pb文件(我们在第1部分用在jupyter笔记本中)，还包含.ckpt文件。在models文件夹中创建另一个名为train的文件夹。

修改配置文件

使用文本编辑器打开新移动的配置文件，在最开始的一行将类别的数量更改为数据集中类别的数量。接下来，将fine_tune_checkpoint路径更改为指向model.ckpt文件。如果您遵循我所建议的模型结构，修改如下：

fine_tune_checkpoint: "models/model.ckpt"

参数num_steps决定在完成之前将运行的训练步数。这个数字实际上取决于数据集的大小以及许多其他因素(包括您愿意让模型训练多长时间)。一旦启动训练，我建议您看看每个训练步骤花了多长时间并相应地调整num_steps。

接下来，您需要更改训练和评估数据集的input_path和label_map_path。input_path指向TFRecord文件。在为label_map_path设置路径之前，我们需要创建它应该指向的文件。它查找的是一个.pbtxt文件，其中包含数据集的每个标签的ID和名称。您可以按照以下格式在文本文件中创建它。

item {
  id: 1
  name: 'Green'
} 
item {
  id: 2
  name: 'Red'
}

确保以id: 1而不是0开头。我建议将此文件放在数据文件夹中。最后将num_examples设置为您拥有的评估样本数。

训练

进入到object_detection文件夹并将train.py复制到新创建的训练文件夹中。要开始训练，只需打开终端窗口，导航到此文件夹(确保已按照第1部分中的安装说明操作)并输入命令行:

python train.py --logtostderr --train_dir=./models/train --pipeline_config_path=rfcn_resnet101_coco.config

其中pipline_config_path指向您的配置文件。训练现在开始了。请注意，根据您的系统，训练可能需要几分钟才能开始，所以如果它没有崩溃或停止，请给它一些时间。

如果因内存不足而导致训练失败，您可以尝试多种解决方案。首先尝试添加参数

batch_queue_capacity: 2
prefetch_queue_capacity: 2

到train_config部分的配置文件。例如，将这两行放到gradient_clipping_by_norm和fine_tune_checkpoint之间，上面的数字2只是开始训练的起步值。这些值的默认值分别为8和10，增加这些值有助于加快训练速度。

就这样，您现在已经开始训练模型了！如果您想更好地了解训练的进展情况，请查看TensorBoard。

在下一篇文章中，我将向您展示如何保存训练好的模型并将其部署到项目中！