def __init__(self): self.input_size = cfg.TEST.INPUT_SIZE self.anchor_per_scale = cfg.YOLO.ANCHOR_PER_SCALE self.classes = {0: 'front', 1: 'rear'} self.num_classes = len(self.classes) self.anchors = np.array(utils.get_anchors(cfg.YOLO.ANCHORS)) self.score_threshold = cfg.TEST.SCORE_THRESHOLD self.iou_threshold = cfg.TEST.IOU_THRESHOLD self.moving_ave_decay = cfg.YOLO.MOVING_AVE_DECAY self.annotation_path = cfg.TEST.ANNOT_PATH self.weight_file = cfg.TEST.WEIGHT_FILE self.write_image = cfg.TEST.WRITE_IMAGE self.write_image_path = cfg.TEST.WRITE_IMAGE_PATH self.test_video_path = cfg.TEST.TEST_VIDEO_PATH self.predict_video_path = cfg.TEST.PREDICT_VIDEO_PATH with tf.name_scope('input'): self.input_data = tf.placeholder(dtype=tf.float32, shape=[1, 544, 544, 3], name='input_data') self.trainable = tf.placeholder(dtype=tf.bool, name='trainable') model = YOLOV3(self.input_data, self.trainable, 1) self.pred_sbbox, self.pred_mbbox, self.pred_lbbox = model.pred_sbbox, model.pred_mbbox, model.pred_lbbox self.conv_sbbox, self.conv_mbbox, self.conv_lbbox = model.conv_sbbox, model.conv_mbbox, model.conv_lbbox with tf.name_scope('ema'): ema_obj = tf.train.ExponentialMovingAverage(self.moving_ave_decay) self.sess = tf.Session(config=tf.ConfigProto( allow_soft_placement=True)) self.saver = tf.train.Saver() self.saver.restore(self.sess, self.weight_file)
def create_pb(ckpt_file): pb_file = "./yolov3_coco.pb" output_node_names = [ "input/input_data", "pred_sbbox/concat_2", "pred_mbbox/concat_2", "pred_lbbox/concat_2" ] with tf.name_scope('input'): input_data = tf.placeholder(dtype=tf.float32, name='input_data') model = YOLOV3(input_data, trainable=False) print(model.conv_sbbox, model.conv_mbbox, model.conv_lbbox) sess = tf.Session(config=tf.ConfigProto(allow_soft_placement=True)) saver = tf.train.Saver() saver.restore(sess, ckpt_file) converted_graph_def = tf.graph_util.convert_variables_to_constants( sess, input_graph_def=sess.graph.as_graph_def(), output_node_names=output_node_names) with tf.gfile.GFile(pb_file, "wb") as f: f.write(converted_graph_def.SerializeToString())
def __init__(self): self.input_size = cfg.TEST.INPUT_SIZE self.classes = utils.read_class_names(cfg.YOLO.CLASSES) self.num_classes = len(self.classes) self.score_threshold = cfg.TEST.SCORE_THRESHOLD self.iou_threshold = cfg.TEST.IOU_THRESHOLD self.moving_ave_decay = cfg.YOLO.MOVING_AVE_DECAY self.weight_file = cfg.TEST.WEIGHT_FILE self.show_label = cfg.TEST.SHOW_LABEL with tf.name_scope('input'): self.input_data = tf.placeholder(dtype=tf.float32, name='input_data') self.trainable = tf.placeholder(dtype=tf.bool, name='trainable') model = YOLOV3(self.input_data, self.trainable) self.pred_sbbox, self.pred_mbbox, self.pred_lbbox = model.pred_sbbox, model.pred_mbbox, model.pred_lbbox with tf.name_scope('ema'): ema_obj = tf.train.ExponentialMovingAverage(self.moving_ave_decay) self.sess = tf.Session(config=tf.ConfigProto( allow_soft_placement=True)) self.saver = tf.train.Saver(ema_obj.variables_to_restore()) self.saver.restore(self.sess, self.weight_file)
def __init__(self): self.input_size = cfg.TEST.INPUT_SIZE self.anchor_per_scale = cfg.YOLO.ANCHOR_PER_SCALE self.classes = utils.read_class_names(cfg.YOLO.CLASSES) self.num_classes = len(self.classes) self.anchors = np.array(utils.get_anchors(cfg.YOLO.ANCHORS)) self.score_threshold = cfg.TEST.SCORE_THRESHOLD self.iou_threshold = cfg.TEST.IOU_THRESHOLD self.moving_ave_decay = cfg.YOLO.MOVING_AVE_DECAY self.annotation_path = cfg.TEST.ANNOT_PATH self.weight_file = cfg.TEST.WEIGHT_FILE self.write_image = cfg.TEST.WRITE_IMAGE self.write_image_path = cfg.TEST.WRITE_IMAGE_PATH self.show_label = cfg.TEST.SHOW_LABEL with tf.name_scope('input'): self.input_data = tf.placeholder(dtype=tf.float32, name='input_data') self.trainable = tf.placeholder(dtype=tf.bool, name='trainable') model = YOLOV3() #(self.input_data, self.trainable) sqback = backbone.squeezenet('sqz_full.mat') sqback.forward(imgs=self.input_data, trainable=False) model.forward(sqback, False) self.pred_sbbox, self.pred_mbbox, self.pred_lbbox = model.pred_sbbox, model.pred_mbbox, model.pred_lbbox with tf.name_scope('ema'): ema_obj = tf.train.ExponentialMovingAverage(self.moving_ave_decay) self.sess = tf.Session(config=tf.ConfigProto(allow_soft_placement=True)) self.saver = tf.train.Saver(ema_obj.variables_to_restore()) self.saver.restore(self.sess, self.weight_file)
def main(export_path='/tmp/yolo/1'): with tf.name_scope('input'): input_data = tf.placeholder(dtype=tf.float32, name='input_data', shape=(1, 416, 416, 3)) model = YOLOV3(input_data, trainable=False) print(model.conv_sbbox, model.conv_mbbox, model.conv_lbbox) sess = tf.Session(config=tf.ConfigProto(allow_soft_placement=True)) saver = tf.train.Saver() saver.restore(sess, ckpt_file) converted_graph_def = tf.graph_util.convert_variables_to_constants( sess, input_graph_def=sess.graph.as_graph_def(), output_node_names=output_node_names) tf.saved_model.simple_save( sess, export_path, inputs={'input_data': input_data}, outputs={ 'conv_sbbox': tf.get_default_graph().get_tensor_by_name('pred_sbbox/concat_2:0'), 'conv_mbbox': tf.get_default_graph().get_tensor_by_name('pred_mbbox/concat_2:0'), 'conv_lbbox': tf.get_default_graph().get_tensor_by_name('pred_lbbox/concat_2:0'), }, )
def main(_): with tf.name_scope('input'): input_data = tf.placeholder(dtype=tf.float32, name='input_data') serialized_tf_example = tf.placeholder(tf.string, name='tf_example') feature_configs = { 'x': tf.FixedLenFeature(shape=[], dtype=tf.float32), } tf_example = tf.parse_example(serialized_tf_example, feature_configs) tf_example['x'] = tf.reshape(tf_example['x'], (1, 416, 416, 3)) input_tensor = tf.identity(tf_example['x'], name='x') # use tf.identity() to assign name model = YOLOV3(input_data, trainable=False) print(model.conv_sbbox, model.conv_mbbox, model.conv_lbbox) sess = tf.Session(config=tf.ConfigProto(allow_soft_placement=True)) saver = tf.train.Saver() saver.restore(sess, ckpt_file) # Create SavedModelBuilder class # defines where the model will be exported export_path_base = FLAGS.export_model_dir export_path = os.path.join(tf.compat.as_bytes(export_path_base), tf.compat.as_bytes(str(FLAGS.model_version))) print('Exporting trained model to', export_path) builder = tf.saved_model.builder.SavedModelBuilder(export_path) tensor_info_input = tf.saved_model.utils.build_tensor_info(input_tensor) tensor_conv_sbbox_output = tf.saved_model.utils.build_tensor_info( model.conv_sbbox) tensor_conv_mbbox_output = tf.saved_model.utils.build_tensor_info( model.conv_mbbox) tensor_conv_lbbox_output = tf.saved_model.utils.build_tensor_info( model.conv_lbbox) prediction_signature = ( tf.saved_model.signature_def_utils.build_signature_def( inputs={'images': tensor_info_input}, outputs={ 'conv_sbbox': tensor_conv_sbbox_output, 'conv_mbbox': tensor_conv_mbbox_output, 'conv_lbbox': tensor_conv_lbbox_output }, method_name=tf.saved_model.signature_constants.PREDICT_METHOD_NAME) ) builder.add_meta_graph_and_variables( sess, [tf.saved_model.tag_constants.SERVING], signature_def_map={ 'predict_images': prediction_signature, }) # export the model builder.save(as_text=True) print('Done exporting!')
def main(argv=None): input_data = tf.placeholder(dtype=tf.float32, shape=(None, 608, 608, 3), name='inputs') model = YOLOV3(input_data, trainable=False, tiny=FLAGS.tiny) load_ops = load_weights(tf.global_variables(), FLAGS.weights_file) with tf.Session() as sess: sess.run(load_ops) freeze_graph(sess, FLAGS.output_graph)
def __init__(self): self.input_size = 416 self.anchor_per_scale = 3 self.path1 = r"./data/classes/antenna.names" #修改 self.classes = utils.read_class_names(self.path1) self.num_classes = len(self.classes) self.path2 = r"./data/anchors/basline_anchors.txt" self.anchors = np.array(utils.get_anchors(self.path2)) self.score_threshold = 0.4 self.iou_threshold = 0.5 self.moving_ave_decay = 0.9995 self.path3 = "./data/dataset/antenna_train.txt" self.annotation_path = self.path3 self.path4 = r'checkpoint1/yolov3_test_loss=0.9668.ckpt-705' # self.path4 = r'checkpoint/model/yolov3_test_loss=49.6633.ckpt-56' self.weight_file = self.path4 self.write_image = True self.path5 = r"./data/detection1/" self.write_image_path = self.path5 # 是否将图片的预测结果保存 self.show_label = True with tf.name_scope('input'): self.input_data = tf.placeholder(dtype=tf.float32, name='input_data') self.trainable = tf.placeholder(dtype=tf.bool, name='trainable') print("in", self.input_data.shape) model = YOLOV3(self.input_data, self.trainable) self.pred_sbbox, self.pred_mbbox, self.pred_lbbox = model.pred_sbbox, model.pred_mbbox, model.pred_lbbox print(1 + 2) with tf.name_scope('ema'): ema_obj = tf.train.ExponentialMovingAverage(self.moving_ave_decay) self.sess = tf.Session(config=tf.ConfigProto( allow_soft_placement=True)) self.saver = tf.train.Saver(ema_obj.variables_to_restore()) self.saver.restore(self.sess, self.weight_file)
def inference(input_data): model = YOLOV3(False, opts) # construct model # we will put whole network on one ipu layers = [] # build layer functions for backbone and upsample layers.extend(model.build_backbone()) # last layer of darknet53 is classification layer, so it have 52 conv layers assert len(layers) == 52 layers.extend(model.build_upsample()) # there is 25 conv layers if we count upsmaple as a conv layer assert len(layers) == 52 + 25 # decoding layer and loss layer is always put on last IPU layers.append(model.decode_boxes) # reuse stages_constructor so we don't need to pass params by hand network_func = stages_constructor( [layers], ["input_data", "nums"], ["pred_sbbox", "pred_mbbox", "pred_lbbox", "nums"])[0] return network_func(input_data)
def __init__(self, opts): """Create a training class The constructor init all needed parameters """ self.opts = opts self.learn_rate_init = opts["train"]["learn_rate_init"] self.learn_rate_end = opts["train"]["learn_rate_end"] self.epochs = opts["train"]["epochs"] self.warmup_epochs = opts["train"]["warmup_epochs"] self.initial_weight = opts["train"]["initial_weight"] self.moving_avg_decay = opts["yolo"]["moving_avg_decay"] self.trainset = Dataset("train", self.opts) self.steps_per_epoch = len( self.trainset) / opts["train"]["total_replicas"] self.precision = tf.float16 if opts["yolo"][ "precision"] == "fp16" else tf.float32 self.model = YOLOV3(opts["train"]["bn_trainable"], opts) self.batch_size = opts["train"]["batch_size"] self.data_threads_number = opts["train"]["data_threads_number"] self.loss_scaling = opts["train"]["loss_scaling"] self.repeat_count = opts["train"]["repeat_count"] self.for_speed_test = opts["train"]["for_speed_test"]
def __init__(self): self.input_size = cfg.TEST.INPUT_SIZE self.anchor_per_scale = cfg.YOLO.ANCHOR_PER_SCALE self.classes = utils.read_class_names(cfg.YOLO.CLASSES) self.num_classes = len(self.classes) self.anchors = np.array(utils.get_anchors(cfg.YOLO.ANCHORS)) self.score_threshold = cfg.TEST.SCORE_THRESHOLD self.iou_threshold = cfg.TEST.IOU_THRESHOLD self.moving_ave_decay = cfg.YOLO.MOVING_AVE_DECAY self.annotation_path = cfg.TEST.ANNOT_PATH self.weight_file = cfg.TEST.WEIGHT_FILE self.write_image = cfg.TEST.WRITE_IMAGE self.write_image_path = cfg.TEST.WRITE_IMAGE_PATH self.show_label = cfg.TEST.SHOW_LABEL # model_dir = os.path.abspath(os.path.dirname(self.weight_file)) # self.weight_file = os.path.join(model_dir, 'yolov3_test_loss=8.1157.ckpt-24') # self.weight_file = tf.train.latest_checkpoint(model_dir) # print(self.weight_file) with tf.name_scope('input'): self.input_data = tf.placeholder(dtype=tf.float32, name='input_data') self.trainable = tf.placeholder(dtype=tf.bool, name='trainable') model = YOLOV3(self.input_data, self.trainable) self.pred_sbbox, self.pred_mbbox, self.pred_lbbox = model.pred_sbbox, model.pred_mbbox, model.pred_lbbox # with tf.name_scope('ema'): # ema_obj = tf.train.ExponentialMovingAverage(self.moving_ave_decay) self.sess = tf.Session(config=tf.ConfigProto( allow_soft_placement=True)) # self.sess = tf.Session() # self.saver = tf.train.Saver(ema_obj.variables_to_restore()) self.saver = tf.train.Saver(tf.global_variables()) self.saver.restore(self.sess, self.weight_file)
def __init__(self): self.anchor_per_scale = cfg.YOLO.ANCHOR_PER_SCALE self.classes = utils.read_class_names(cfg.YOLO.CLASSES) self.num_classes = len(self.classes) self.learn_rate_init = cfg.TRAIN.LEARN_RATE_INIT self.learn_rate_end = cfg.TRAIN.LEARN_RATE_END self.first_stage_epochs = cfg.TRAIN.FISRT_STAGE_EPOCHS self.second_stage_epochs = cfg.TRAIN.SECOND_STAGE_EPOCHS self.warmup_periods = cfg.TRAIN.WARMUP_EPOCHS self.initial_weight = cfg.TRAIN.INITIAL_WEIGHT self.time = time.strftime('%Y-%m-%d-%H-%M-%S', time.localtime(time.time())) self.moving_ave_decay = cfg.YOLO.MOVING_AVE_DECAY self.max_bbox_per_scale = 100 self.train_logdir = "./data/log/train" self.trainset = Dataset('train') self.testset = Dataset('test') self.steps_per_period = len(self.trainset) self.sess = tf.Session(config=tf.ConfigProto( allow_soft_placement=True)) with tf.name_scope('define_input'): self.input_data = tf.placeholder(dtype=tf.float32, name='input_data') self.label_sbbox = tf.placeholder(dtype=tf.float32, name='label_sbbox') self.label_mbbox = tf.placeholder(dtype=tf.float32, name='label_mbbox') self.label_lbbox = tf.placeholder(dtype=tf.float32, name='label_lbbox') self.true_sbboxes = tf.placeholder(dtype=tf.float32, name='sbboxes') self.true_mbboxes = tf.placeholder(dtype=tf.float32, name='mbboxes') self.true_lbboxes = tf.placeholder(dtype=tf.float32, name='lbboxes') self.input_image_with_bboxes = tf.placeholder(dtype=tf.float32, name='input_image') self.trainable = tf.placeholder(dtype=tf.bool, name='training') with tf.name_scope("define_loss"): self.model = YOLOV3(self.input_data, trainable=True, tiny=TINY) self.net_var = tf.global_variables() self.giou_loss, self.conf_loss, self.prob_loss = self.model.compute_loss( self.label_sbbox, self.label_mbbox, self.label_lbbox, self.true_sbboxes, self.true_mbboxes, self.true_lbboxes) self.loss = self.giou_loss + 2 * self.conf_loss + 5 * self.prob_loss # 增加权重,调整了一下3种loss的影响力度 with tf.name_scope('learn_rate'): self.global_step = tf.Variable(1.0, dtype=tf.float64, trainable=False, name='global_step') warmup_steps = tf.constant(self.warmup_periods * self.steps_per_period, dtype=tf.float64, name='warmup_steps') train_steps = tf.constant( (self.first_stage_epochs + self.second_stage_epochs) * self.steps_per_period, dtype=tf.float64, name='train_steps') # tf.cond : if pred, return fn1 ,else, return fn2 ; warm_up对应原版的burn_in(power=1),learn rate随着epoch线性增长 , # 之后的learn rate更新策略是cos曲线下降 self.learn_rate = tf.cond( pred=self.global_step < warmup_steps, true_fn=lambda: self.global_step / warmup_steps * self. learn_rate_init, false_fn=lambda: self.learn_rate_end + 0.5 * (self.learn_rate_init - self.learn_rate_end) * (1 + tf.cos( (self.global_step - warmup_steps) / (train_steps - warmup_steps) * np.pi))) global_step_update = tf.assign_add(self.global_step, 1.0) with tf.name_scope("define_weight_decay"): # ExponentialMovingAverage 滑动平均的方法更新参数 moving_ave = tf.train.ExponentialMovingAverage( self.moving_ave_decay).apply(tf.trainable_variables()) with tf.name_scope("define_first_stage_train"): # 第一阶段训练:仅仅训练三个分支的最后卷积层 self.first_stage_trainable_var_list = [] for var in tf.trainable_variables(): var_name = var.op.name var_name_mess = str(var_name).split('/') if var_name_mess[0] in [ 'conv_sbbox', 'conv_mbbox', 'conv_lbbox' ]: self.first_stage_trainable_var_list.append(var) first_stage_optimizer = tf.train.AdamOptimizer( self.learn_rate).minimize( self.loss, var_list=self.first_stage_trainable_var_list) # tf.control_dependencies(), 指定某些操作执行的依赖关系,后面的操作要在()的操作之后执行 with tf.control_dependencies( tf.get_collection(tf.GraphKeys.UPDATE_OPS)): with tf.control_dependencies( [first_stage_optimizer, global_step_update]): with tf.control_dependencies([moving_ave]): self.train_op_with_frozen_variables = tf.no_op() with tf.name_scope("define_second_stage_train"): # 第二阶段训练:训练所有的层,其实也就是 finetunning second_stage_trainable_var_list = tf.trainable_variables() second_stage_optimizer = tf.train.AdamOptimizer( self.learn_rate).minimize( self.loss, var_list=second_stage_trainable_var_list) with tf.control_dependencies( tf.get_collection(tf.GraphKeys.UPDATE_OPS)): with tf.control_dependencies( [second_stage_optimizer, global_step_update]): with tf.control_dependencies([moving_ave]): self.train_op_with_all_variables = tf.no_op() with tf.name_scope('loader_and_saver'): self.loader = tf.train.Saver(self.net_var) self.saver = tf.train.Saver(tf.global_variables(), max_to_keep=10) with tf.name_scope('learning_rate'): tf.summary.scalar("learn_rate", self.learn_rate) with tf.name_scope('loss'): tf.summary.scalar("giou_loss", self.giou_loss) tf.summary.scalar("conf_loss", self.conf_loss) tf.summary.scalar("prob_loss", self.prob_loss) tf.summary.scalar("total_loss", self.loss) with tf.name_scope('images'): tf.summary.image("input_data", self.input_data) tf.summary.image("input_image", self.input_image_with_bboxes) logdir = "./data/log/" if os.path.exists(logdir): shutil.rmtree(logdir) os.mkdir(logdir) self.write_op = tf.summary.merge_all() loss = tf.summary.scalar("test_loss", self.loss) test_img = tf.summary.image("input_data", self.input_data) self.write_op_test = tf.summary.merge([loss, test_img]) self.summary_writer = tf.summary.FileWriter(logdir, graph=self.sess.graph)
# ================================================================ import tensorflow as tf from core.yolov3 import YOLOV3 TINY = True INPUT_SIZE = 608 pb_file = "./pb/testtiny.pb" ckpt_file = "./checkpoint/yolov3_test_loss=3.5116.ckpt-200" # output_node_names = ["input/input_data", "pred_sbbox/concat_2", "pred_mbbox/concat_2", "pred_lbbox/concat_2"] output_node_names = ['inputs', 'output_boxes'] input_data = tf.placeholder(dtype=tf.float32, shape=(None, INPUT_SIZE, INPUT_SIZE, 3), name='inputs') model = YOLOV3(input_data, trainable=False, tiny=TINY) # print(model.conv_sbbox, model.conv_mbbox, model.conv_lbbox) sess = tf.Session(config=tf.ConfigProto(allow_soft_placement=True)) saver = tf.train.Saver() saver.restore(sess, ckpt_file) converted_graph_def = tf.graph_util.convert_variables_to_constants( sess, input_graph_def=sess.graph.as_graph_def(), output_node_names=output_node_names) with tf.gfile.GFile(pb_file, "wb") as f: f.write(converted_graph_def.SerializeToString())
def __init__(self): # 从config文件获取到一些变量 self.anchor_per_scale = cfg.YOLO.ANCHOR_PER_SCALE self.classes = utils.read_class_names(cfg.YOLO.CLASSES) self.num_classes = len(self.classes) self.learn_rate_init = cfg.TRAIN.LEARN_RATE_INIT self.learn_rate_end = cfg.TRAIN.LEARN_RATE_END self.first_stage_epochs = cfg.TRAIN.FISRT_STAGE_EPOCHS self.second_stage_epochs = cfg.TRAIN.SECOND_STAGE_EPOCHS self.warmup_periods = cfg.TRAIN.WARMUP_EPOCHS self.initial_weight = cfg.TRAIN.INITIAL_WEIGHT self.time = time.strftime('%Y-%m-%d-%H-%M-%S', time.localtime(time.time())) self.moving_ave_decay = cfg.YOLO.MOVING_AVE_DECAY self.max_bbox_per_scale = 150 self.train_logdir = "./data/log/train" # 日志保存地址 self.trainset = Dataset('train') self.testset = Dataset('test') self.steps_per_period = len(self.trainset) self.sess = tf.Session(config=tf.ConfigProto( allow_soft_placement=True)) with tf.name_scope('define_input'): # 定义输入层 self.input_data = tf.placeholder(dtype=tf.float32, name='input_data') self.label_sbbox = tf.placeholder(dtype=tf.float32, name='label_sbbox') self.label_mbbox = tf.placeholder(dtype=tf.float32, name='label_mbbox') self.label_lbbox = tf.placeholder(dtype=tf.float32, name='label_lbbox') self.true_sbboxes = tf.placeholder(dtype=tf.float32, name='sbboxes') self.true_mbboxes = tf.placeholder(dtype=tf.float32, name='mbboxes') self.true_lbboxes = tf.placeholder(dtype=tf.float32, name='lbboxes') self.trainable = tf.placeholder(dtype=tf.bool, name='training') with tf.name_scope("define_loss"): # 定义损失函数 self.model = YOLOV3(self.input_data, self.trainable) self.net_var = tf.global_variables() self.giou_loss, self.conf_loss, self.prob_loss = self.model.compute_loss( self.label_sbbox, self.label_mbbox, self.label_lbbox, self.true_sbboxes, self.true_mbboxes, self.true_lbboxes) self.loss = self.giou_loss + self.conf_loss + self.prob_loss with tf.name_scope('learn_rate'): # 定义学习率 self.global_step = tf.Variable(1.0, dtype=tf.float64, trainable=False, name='global_step') warmup_steps = tf.constant(self.warmup_periods * self.steps_per_period, dtype=tf.float64, name='warmup_steps') train_steps = tf.constant( (self.first_stage_epochs + self.second_stage_epochs) * self.steps_per_period, dtype=tf.float64, name='train_steps') self.learn_rate = tf.cond( pred=self.global_step < warmup_steps, true_fn=lambda: self.global_step / warmup_steps * self. learn_rate_init, false_fn=lambda: self.learn_rate_end + 0.5 * (self.learn_rate_init - self.learn_rate_end) * (1 + tf.cos( (self.global_step - warmup_steps) / (train_steps - warmup_steps) * np.pi))) global_step_update = tf.assign_add(self.global_step, 1.0) ''' warmup_steps作用: 神经网络在刚开始训练的过程中容易出现loss=NaN的情况,为了尽量避免这个情况,因此初始的学习率设置得很低 但是这又使得训练速度变慢了。因此,采用逐渐增大的学习率,从而达到既可以尽量避免出现nan,又可以等训练过程稳定了再增大训练速度的目的。 ''' with tf.name_scope( "define_weight_decay"): # 指数平滑,可以让算法在最后不那么震荡,结果更有鲁棒性 moving_ave = tf.train.ExponentialMovingAverage( self.moving_ave_decay).apply(tf.trainable_variables()) # 指定需要恢复的参数。层等信息, 位置提前,减少模型体积。 with tf.name_scope('loader_and_saver'): variables_to_restore = [ v for v in self.net_var if v.name.split('/')[0] not in ['conv_sbbox', 'conv_mbbox', 'conv_lbbox'] ] self.loader = tf.train.Saver(variables_to_restore) self.saver = tf.train.Saver(tf.global_variables(), max_to_keep=1) with tf.name_scope("define_first_stage_train"): # 第一阶段训练,只训练指定层 self.first_stage_trainable_var_list = [] for var in tf.trainable_variables(): var_name = var.op.name var_name_mess = str(var_name).split('/') if var_name_mess[0] in [ 'conv_sbbox', 'conv_mbbox', 'conv_lbbox' ]: self.first_stage_trainable_var_list.append(var) first_stage_optimizer = tf.train.AdamOptimizer( self.learn_rate).minimize( self.loss, var_list=self.first_stage_trainable_var_list) with tf.control_dependencies( tf.get_collection(tf.GraphKeys.UPDATE_OPS)): with tf.control_dependencies( [first_stage_optimizer, global_step_update]): with tf.control_dependencies([moving_ave]): self.train_op_with_frozen_variables = tf.no_op() with tf.name_scope("define_second_stage_train"): # 第二阶段训练,释放所有层 second_stage_trainable_var_list = tf.trainable_variables() second_stage_optimizer = tf.train.AdamOptimizer( self.learn_rate).minimize( self.loss, var_list=second_stage_trainable_var_list) with tf.control_dependencies( tf.get_collection(tf.GraphKeys.UPDATE_OPS)): with tf.control_dependencies( [second_stage_optimizer, global_step_update]): with tf.control_dependencies([moving_ave]): self.train_op_with_all_variables = tf.no_op() with tf.name_scope('summary'): tf.summary.scalar("learn_rate", self.learn_rate) tf.summary.scalar("giou_loss", self.giou_loss) tf.summary.scalar("conf_loss", self.conf_loss) tf.summary.scalar("prob_loss", self.prob_loss) tf.summary.scalar("total_loss", self.loss) logdir = "./data/log/" # 日志保存地址 if os.path.exists(logdir): shutil.rmtree(logdir) os.mkdir(logdir) self.write_op = tf.summary.merge_all() self.summary_writer = tf.summary.FileWriter(logdir, graph=self.sess.graph)
def __init__(self): self.anchor_per_scale = cfg.YOLO.ANCHOR_PER_SCALE self.classes = utils.read_class_names( cfg.YOLO.CLASSES) #dict类型,ID---name self.num_classes = len(self.classes) #检测类别数量 self.learn_rate_init = cfg.TRAIN.LEARN_RATE_INIT self.learn_rate_end = cfg.TRAIN.LEARN_RATE_END self.first_stage_epochs = cfg.TRAIN.FISRT_STAGE_EPOCHS self.second_stage_epochs = cfg.TRAIN.SECOND_STAGE_EPOCHS self.warmup_periods = cfg.TRAIN.WARMUP_EPOCHS self.initial_weight = cfg.TRAIN.INITIAL_WEIGHT self.time = time.strftime('%Y-%m-%d-%H-%M-%S', time.localtime(time.time())) self.moving_ave_decay = cfg.YOLO.MOVING_AVE_DECAY #滑动平均时的decay值 self.max_bbox_per_scale = 150 #每个尺度上检测目标的数量 self.train_logdir = "./data/log/train" # 训练日志保存路径 self.trainset = Dataset('train') self.testset = Dataset('test') self.steps_per_period = len(self.trainset) self.config = tf.ConfigProto(allow_soft_placement=True) #self.config.gpu_options.per_process_gpu_memory_fraction = 0.2 # 占用40%显存 self.sess = tf.Session(config=self.config) with tf.name_scope('define_input'): self.input_data = tf.placeholder(dtype=tf.float32, name='input_data') self.label_sbbox = tf.placeholder(dtype=tf.float32, name='label_sbbox') self.label_mbbox = tf.placeholder(dtype=tf.float32, name='label_mbbox') self.label_lbbox = tf.placeholder(dtype=tf.float32, name='label_lbbox') self.true_sbboxes = tf.placeholder(dtype=tf.float32, name='sbboxes') self.true_mbboxes = tf.placeholder(dtype=tf.float32, name='mbboxes') self.true_lbboxes = tf.placeholder(dtype=tf.float32, name='lbboxes') self.trainable = tf.placeholder( dtype=tf.bool, name='training') # 占位符,对训练时,为True,验证时为False with tf.name_scope("define_loss"): self.model = YOLOV3(self.input_data, self.trainable) self.net_var = tf.global_variables() self.giou_loss, self.conf_loss, self.prob_loss = self.model.compute_loss( self.label_sbbox, self.label_mbbox, self.label_lbbox, self.true_sbboxes, self.true_mbboxes, self.true_lbboxes) self.loss = self.giou_loss + self.conf_loss + self.prob_loss # 损失函数 with tf.name_scope('learn_rate'): self.global_step = tf.Variable(1.0, dtype=tf.float64, trainable=False, name='global_step') warmup_steps = tf.constant(self.warmup_periods * self.steps_per_period, dtype=tf.float64, name='warmup_steps') train_steps = tf.constant( (self.first_stage_epochs + self.second_stage_epochs) * self.steps_per_period, dtype=tf.float64, name='train_steps') self.learn_rate = tf.cond( pred=self.global_step < warmup_steps, true_fn=lambda: self.global_step / warmup_steps * self. learn_rate_init, false_fn=lambda: self.learn_rate_end + 0.5 * (self.learn_rate_init - self.learn_rate_end) * (1 + tf.cos( (self.global_step - warmup_steps) / (train_steps - warmup_steps) * np.pi))) global_step_update = tf.assign_add(self.global_step, 1.0) with tf.name_scope("define_weight_decay"): moving_ave = tf.train.ExponentialMovingAverage( self.moving_ave_decay).apply( tf.trainable_variables()) #给模型中的变量创建滑动平均(滑动平均,作用于模型中的变量) # 第一阶段训练:仅仅训练三个分支的最后卷积层 with tf.name_scope("define_first_stage_train"): self.first_stage_trainable_var_list = [] for var in tf.trainable_variables(): var_name = var.op.name var_name_mess = str(var_name).split('/') if var_name_mess[0] in [ 'conv_sbbox', 'conv_mbbox', 'conv_lbbox' ]: self.first_stage_trainable_var_list.append(var) first_stage_optimizer = tf.train.AdamOptimizer( self.learn_rate).minimize( self.loss, var_list=self.first_stage_trainable_var_list) with tf.control_dependencies( tf.get_collection(tf.GraphKeys.UPDATE_OPS)): with tf.control_dependencies( [first_stage_optimizer, global_step_update]): with tf.control_dependencies([moving_ave]): self.train_op_with_frozen_variables = tf.no_op() # 第二阶段训练:训练所有的层,其实也就是 fine-tunning 阶段 with tf.name_scope("define_second_stage_train"): second_stage_trainable_var_list = tf.trainable_variables() second_stage_optimizer = tf.train.AdamOptimizer( self.learn_rate).minimize( self.loss, var_list=second_stage_trainable_var_list) with tf.control_dependencies( tf.get_collection(tf.GraphKeys.UPDATE_OPS)): with tf.control_dependencies( [second_stage_optimizer, global_step_update]): with tf.control_dependencies([moving_ave]): self.train_op_with_all_variables = tf.no_op() with tf.name_scope('loader_and_saver'): self.loader = tf.train.Saver(self.net_var) self.saver = tf.train.Saver(tf.global_variables(), max_to_keep=10) with tf.name_scope('summary'): tf.summary.scalar("learn_rate", self.learn_rate) tf.summary.scalar("giou_loss", self.giou_loss) tf.summary.scalar("conf_loss", self.conf_loss) tf.summary.scalar("prob_loss", self.prob_loss) tf.summary.scalar("total_loss", self.loss) #tf.summary.image("input_image", self.input_data) logdir = "./data/log/" if os.path.exists(logdir): shutil.rmtree(logdir) #递归删除文件夹下的所有子文件夹和子文件 os.mkdir(logdir) self.write_op = tf.summary.merge_all( ) #将所有summary全部保存到磁盘,以便tensorboard显示 self.summary_writer = tf.summary.FileWriter( logdir, graph=self.sess.graph) #指定一个文件用来保存图
def export(): serialized_tf_example = tf.placeholder(tf.string, shape=[None], name='encoded_image_tensor') images = tf.map_fn(preprocess_image, serialized_tf_example, tf.float32) model = YOLOV3(images, trainable=False) print(model.conv_sbbox, model.conv_mbbox, model.conv_lbbox) sess = tf.Session(config=tf.ConfigProto(allow_soft_placement=True)) saver = tf.train.Saver() ckpt = tf.train.get_checkpoint_state(args.ckpt_dir) saver.restore(sess, ckpt.model_checkpoint_path) converted_graph_def = tf.graph_util.convert_variables_to_constants( sess, input_graph_def=sess.graph.as_graph_def(), output_node_names=output_node_names) with tf.gfile.GFile(args.freeze_graph_dir, "wb") as f: f.write(converted_graph_def.SerializeToString()) optimize_graph('', args.freeze_graph_dir, transforms, output_node_names, outname=args.optimizer_graph_dir) graph_def = get_graph_def_from_file(args.optimizer_graph_dir) with tf.Graph().as_default(): tf.import_graph_def(graph_def, input_map=None, return_elements=None, name="") with tf.Session() as sess: # Export inference model. output_path = os.path.join( tf.compat.as_bytes(args.output_dir), tf.compat.as_bytes(str(args.model_version))) print('Exporting trained model to', output_path) builder = tf.saved_model.builder.SavedModelBuilder(output_path) input_tensor = tf.get_default_graph().get_tensor_by_name( args.input_tensor) sbbox_tensor = tf.get_default_graph().get_tensor_by_name( args.sbbox_tensor) mbbox_tensor = tf.get_default_graph().get_tensor_by_name( args.mbbox_tensor) lbbox_tensor = tf.get_default_graph().get_tensor_by_name( args.lbbox_tensor) class_num = args.class_num sbbox_tensor = tf.reshape(tf.convert_to_tensor(sbbox_tensor), shape=[-1, 5 + class_num]) mbbox_tensor = tf.reshape(tf.convert_to_tensor(mbbox_tensor), shape=[-1, 5 + class_num]) lbbox_tensor = tf.reshape(tf.convert_to_tensor(lbbox_tensor), shape=[-1, 5 + class_num]) output_tensor = tf.concat( [sbbox_tensor, mbbox_tensor, lbbox_tensor], 0) # top_100 = tf.nn.top_k(output_tensor[:, 4], 100) # output_tensor = tf.gather(output_tensor, top_100.indices) num_tensor = tf.constant(100) classes_tensor = tf.argmax(output_tensor[:, 5:], axis=1) scores_tensor = output_tensor[:, 4] raw_boxs_tensor = output_tensor[:, 0:4] / IMAGE_SIZE print(raw_boxs_tensor.shape) boxs_tensor_com = raw_boxs_tensor print(boxs_tensor_com.shape) boxs_tensor_minx = raw_boxs_tensor[:, 0] - raw_boxs_tensor[:, 2] * 0.5 boxs_tensor_miny = raw_boxs_tensor[:, 1] - raw_boxs_tensor[:, 3] * 0.5 boxs_tensor_maxx = raw_boxs_tensor[:, 0] + raw_boxs_tensor[:, 2] * 0.5 boxs_tensor_maxy = raw_boxs_tensor[:, 1] + raw_boxs_tensor[:, 3] * 0.5 boxs_tensor_minx = tf.expand_dims(boxs_tensor_minx, 1) boxs_tensor_miny = tf.expand_dims(boxs_tensor_miny, 1) boxs_tensor_maxx = tf.expand_dims(boxs_tensor_maxx, 1) boxs_tensor_maxy = tf.expand_dims(boxs_tensor_maxy, 1) boxs_tensor_com = tf.concat([ boxs_tensor_miny, boxs_tensor_minx, boxs_tensor_maxy, boxs_tensor_maxx ], 1) boxs_tensor_indices = tf.image.non_max_suppression( boxs_tensor_com, scores_tensor, 100, 0.5) boxs_tensor_com = tf.gather(boxs_tensor_com, boxs_tensor_indices) classes_tensor = tf.gather(classes_tensor, boxs_tensor_indices) scores_tensor = tf.gather(scores_tensor, boxs_tensor_indices) scores_tensor_info = tf.saved_model.utils.build_tensor_info( tf.expand_dims(scores_tensor, 0)) classes_tensor_info = tf.saved_model.utils.build_tensor_info( tf.expand_dims(classes_tensor, 0)) boxes_tensor_info = tf.saved_model.utils.build_tensor_info( tf.expand_dims(boxs_tensor_com, 0)) raw_boxes_tensor_info = tf.saved_model.utils.build_tensor_info( tf.expand_dims(raw_boxs_tensor, 0)) num_tensor_info = tf.saved_model.utils.build_tensor_info( tf.expand_dims(num_tensor, 0)) inputs_tensor_info = tf.saved_model.utils.build_tensor_info( input_tensor) tensor_info_inputs = {'inputs': inputs_tensor_info} print(scores_tensor_info, inputs_tensor_info, classes_tensor_info, boxes_tensor_info, num_tensor_info) prediction_signature = ( tf.saved_model.signature_def_utils.build_signature_def( inputs=tensor_info_inputs, # outputs=tensor_info_outputs, outputs={ 'detection_scores': scores_tensor_info, 'detection_classes': classes_tensor_info, 'detection_boxes': boxes_tensor_info, # 'raw_boxes': raw_boxes_tensor_info, 'num_detections': num_tensor_info, }, method_name=tf.saved_model.signature_constants. PREDICT_METHOD_NAME)) builder.add_meta_graph_and_variables( sess, [tf.saved_model.tag_constants.SERVING], signature_def_map={ tf.saved_model.signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY: prediction_signature, }, ) builder.save() print('Successfully exported model to %s' % args.output_dir)
import tensorflow as tf from core.backbone import darknet53 from core.yolov3 import YOLOV3 def stats_graph(graph): flops = tf.profiler.profile( graph, options=tf.profiler.ProfileOptionBuilder.float_operation()) params = tf.profiler.profile(graph, options=tf.profiler.ProfileOptionBuilder. trainable_variables_parameter()) print('FLOPs: {}; Trainable params: {}'.format(flops.total_float_ops, params.total_parameters)) with tf.Graph().as_default() as graph: image = tf.placeholder(dtype=tf.float32, shape=[1, 544, 544, 3], name="input") out = YOLOV3(image, True, 1) stats_graph(graph)
def __init__(self): self.anchor_per_scale = 3 self.classes = utils.read_class_names( "/home/Pedestrian/Documents/TensorFlow_YOLOv3-master/LabelImage_v1.8.1/data/predefined_classes.txt" ) self.num_classes = len(self.classes) self.train_epochs = 160 self.max_bbox_per_scale = 150 self.trainset = Dataset() self.steps_per_period = len(self.trainset) self.sess = tf.Session(config=tf.ConfigProto( log_device_placement=True)) with tf.name_scope('define_input'): self.input_data = tf.placeholder(dtype=tf.float32, name='input_data') self.label_sbbox = tf.placeholder(dtype=tf.float32, name='label_sbbox') self.label_mbbox = tf.placeholder(dtype=tf.float32, name='label_mbbox') self.label_lbbox = tf.placeholder(dtype=tf.float32, name='label_lbbox') self.true_sbboxes = tf.placeholder(dtype=tf.float32, name='sbboxes') self.true_mbboxes = tf.placeholder(dtype=tf.float32, name='mbboxes') self.true_lbboxes = tf.placeholder(dtype=tf.float32, name='lbboxes') self.trainable = tf.placeholder(dtype=tf.bool, name='training') with tf.name_scope("define_loss"): self.model = YOLOV3(self.input_data, self.trainable) self.net_var = tf.global_variables() self.giou_loss, self.conf_loss, self.prob_loss = self.model.compute_loss( self.label_sbbox, self.label_mbbox, self.label_lbbox, self.true_sbboxes, self.true_mbboxes, self.true_lbboxes) self.loss = self.giou_loss + self.conf_loss + self.prob_loss with tf.name_scope('learn_rate'): self.global_step = tf.Variable( 1.0, dtype=tf.float32, trainable=False, name='global_step' ) # step 1: 1.0 ==> num * self.steps_per_period self.update_global_step = tf.assign_add(self.global_step, 1.0) self.learn_rate = tf.Variable(1e-4, dtype=tf.float32, trainable=False, name='learn_rate') self.update_learn_rate = tf.assign(self.learn_rate, self.learn_rate * 0.8) with tf.name_scope("train"): self.optimizer = tf.train.AdamOptimizer(self.learn_rate).minimize( self.loss, var_list=tf.trainable_variables()) with tf.control_dependencies( tf.get_collection(tf.GraphKeys.UPDATE_OPS)): with tf.control_dependencies( [self.optimizer, self.update_global_step]): self.train_op_with_trainable_variables = tf.no_op() with tf.name_scope('summary'): tf.summary.scalar("learn_rate", self.learn_rate) tf.summary.scalar("giou_loss", self.giou_loss) tf.summary.scalar("conf_loss", self.conf_loss) tf.summary.scalar("prob_loss", self.prob_loss) tf.summary.scalar("total_loss", self.loss) logdir = "log" if os.path.exists(logdir): shutil.rmtree(logdir) os.mkdir(logdir) self.write_op = tf.summary.merge_all() self.summary_writer = tf.summary.FileWriter(logdir, graph=self.sess.graph)
def __init__(self): # 每一层上Anchor的数目 self.anchor_per_scale = cfg.YOLO.ANCHOR_PER_SCALE # 类别 self.classes = utils.read_class_names(cfg.YOLO.CLASSES) # 类别数目 self.num_classes = len(self.classes) # 初始化学习率 self.learn_rate_init = cfg.TRAIN.LEARN_RATE_INIT # 最小学习率 self.learn_rate_end = cfg.TRAIN.LEARN_RATE_END # 第一阶段的阶段数 self.first_stage_epochs = cfg.TRAIN.FISRT_STAGE_EPOCHS # 第二阶段的阶段数 self.second_stage_epochs = cfg.TRAIN.SECOND_STAGE_EPOCHS # 学习率变化参数 self.warmup_periods = cfg.TRAIN.WARMUP_EPOCHS # 模型训练持久化恢复对应的文件路径 self.initial_weight = cfg.TRAIN.INITIAL_WEIGHT # 当前时间信息 self.time = time.strftime('%Y-%m-%d-%H-%M-%S', time.localtime(time.time())) # 滑动平均的系数值(用于训练数据对应的滑动平均系数) self.moving_ave_decay = cfg.YOLO.MOVING_AVE_DECAY # 每个尺度最多允许的边框数目 self.max_bbox_per_scale = 150 # 训练数据所在的文件路径 self.train_logdir = "./data/log/train" # 加载训练数据 # TODO: 修改为支持自己数据的形式 self.trainset = Dataset('train') # 加载测试数据 self.testset = Dataset('test') # 总训练数据 self.steps_per_period = len(self.trainset) # 构建会话对象 self.sess = tf.Session(config=tf.ConfigProto( allow_soft_placement=True)) # 定义网络输入 # 注输出是置信度和边框,也就是一个输入对应两个输出,在这里就是6个输出 with tf.name_scope('define_input'): self.input_data = tf.placeholder(dtype=tf.float32, name='input_data') self.label_sbbox = tf.placeholder(dtype=tf.float32, name='label_sbbox') # 候选框 self.label_mbbox = tf.placeholder(dtype=tf.float32, name='label_mbbox') self.label_lbbox = tf.placeholder(dtype=tf.float32, name='label_lbbox') self.true_sbboxes = tf.placeholder(dtype=tf.float32, name='sbboxes') # 位置信息 self.true_mbboxes = tf.placeholder(dtype=tf.float32, name='mbboxes') self.true_lbboxes = tf.placeholder(dtype=tf.float32, name='lbboxes') self.trainable = tf.placeholder(dtype=tf.bool, name='training') #是训练还是测试 # 定义网络及损失函数 with tf.name_scope("define_loss"): # 构建模型 self.model = YOLOV3(self.input_data, self.trainable) # 获取网络中所有的模型变量 self.net_var = tf.global_variables() # 构建计算位置、可信度以及概率损失函数 self.giou_loss, self.conf_loss, self.prob_loss = self.model.compute_loss( self.label_sbbox, self.label_mbbox, self.label_lbbox, self.true_sbboxes, self.true_mbboxes, self.true_lbboxes) # 三个损失函数合并成一个损失函数 self.loss = self.giou_loss + self.conf_loss + self.prob_loss with tf.name_scope('learn_rate'): # 构建一个全局变量对象 self.global_step = tf.Variable(1.0, dtype=tf.float64, trainable=False, name='global_step') # 学习率变化阈值(默认两个批次之前) warmup_steps = tf.constant(self.warmup_periods * self.steps_per_period, dtype=tf.float64, name='warmup_steps') # 总训练步骤 train_steps = tf.constant( (self.first_stage_epochs + self.second_stage_epochs) * self.steps_per_period, dtype=tf.float64, name='train_steps') # 在两个批次之前,如果训练迭代次数小于阈值,那么设置为初始概率的一部分。 # 当批次超过阈值之后,做一个学习率转换 self.learn_rate = tf.cond( pred=self.global_step < warmup_steps, #刚开始学习率大,train_steps和warmup_steps为固定值 true_fn=lambda: self.global_step / warmup_steps * self. learn_rate_init, false_fn=lambda: self.learn_rate_end + 0.5 * (self.learn_rate_init - self.learn_rate_end) * (1 + tf.cos( (self.global_step - warmup_steps) / (train_steps - warmup_steps) * np.pi))) # 定义全局步骤变量更新参数 global_step_update = tf.assign_add(self.global_step, 1.0) # 定义模型参数滑动平均更新(目的:为了让模型参数更加平滑) with tf.name_scope("define_weight_decay"): moving_ave = tf.train.ExponentialMovingAverage(self.moving_ave_decay) \ .apply(tf.trainable_variables()) # 第一阶段的模型训练相关参数设置 with tf.name_scope("define_first_stage_train"): # 获取第一阶段模型训练相关参数 self.first_stage_trainable_var_list = [] for var in tf.trainable_variables(): var_name = var.op.name var_name_mess = str(var_name).split('/') if var_name_mess[0] in [ 'conv_sbbox', 'conv_mbbox', 'conv_lbbox' ]: self.first_stage_trainable_var_list.append(var) # 第一阶段的模型优化器 first_stage_optimizer = tf.train.AdamOptimizer(self.learn_rate) \ .minimize(self.loss, var_list=self.first_stage_trainable_var_list) # 加入依赖控制 # 先做BN的更新操作、在进行模型训练/步骤参数更新、最终更新模型参数值 with tf.control_dependencies( tf.get_collection(tf.GraphKeys.UPDATE_OPS)): # 批归一化 with tf.control_dependencies( [first_stage_optimizer, global_step_update]): with tf.control_dependencies([moving_ave]): # 定义一个训练操作(实际上不做任何操作,仅仅触发所有操作) self.train_op_with_frozen_variables = tf.no_op() # 第二阶段的模型训练相关参数设置 with tf.name_scope("define_second_stage_train"): # 获取第二阶段的相关训练参数(所有训练参数) second_stage_trainable_var_list = tf.trainable_variables() # 构建第二阶段的优化器对象 second_stage_optimizer = tf.train.AdamOptimizer(self.learn_rate) \ .minimize(self.loss, var_list=second_stage_trainable_var_list) # 加入依赖控制 # 先做BN的更新操作、在进行模型训练/步骤参数更新、最终更新模型参数值 with tf.control_dependencies( tf.get_collection(tf.GraphKeys.UPDATE_OPS)): with tf.control_dependencies( [second_stage_optimizer, global_step_update]): with tf.control_dependencies([moving_ave]): # 定义一个训练操作(实际上不做任何操作,仅仅触发所有操作) self.train_op_with_all_variables = tf.no_op() with tf.name_scope('loader_and_saver'): # 定义模型加载对象以及模型持久化对象 self.loader = tf.train.Saver( self.net_var) # 仅加载网络参数(也就是前向过程中的所有参数) self.saver = tf.train.Saver(tf.global_variables(), max_to_keep=3) with tf.name_scope('summary'): # 可视化相关内容 tf.summary.scalar("learn_rate", self.learn_rate) tf.summary.scalar("giou_loss", self.giou_loss) tf.summary.scalar("conf_loss", self.conf_loss) tf.summary.scalar("prob_loss", self.prob_loss) tf.summary.scalar("total_loss", self.loss) # 可视化日志文件构建 logdir = "./data/log2/" if os.path.exists(logdir): shutil.rmtree(logdir) os.mkdir(logdir) # 输出可视化对象 self.write_op = tf.summary.merge_all() self.summary_writer = tf.summary.FileWriter(logdir, graph=self.sess.graph) print("初始化完成.....")
def __init__(self): # Initialize Horovod hvd.init() config=tf.ConfigProto(allow_soft_placement=True) config.gpu_options.visible_device_list = str(hvd.local_rank()) self.anchor_per_scale = cfg.YOLO.ANCHOR_PER_SCALE self.classes = utils.read_class_names(cfg.YOLO.CLASSES) self.num_classes = len(self.classes) self.learn_rate_init = cfg.TRAIN.LEARN_RATE_INIT self.learn_rate_end = cfg.TRAIN.LEARN_RATE_END self.first_stage_epochs = cfg.TRAIN.FISRT_STAGE_EPOCHS self.second_stage_epochs = cfg.TRAIN.SECOND_STAGE_EPOCHS self.warmup_periods = cfg.TRAIN.WARMUP_EPOCHS self.initial_weight = cfg.TRAIN.INITIAL_WEIGHT self.time = time.strftime('%Y-%m-%d-%H-%M-%S', time.localtime(time.time())) self.moving_ave_decay = cfg.YOLO.MOVING_AVE_DECAY self.max_bbox_per_scale = 150 self.train_logdir = "./data/log/train" self.trainset = Dataset('train') self.testset = Dataset('test') self.steps_per_period = len(self.trainset) self.sess = tf.Session(config=config) with tf.name_scope('define_input'): self.input_data = tf.placeholder(dtype=tf.float32, name='input_data') self.label_sbbox = tf.placeholder(dtype=tf.float32, name='label_sbbox') self.label_mbbox = tf.placeholder(dtype=tf.float32, name='label_mbbox') self.label_lbbox = tf.placeholder(dtype=tf.float32, name='label_lbbox') self.true_sbboxes = tf.placeholder(dtype=tf.float32, name='sbboxes') self.true_mbboxes = tf.placeholder(dtype=tf.float32, name='mbboxes') self.true_lbboxes = tf.placeholder(dtype=tf.float32, name='lbboxes') self.trainable = tf.placeholder(dtype=tf.bool, name='training') with tf.name_scope("define_loss"): self.model = YOLOV3(self.input_data, self.trainable) self.net_var = tf.global_variables() self.giou_loss, self.conf_loss, self.prob_loss = self.model.compute_loss( self.label_sbbox, self.label_mbbox, self.label_lbbox, self.true_sbboxes, self.true_mbboxes, self.true_lbboxes) self.loss = self.giou_loss + self.conf_loss + self.prob_loss with tf.name_scope('learn_rate'): self.global_step = tf.Variable(1.0, dtype=tf.float64, trainable=False, name='global_step') warmup_steps = tf.constant(self.warmup_periods * self.steps_per_period, dtype=tf.float64, name='warmup_steps') train_steps = tf.constant( (self.first_stage_epochs + self.second_stage_epochs)* self.steps_per_period, dtype=tf.float64, name='train_steps') self.learn_rate = tf.cond( pred=self.global_step < warmup_steps, true_fn=lambda: self.global_step / warmup_steps * self.learn_rate_init, false_fn=lambda: self.learn_rate_end + 0.5 * (self.learn_rate_init - self.learn_rate_end) * (1 + tf.cos( (self.global_step - warmup_steps) / (train_steps - warmup_steps) * np.pi)) ) global_step_update = tf.assign_add(self.global_step, 1.0) #for Horovod expand learning rate self.learn_rate = self.learn_rate * hvd.size() with tf.name_scope("define_weight_decay"): moving_ave = tf.train.ExponentialMovingAverage(self.moving_ave_decay).apply(tf.trainable_variables()) with tf.name_scope("define_first_stage_train"): self.first_stage_trainable_var_list = [] for var in tf.trainable_variables(): var_name = var.op.name var_name_mess = str(var_name).split('/') if var_name_mess[0] in ['conv_sbbox', 'conv_mbbox', 'conv_lbbox']: self.first_stage_trainable_var_list.append(var) first_opt = tf.train.AdamOptimizer(self.learn_rate) #for Horovod first_opt = hvd.DistributedOptimizer(first_opt) #for Horovod hooks = [hvd.BroadcastGlobalVariablesHook(0)] first_stage_optimizer = first_opt.minimize(self.loss, var_list=self.first_stage_trainable_var_list) with tf.control_dependencies(tf.get_collection(tf.GraphKeys.UPDATE_OPS)): with tf.control_dependencies([first_stage_optimizer, global_step_update]): with tf.control_dependencies([moving_ave]): self.train_op_with_frozen_variables = tf.no_op() with tf.name_scope("define_second_stage_train"): second_stage_trainable_var_list = tf.trainable_variables() second_opt = tf.train.AdamOptimizer(self.learn_rate) #for Horovod second_opt = hvd.DistributedOptimizer(second_opt) second_stage_optimizer = second_opt.minimize(self.loss, var_list=second_stage_trainable_var_list) with tf.control_dependencies(tf.get_collection(tf.GraphKeys.UPDATE_OPS)): with tf.control_dependencies([second_stage_optimizer, global_step_update]): with tf.control_dependencies([moving_ave]): self.train_op_with_all_variables = tf.no_op() with tf.name_scope('loader_and_saver'): self.loader = tf.train.Saver(self.net_var) self.saver = tf.train.Saver(tf.global_variables(), max_to_keep=10) with tf.name_scope('summary'): tf.summary.scalar("learn_rate", self.learn_rate) tf.summary.scalar("giou_loss", self.giou_loss) tf.summary.scalar("conf_loss", self.conf_loss) tf.summary.scalar("prob_loss", self.prob_loss) tf.summary.scalar("total_loss", self.loss) logdir = "./data/log/" if os.path.exists(logdir): shutil.rmtree(logdir) os.mkdir(logdir) self.write_op = tf.summary.merge_all() self.summary_writer = tf.summary.FileWriter(logdir, graph=self.sess.graph)
i += 1 print("ptr:", ptr) return assign_ops with open(arguments.config, "r") as f: opts = json.load(f) with tf.name_scope("input"): input_data = tf.placeholder(dtype=tf.float32, shape=(1, 320, 320, 3), name="input_data") # converge to a ckpt with fp32 # then use another script to convert to fp16 if needed opts["yolo"]["precision"] = "fp32" model = YOLOV3(False, opts) # construct model layers = [ model.build_backbone_part1, model.build_backbone_part2, model.build_backbone_part3, model.build_backbone_part4, model.build_backbone_part5, model.build_backbone_part6, ] # reuse stages_constructor so we don't need to pass params by hand network_func = stages_constructor( [layers], ["image"], ["pred_sbbox", "pred_mbbox", "pred_lbbox"])[0] network_func(input_data)
if not os.path.exists(ckpt_file + '.index'): print('freeze_ckpt_to_pb ckpt_file=', ckpt_file, ' not exist') sys.exit() pb_file = ckpt_file + '.pb' #argv[4] print('freeze_ckpt_to_pb gpu_id=%s, net_type=%s, ckpt_file=%s, pb_file=%s' % (gpu_id, net_type, ckpt_file, pb_file)) os.environ['CUDA_VISIBLE_DEVICES'] = str(gpu_id) output_node_names = ["input/input_data", "pred_sbbox/concat_2", "pred_mbbox/concat_2", "pred_lbbox/concat_2"] #output_node_names = ["input/input_data", "conv_sbbox/BiasAdd", "conv_mbbox/BiasAdd", "conv_lbbox/BiasAdd"] with tf.name_scope('input'): input_data = tf.placeholder(dtype=tf.float32, name='input_data') if net_type == 'yolov3': model = YOLOV3(input_data, trainable=False, freeze_pb=False) elif net_type == 'yolov4': model = YOLOV4(input_data, trainable=False) elif net_type == 'yolov5': model = YOLOV5(input_data, trainable=False) else: print('net_type=', net_type, ' error') sess = tf.Session(config=tf.ConfigProto(allow_soft_placement=True)) saver = tf.train.Saver() saver.restore(sess, ckpt_file) converted_graph_def = tf.graph_util.convert_variables_to_constants(sess, input_graph_def=sess.graph.as_graph_def(), output_node_names=output_node_names) with tf.gfile.GFile(pb_file, "wb") as f: f.write(converted_graph_def.SerializeToString())
with tf.name_scope('input'): serialized_input_data = tf.placeholder(dtype=tf.string, name='input_data') feature_configs = { 'image/encoded': tf.FixedLenFeature(shape=[], dtype=tf.string), } input_data = tf.parse_example(serialized_input_data, feature_configs) jpegs = input_data['image/encoded'] image_string = tf.reshape(jpegs, shape=[]) image_tensor = tf.image.decode_image(image_string, channels=3) image_tensor.set_shape((None, None, 3)) image_input = image_tensor[np.newaxis, ...] / 255 model = YOLOV3(image_input, trainable=False) print(model.conv_sbbox, model.conv_mbbox, model.conv_lbbox) sess = tf.Session(config=tf.ConfigProto(allow_soft_placement=True)) saver = tf.train.Saver() saver.restore(sess, ckpt_file) converted_graph_def = tf.graph_util.convert_variables_to_constants( sess, input_graph_def=sess.graph.as_graph_def(), output_node_names=output_node_names) # Export inference model. output_path = os.path.join(tf.compat.as_bytes(output_dir), tf.compat.as_bytes(str(model_version))) print('Exporting trained model to', output_path)
def __init__(self, net_type): self.net_type = net_type self.anchor_per_scale = cfg.YOLO.ANCHOR_PER_SCALE self.classes = utils.read_class_names(cfg.YOLO.CLASSES) self.num_classes = len(self.classes) self.learn_rate_init = cfg.TRAIN.LEARN_RATE_INIT self.learn_rate_end = cfg.TRAIN.LEARN_RATE_END self.first_stage_epochs = cfg.TRAIN.FISRT_STAGE_EPOCHS self.second_stage_epochs = cfg.TRAIN.SECOND_STAGE_EPOCHS self.warmup_periods = cfg.TRAIN.WARMUP_EPOCHS self.initial_weight = cfg.TRAIN.INITIAL_WEIGHT self.ckpt_path = cfg.TRAIN.CKPT_PATH if not os.path.exists(self.ckpt_path): os.makedirs(self.ckpt_path) self.time = time.strftime('%Y-%m-%d-%H-%M-%S', time.localtime(time.time())) self.moving_ave_decay = cfg.YOLO.MOVING_AVE_DECAY self.max_bbox_per_scale = 150 self.log_path = ('log/%s' % net_type) if os.path.exists(self.log_path): shutil.rmtree(self.log_path) #os.removedirs(self.log_path) os.makedirs(self.log_path) self.trainset = Dataset('train', self.net_type) self.testset = Dataset('test', self.net_type) self.steps_per_period = len(self.trainset) config = tf.ConfigProto() config.gpu_options.allow_growth = True self.sess = tf.Session(config=config) with tf.name_scope('input'): if net_type == 'tiny': self.input_data = tf.placeholder(dtype=tf.float32, name='input_data') self.label_mbbox = tf.placeholder(dtype=tf.float32, name='label_mbbox') self.label_lbbox = tf.placeholder(dtype=tf.float32, name='label_lbbox') self.true_mbboxes = tf.placeholder(dtype=tf.float32, name='mbboxes') self.true_lbboxes = tf.placeholder(dtype=tf.float32, name='lbboxes') self.trainable = tf.placeholder(dtype=tf.bool, name='training') else: self.input_data = tf.placeholder(dtype=tf.float32, name='input_data') self.label_sbbox = tf.placeholder(dtype=tf.float32, name='label_sbbox') self.label_mbbox = tf.placeholder(dtype=tf.float32, name='label_mbbox') self.label_lbbox = tf.placeholder(dtype=tf.float32, name='label_lbbox') self.true_sbboxes = tf.placeholder(dtype=tf.float32, name='sbboxes') self.true_mbboxes = tf.placeholder(dtype=tf.float32, name='mbboxes') self.true_lbboxes = tf.placeholder(dtype=tf.float32, name='lbboxes') self.trainable = tf.placeholder(dtype=tf.bool, name='training') with tf.name_scope('define_loss'): if self.net_type == 'tiny': self.model = YOLOV3Tiny(self.input_data, self.trainable) self.net_var = tf.global_variables() self.iou_loss, self.conf_loss, self.prob_loss = self.model.compute_loss(self.label_mbbox, self.label_lbbox, self.true_mbboxes, self.true_lbboxes) self.loss = self.iou_loss + self.conf_loss + self.prob_loss elif self.net_type == 'yolov3': self.model = YOLOV3(self.input_data, self.trainable) self.net_var = tf.global_variables() self.iou_loss, self.conf_loss, self.prob_loss = self.model.compute_loss(self.label_sbbox, self.label_mbbox, self.label_lbbox, self.true_sbboxes, self.true_mbboxes, self.true_lbboxes) self.loss = self.iou_loss + self.conf_loss + self.prob_loss elif self.net_type == 'yolov4' or self.net_type == 'yolov5': iou_use = 1 # (0, 1, 2) ==> (giou_loss, diou_loss, ciou_loss) focal_use = False # (False, True) ==> (normal, focal_loss) label_smoothing = 0 if self.net_type == 'yolov4': self.model = YOLOV4(self.input_data, self.trainable) else: self.model = YOLOV5(self.input_data, self.trainable) self.net_var = tf.global_variables() self.iou_loss, self.conf_loss, self.prob_loss = self.model.compute_loss(self.label_sbbox, self.label_mbbox, self.label_lbbox, self.true_sbboxes, self.true_mbboxes, self.true_lbboxes, iou_use, focal_use, label_smoothing) self.loss = self.iou_loss + self.conf_loss + self.prob_loss # self.loss = tf.Print(self.loss, [self.iou_loss, self.conf_loss, self.prob_loss], message='loss: ') else: print('self.net_type=%s error' % self.net_type) with tf.name_scope('learn_rate'): self.global_step = tf.Variable(1.0, dtype=tf.float64, trainable=False, name='global_step') warmup_steps = tf.constant(self.warmup_periods * self.steps_per_period, dtype=tf.float64, name='warmup_steps') train_steps = tf.constant((self.first_stage_epochs + self.second_stage_epochs) * self.steps_per_period, dtype=tf.float64, name='train_steps') self.learn_rate = tf.cond(pred=self.global_step < warmup_steps, true_fn=lambda: self.global_step / warmup_steps * self.learn_rate_init, false_fn=lambda: self.learn_rate_end + 0.5 * (self.learn_rate_init - self.learn_rate_end) * \ (1 + tf.cos((self.global_step - warmup_steps) / (train_steps - warmup_steps) * np.pi))) global_step_update = tf.assign_add(self.global_step, 1.0) with tf.name_scope('define_weight_decay'): moving_ave = tf.train.ExponentialMovingAverage(self.moving_ave_decay).apply(tf.trainable_variables()) with tf.name_scope('define_first_stage_train'): self.first_stage_trainable_var_list = [] for var in tf.trainable_variables(): var_name = var.op.name var_name_mess = str(var_name).split('/') if net_type == 'tiny': bboxes = ['conv_mbbox', 'conv_lbbox'] else: bboxes = ['conv_sbbox', 'conv_mbbox', 'conv_lbbox'] if var_name_mess[0] in bboxes: self.first_stage_trainable_var_list.append(var) first_stage_optimizer = tf.train.AdamOptimizer(self.learn_rate).minimize(self.loss, var_list=self.first_stage_trainable_var_list) with tf.control_dependencies(tf.get_collection(tf.GraphKeys.UPDATE_OPS)): with tf.control_dependencies([first_stage_optimizer, global_step_update]): with tf.control_dependencies([moving_ave]): self.train_op_with_frozen_variables = tf.no_op() with tf.name_scope('define_second_stage_train'): second_stage_trainable_var_list = tf.trainable_variables() second_stage_optimizer = tf.train.AdamOptimizer(self.learn_rate).minimize(self.loss, var_list=second_stage_trainable_var_list) with tf.control_dependencies(tf.get_collection(tf.GraphKeys.UPDATE_OPS)): with tf.control_dependencies([second_stage_optimizer, global_step_update]): with tf.control_dependencies([moving_ave]): self.train_op_with_all_variables = tf.no_op() with tf.name_scope('loader_and_saver'): self.loader = tf.train.Saver(self.net_var) self.saver = tf.train.Saver(tf.global_variables(), max_to_keep=1000) with tf.name_scope('summary'): tf.summary.scalar('learn_rate', self.learn_rate) tf.summary.scalar('iou_loss', self.iou_loss) tf.summary.scalar('conf_loss', self.conf_loss) tf.summary.scalar('prob_loss', self.prob_loss) tf.summary.scalar('total_loss', self.loss) self.write_op = tf.summary.merge_all() self.summary_writer = tf.summary.FileWriter(self.log_path, graph=self.sess.graph)
var_name = var.op.name var_name_mess = str(var_name).split('/') var_shape = var.shape org_weights_mess.append([var_name, var_shape]) print("=> " + str(var_name).ljust(50), var_shape) print() tf.reset_default_graph() cur_weights_mess = [] tf.Graph().as_default() with tf.name_scope('input'): input_data = tf.placeholder(dtype=tf.float32, shape=(1, 416, 416, 3), name='input_data') training = tf.placeholder(dtype=tf.bool, name='trainable') model = YOLOV3(input_data, training) for var in tf.global_variables(): var_name = var.op.name var_name_mess = str(var_name).split('/') var_shape = var.shape print(var_name_mess[0]) cur_weights_mess.append([var_name, var_shape]) print("=> " + str(var_name).ljust(50), var_shape) org_weights_num = len(org_weights_mess) cur_weights_num = len(cur_weights_mess) if cur_weights_num != org_weights_num: raise RuntimeError print('=> Number of weights that will rename:\t%d' % cur_weights_num) cur_to_org_dict = {}
# #================================================================ import tensorflow as tf from core.yolov3 import YOLOV3 pb_file = "./yolov3_coco.pb" ckpt_file = "./checkpoint/yolov3_test_loss=10.6896.ckpt-50" output_node_names = [ "input/input_data", "pred_sbbox/concat_2", "pred_mbbox/concat_2", "pred_lbbox/concat_2" ] with tf.name_scope('input'): input_data = tf.placeholder(dtype=tf.float32, name='input_data') model = YOLOV3(input_data, trainable=tf.cast(False, tf.bool)) print(model.conv_sbbox, model.conv_mbbox, model.conv_lbbox) sess = tf.Session(config=tf.ConfigProto(allow_soft_placement=True)) saver = tf.train.Saver() saver.restore(sess, ckpt_file) converted_graph_def = tf.graph_util.convert_variables_to_constants( sess, input_graph_def=sess.graph.as_graph_def(), output_node_names=output_node_names) with tf.gfile.GFile(pb_file, "wb") as f: f.write(converted_graph_def.SerializeToString())
# #================================================================ import tensorflow as tf from core.yolov3 import YOLOV3 from core.config import cfg pb_file = "./yolov3_mark_no_manualFlip_no_codeFlip.pb" ckpt_file = cfg.TEST.WEIGHT_FILE output_node_names = ["input/input_data", "pred_sbbox/concat_2", "pred_mbbox/concat_2", "pred_lbbox/concat_2"] with tf.name_scope('input'): input_data = tf.placeholder(dtype=tf.float32, name='input_data') model = YOLOV3(input_data, trainable=False) print(model.conv_sbbox, model.conv_mbbox, model.conv_lbbox) sess = tf.Session(config=tf.ConfigProto(allow_soft_placement=True)) saver = tf.train.Saver() saver.restore(sess, ckpt_file) converted_graph_def = tf.graph_util.convert_variables_to_constants(sess, input_graph_def = sess.graph.as_graph_def(), output_node_names = output_node_names) with tf.gfile.GFile(pb_file, "wb") as f: f.write(converted_graph_def.SerializeToString())
def __init__(self): self.anchor_per_scale = 3 self.path1 = r"./data/classes/antenna.names" #修改 self.classes = utils.read_class_names(self.path1) self.num_classes = len(self.classes) self.learn_rate_init = 1e-4 # 1e-4 self.learn_rate_end = 1e-6 # 1e-6 self.first_stage_epochs = 20 # 40 加载不到预训练权重,不进行一阶段训练 self.warmup_periods = 2 # 默认为23444 self.time = time.strftime('%Y-%m-%d-%H-%M-%S', time.localtime(time.time())) # 初始化时间 self.moving_ave_decay = 0.9995 # 默认为0.9995 self.max_bbox_per_scale = 150 # 每个照片最多的ground truth框的数量 self.trainset = Dataset('train') # train self.testset = Dataset('test') self.steps_per_period = len(self.trainset) # 整个数据集迭代一个epoch,需要多少个batch self.sess = tf.Session(config=tf.ConfigProto(allow_soft_placement=True)) with tf.name_scope('define_input'): self.input_data = tf.placeholder(dtype=tf.float32, name='input_data') self.label_sbbox = tf.placeholder(dtype=tf.float32, name='label_sbbox') self.label_mbbox = tf.placeholder(dtype=tf.float32, name='label_mbbox') self.label_lbbox = tf.placeholder(dtype=tf.float32, name='label_lbbox') self.true_sbboxes = tf.placeholder(dtype=tf.float32, name='sbboxes') self.true_mbboxes = tf.placeholder(dtype=tf.float32, name='mbboxes') self.true_lbboxes = tf.placeholder(dtype=tf.float32, name='lbboxes') self.trainable = tf.placeholder(dtype=tf.bool, name='training') # TODO 主要的部分 with tf.name_scope("define_loss"): self.model = YOLOV3(self.input_data, self.trainable) self.net_var = tf.global_variables() self.giou_loss, self.conf_loss, self.prob_loss = self.model.compute_loss( self.label_sbbox, self.label_mbbox, self.label_lbbox, self.true_sbboxes, self.true_mbboxes, self.true_lbboxes) self.loss = self.giou_loss + self.conf_loss + self.prob_loss # warm up with tf.name_scope('learn_rate'): self.global_step = tf.Variable(1.0, dtype=tf.float32, trainable=False, name='global_step') # 指数衰减lr self.learn_rate = tf.train.exponential_decay(learning_rate=self.learn_rate_init, global_step=self.global_step, decay_steps=self.steps_per_period * 20, decay_rate=0.9, staircase=False, name=None) global_step_update = tf.assign_add(self.global_step, 1.0) with tf.name_scope("define_weight_decay"): moving_ave = tf.train.ExponentialMovingAverage(self.moving_ave_decay).apply(tf.trainable_variables()) with tf.name_scope("define_second_stage_train"): second_stage_trainable_var_list = tf.trainable_variables() second_stage_optimizer = tf.train.AdamOptimizer(self.learn_rate).minimize(self.loss, var_list=second_stage_trainable_var_list) with tf.control_dependencies(tf.get_collection(tf.GraphKeys.UPDATE_OPS)): with tf.control_dependencies([second_stage_optimizer, global_step_update]): with tf.control_dependencies([moving_ave]): self.train_op_with_all_variables = tf.no_op() # tf.no_op()表示执行完control_dependencies的变量更新之后,不做任何操作,主要确保control_dependencies的变量更新。 with tf.name_scope('loader_and_saver'): self.loader = tf.train.Saver(self.net_var) self.saver = tf.train.Saver(tf.global_variables(), max_to_keep=5)
def __init__(self): self.anchor_per_scale = cfg.YOLO.ANCHOR_PER_SCALE self.classes = utils.read_class_names(cfg.YOLO.CLASSES) self.num_classes = len(self.classes) self.learn_rate_init = cfg.TRAIN.LEARN_RATE_INIT self.learn_rate_end = cfg.TRAIN.LEARN_RATE_END self.first_stage_epochs = cfg.TRAIN.FISRT_STAGE_EPOCHS self.second_stage_epochs = cfg.TRAIN.SECOND_STAGE_EPOCHS self.warmup_periods = cfg.TRAIN.WARMUP_EPOCHS self.initial_weight = cfg.TRAIN.INITIAL_WEIGHT self.time = time.strftime('%Y-%m-%d-%H-%M-%S', time.localtime(time.time())) self.moving_ave_decay = cfg.YOLO.MOVING_AVE_DECAY self.train_logdir = "./data/log/train" self.sess = tf.Session(config=tf.ConfigProto( allow_soft_placement=True)) #self.train_input_sizes = cfg.TRAIN.INPUT_SIZE ##[320, 352, 384, 416, 448, 480, 512, 544, 576, 608] #self.train_data = Dataset('train',batch_size=1,num_parallel=4) self.yolo_train_data = Dataset('train', num_parallel=1) self.yolo_val_data = Dataset('val', num_parallel=1) self.train_data = self.yolo_train_data.get_dataset() self.val_data = self.yolo_val_data.get_dataset() self.iterator = tf.data.Iterator.from_structure( self.train_data.output_types, self.train_data.output_shapes) self.train_init_op = self.iterator.make_initializer(self.train_data) self.val_init_op = self.iterator.make_initializer(self.val_data) self.steps_per_period = self.yolo_train_data.batches_per_epoch with tf.name_scope('define_input'): # self.input_data = tf.placeholder(dtype=tf.float32, name='input_data') # self.label_sbbox = tf.placeholder(dtype=tf.float32, name='label_sbbox') # self.label_mbbox = tf.placeholder(dtype=tf.float32, name='label_mbbox') # self.label_lbbox = tf.placeholder(dtype=tf.float32, name='label_lbbox') # self.true_sbboxes = tf.placeholder(dtype=tf.float32, name='sbboxes') # self.true_mbboxes = tf.placeholder(dtype=tf.float32, name='mbboxes') # self.true_lbboxes = tf.placeholder(dtype=tf.float32, name='lbboxes') self.trainable = tf.placeholder(dtype=tf.bool, name='training') with tf.name_scope("define_loss"): self.input_data, self.label_sbbox, self.label_mbbox, self.label_lbbox,\ self.true_sbboxes, self.true_mbboxes, self.true_lbboxes = self.iterator.get_next() self.model = YOLOV3(self.input_data, self.trainable) self.net_var = tf.global_variables() self.giou_loss, self.conf_loss, self.prob_loss = self.model.compute_loss( self.label_sbbox, self.label_mbbox, self.label_lbbox, self.true_sbboxes, self.true_mbboxes, self.true_lbboxes) self.loss = self.giou_loss + self.conf_loss + self.prob_loss with tf.name_scope('learn_rate'): self.global_step = tf.Variable(1.0, dtype=tf.float64, trainable=False, name='global_step') warmup_steps = tf.constant(self.warmup_periods * self.steps_per_period, dtype=tf.float64, name='warmup_steps') train_steps = tf.constant( (self.first_stage_epochs + self.second_stage_epochs) * self.steps_per_period, dtype=tf.float64, name='train_steps') self.learn_rate = tf.cond( pred=self.global_step < warmup_steps, true_fn=lambda: self.global_step / warmup_steps * self. learn_rate_init, false_fn=lambda: self.learn_rate_end + 0.5 * (self.learn_rate_init - self.learn_rate_end) * (1 + tf.cos( (self.global_step - warmup_steps) / (train_steps - warmup_steps) * np.pi))) global_step_update = tf.assign_add(self.global_step, 1.0) with tf.name_scope("define_weight_decay"): moving_ave = tf.train.ExponentialMovingAverage( self.moving_ave_decay).apply(tf.trainable_variables()) with tf.name_scope("define_first_stage_train"): self.first_stage_trainable_var_list = [] for var in tf.trainable_variables(): var_name = var.op.name var_name_mess = str(var_name).split('/') if var_name_mess[0] in [ 'conv_sbbox', 'conv_mbbox', 'conv_lbbox' ]: self.first_stage_trainable_var_list.append(var) first_stage_optimizer = tf.train.AdamOptimizer( self.learn_rate).minimize( self.loss, var_list=self.first_stage_trainable_var_list) with tf.control_dependencies( tf.get_collection(tf.GraphKeys.UPDATE_OPS)): with tf.control_dependencies( [first_stage_optimizer, global_step_update]): with tf.control_dependencies([moving_ave]): self.train_op_with_frozen_variables = tf.no_op() with tf.name_scope("define_second_stage_train"): second_stage_trainable_var_list = tf.trainable_variables() second_stage_optimizer = tf.train.AdamOptimizer( self.learn_rate).minimize( self.loss, var_list=second_stage_trainable_var_list) with tf.control_dependencies( tf.get_collection(tf.GraphKeys.UPDATE_OPS)): with tf.control_dependencies( [second_stage_optimizer, global_step_update]): with tf.control_dependencies([moving_ave]): self.train_op_with_all_variables = tf.no_op() with tf.name_scope('loader_and_saver'): self.loader = tf.train.Saver(self.net_var) self.saver = tf.train.Saver(tf.global_variables(), max_to_keep=2) with tf.name_scope('summary'): tf.summary.scalar("learn_rate", self.learn_rate) tf.summary.scalar("giou_loss", self.giou_loss) tf.summary.scalar("conf_loss", self.conf_loss) tf.summary.scalar("prob_loss", self.prob_loss) tf.summary.scalar("total_loss", self.loss) logdir = "./data/log/" if os.path.exists(logdir): shutil.rmtree(logdir) os.mkdir(logdir) self.write_op = tf.summary.merge_all() self.summary_writer = tf.summary.FileWriter(logdir, graph=self.sess.graph)
def create_model( config, max_box_per_image, warmup_batches, multi_gpu, saved_weights_name, lr ): if config["model"]["model_name"] == "yolov3": print('[INFO] YOLOV3 Model Creating...') if multi_gpu > 1: with tf.device('/cpu:0'): yolo_model = YOLOV3( config=config, max_box_per_image=max_box_per_image, batch_size=config["train"]["batch_size"] // multi_gpu, warmup_batches=warmup_batches) template_model, infer_model = yolo_model.model() else: yolo_model = YOLOV3( config=config, max_box_per_image=max_box_per_image, batch_size=config["train"]["batch_size"], warmup_batches=warmup_batches) template_model, infer_model = yolo_model.model() elif config["model"]["model_name"] == "yolov4": print('[INFO] YOLOV4 Model Creating...') if multi_gpu > 1: with tf.device('/cpu:0'): yolo_model = YOLOV4( config=config, max_box_per_image=max_box_per_image, batch_size=config["train"]["batch_size"] // multi_gpu, warmup_batches=warmup_batches) template_model, infer_model = yolo_model.model() else: yolo_model = YOLOV4( config=config, max_box_per_image=max_box_per_image, batch_size=config["train"]["batch_size"], warmup_batches=warmup_batches) template_model, infer_model = yolo_model.model() else: pass # load the pretrained weight if exists, otherwise load the backend weight only if os.path.exists(saved_weights_name): print("[INFO] Find pretrained weights...") print("\n[INFO] Loading pretrained weights...\n") template_model.load_weights(saved_weights_name) # else: # template_model.load_weights("backend.h5", by_name=True) if multi_gpu > 1: train_model = multi_gpu_model(template_model, gpus=multi_gpu) else: train_model = template_model optimizer = Adam(lr=lr, clipnorm=0.001) train_model.compile(loss=dummy_loss, optimizer=optimizer) return train_model, infer_model