TensorFlow2 CPU训练及导出pb并转tflite

1、训练

1> 环境部署

（1）下载tensorflow的models文件，官网地址github，下载下来的文件名models-master.zip：解压文件，并重命名成models，我在ubuntu里面建立了一个名字为tensorflow的文件夹，并把models文件都放在这个tensorflow里面，当前的文件结构是：
/home/jg/tensorflow/models/，models下文件结构如下
在这里插入图片描述

（2）这里假设已经建安装完成所需要的模块。此时，在tensorflow目录下，建立文件夹models_training、my_model、dataset、pig、test、train等文件夹，结构如下：
在这里插入图片描述
xml2csv.py、csv2tfrecord.py文件在网上都能找到，就不列源码了。
需要注意的是pig_map.pbtxt文件是需要自己建立的文件，里面填写classes信息。

item {
    
    
  name: "pig"
  id: 1
  display_name: "pig"
}

解释下各个文件夹的详细内容，如下图：
在这里插入图片描述

（3）准备好要训练的图片，把要训练的图片按7:1的比例放到train和test文件夹下，然后通过labelImg软件标注目标，生成的xml存放到图片所在路径。
例如：
pig_1.jpg在train路径下，labelImg生成的pig_1.xml也放train路径下。
pig_58.jpg在test路径下，labelImg生成的pig_58.xml也放test路径下。
通过xml2csv.py，生成csv也存放在相应图片路径下。
通过csv2tfrecord.py，生成的record文件也放到图片路径下。

2> 预训练型模选择及配置文件修改

到tf2 detection models zoo官网tf2 zoo 下载预训练模型SSD MobileNet v2 320x320 （这里选了一个速度最快的模型），下载完成后解压文件到models_training下，文件名为ssd_mobilenet_v2_320x320_coco17_tpu-8。修改配置文件pipeline.config
（1）num_classes数量填写实际的数据，这里只训练一种，所以填的1

 ssd {
    
    
    num_classes: 1

（2）fine_tune_checkpoint填写的是预训练模型的ckeckpoint文件位置（这里需要注意填的是ckpt-0.data-00000-of-00001文件的文件名，不需要填后缀）
num_steps填写训练步数，根据需要填写，最好参考官方文档修改。
fine_tune_checkpoint_type填写detection

fine_tune_checkpoint: "/home/jg/tensorflow/model_training/ssd_mobilenet_v2_320x320_coco17_tpu-8/checkpoint/ckpt-0"
  num_steps: 6000
  startup_delay_steps: 0.0
  replicas_to_aggregate: 8
  max_number_of_boxes: 100
  unpad_groundtruth_tensors: false
  fine_tune_checkpoint_type: "detection"
  fine_tune_checkpoint_version: V2

（3）
train_input_reader
label_map_path填写上文中pig_map.pbtxt文件；input_path填写train.record路径

train_input_reader {
    
    
  label_map_path: "/home/jg/tensorflow/model_training/pig_map.pbtxt"
  tf_record_input_reader {
    
    
    input_path: "/home/jg/tensorflow/model_training/dataset/pig/train/train.record"
  }
}

eval_input_reader
label_map_path填写上文中pig_map.pbtxt文件；input_path填写test.record路径

eval_input_reader {
    
    
  label_map_path: "/home/jg/tensorflow/model_training/pig_map.pbtxt"
  shuffle: false
  num_epochs: 1
  tf_record_input_reader {
    
    
    input_path: "/home/jg/tensorflow/model_training/dataset/pig/test/test.record"
  }
}

（4）根据需要修改optimizer中的学习率learning_rate。试着用默认的学习率训练，若loss来回变化，或loss无比巨大，则减小学习率。若loss减小很慢，则增大学习率。

optimizer {
    
    
    momentum_optimizer {
    
    
      learning_rate {
    
    
        cosine_decay_learning_rate {
    
    
          learning_rate_base: 0.03
          total_steps: 6000
          warmup_learning_rate: 0.003
          warmup_steps: 1800
        }
      }
      momentum_optimizer_value: 0.8999999761581421
    }
    use_moving_average: false
  }

（5）batch_size根据实际调整，我只有cpu所以就设置成了1

train_config {
    
    
  batch_size: 1

3> 模型训练及导出

（1）训练程序model_main_tf2.py、
导出程序export_tflite_graph_tf2.py在/home/jg/tensorflow/models/research/object_detection/下可以找到
通过以下命令进行训练。为了方便我把文件拷到了models_trainning路径下。

python model_main_tf2.py --pipeline_config_path=/home/jg/tensorflow/model_training/ssd_mobilenet_v2_320x320_coco17_tpu-8/pipeline.config --model_dir=/home/jg/tensorflow/model_training/my_model/ --alsologtostderr

训练完成后，在my_model下会生成一堆文件
在这里插入图片描述
可以通过浏览器查看训练情况：

tensorboard --logdir=./

注意我在my_model路径下，所以用的./，实际用时候请注意–logdir是训练生成文件的路径。
命令执行后会出现个链接地址，把链接地址复制到浏览器就可以查看训练情况了。如下图（图是后面补的问题训练图，仅参考）：
在这里插入图片描述

（2）导出文件

python export_tflite_graph_tf2.py     --pipeline_config_path ssd_mobilenet_v2_320x320_coco17_tpu-8/pipeline.config     --trained_checkpoint_dir my_model/ --output_directory ./exp

执行导出命令后，会生成一个./exp的文件夹，文件夹内包含生成的pb文件
在这里插入图片描述
（3）转换成tflite文件
拷贝pig_map.pbtxt到saved_model下，然后执行如下命令

tflite_convert --output_file=./saved.tflite --saved_model_dir=/home/jg/tensorflow/model_training/exp/saved_model

成功生成saved.tflite文件，这个时候就可以把saved.tflite文件放到树莓派上跑预测了。因为.tflite是要搭配一个labelmap.txt使用的，可以自己新建一个labelmap.txt里面的内容填写classes的类名。
在这里插入图片描述
最后：因为篇幅有限，剩下的后面继续…

本文参考链接：
https://blog.csdn.net/weixin_44238733/article/details/109861924
特别鸣谢该文博主，如有侵权，联系删。