edit_pipeline.py

# Edit Pipeline config to load in our new tfrecord that we just created and add quantization aware training.
import tensorflow as tf
from google.protobuf import text_format
from object_detection.protos import pipeline_pb2

# Hack to find out if you have colab pro or not :)
gpu_info = !nvidia-smi
gpu_info = '\n'.join(gpu_info)
print(gpu_info)
gpu_name = !nvidia-smi --query-gpu=gpu_name --format=csv
# You get Tesla T4 with free colab and Tesla P100-PCIe with colab pro.
colab_pro = False if 'T4' in gpu_name else True

pipeline = pipeline_pb2.TrainEvalPipelineConfig()                                                                                                                                                                                                          
config_path = '/content/models/research/object_detection/samples/configs/ssdlite_mobiledet_edgetpu_320x320_coco_sync_4x4.config'
with tf.gfile.GFile(config_path, "r") as f:                                                                                                                                                                                                                     
    proto_str = f.read()                                                                                                                                                                                                                                          
    text_format.Merge(proto_str, pipeline)

pipeline.train_input_reader.tf_record_input_reader.input_path[:] = ['/content/dataset/pet_faces_train.record-?????-of-00010']
pipeline.train_input_reader.label_map_path = '/content/dataset/pet_label_map.pbtxt'
pipeline.eval_input_reader[0].tf_record_input_reader.input_path[:] = ['/content/dataset/pet_faces_val.record-?????-of-00010']
pipeline.eval_input_reader[0].label_map_path = '/content/dataset/pet_label_map.pbtxt'
pipeline.train_config.fine_tune_checkpoint = '/content/pretrained_model/ssdlite_mobiledet_edgetpu_320x320_coco_2020_05_19/fp32/model.ckpt'
pipeline.train_config.batch_size = 64 if colab_pro else 32 # Smaller batch size on free gpu to avoid OOM Killer
pipeline.train_config.num_steps = 25000 if colab_pro else 10000 # Less steps with free gpu but 10k should be good enough
pipeline.model.ssd.num_classes = 2
# Enable ssdlite, this should already be enabled in the config we downloaded, but this is just to make sure.
pipeline.model.ssd.box_predictor.convolutional_box_predictor.kernel_size = 3
pipeline.model.ssd.box_predictor.convolutional_box_predictor.use_depthwise = True
pipeline.model.ssd.feature_extractor.use_depthwise = True
# Quantization Aware Training
pipeline.graph_rewriter.quantization.delay = 0
pipeline.graph_rewriter.quantization.weight_bits = 8
pipeline.graph_rewriter.quantization.activation_bits = 8

config_text = text_format.MessageToString(pipeline)                                                                                                                                                                                                        
with tf.gfile.Open(config_path, "wb") as f:                                                                                                                                                                                                                       
    f.write(config_text)

# This is out config after modifying.
!cat /content/models/research/object_detection/samples/configs/ssdlite_mobiledet_edgetpu_320x320_coco_sync_4x4.config