def __init__(self, **kwargs): self.__dict__.update(self._defaults) for name, value in kwargs.items(): setattr(self, name, value) #---------------------------------------------------# # 获得种类和先验框的数量 #---------------------------------------------------# self.class_names, self.num_classes = get_classes(self.classes_path) self.anchors, self.num_anchors = get_anchors(self.anchors_path) #---------------------------------------------------# # 画框设置不同的颜色 #---------------------------------------------------# hsv_tuples = [(x / self.num_classes, 1., 1.) for x in range(self.num_classes)] self.colors = list(map(lambda x: colorsys.hsv_to_rgb(*x), hsv_tuples)) self.colors = list( map(lambda x: (int(x[0] * 255), int(x[1] * 255), int(x[2] * 255)), self.colors)) self.input_image_shape = K.placeholder(shape=(2, )) self.sess = K.get_session() self.boxes, self.scores, self.classes = self.generate() show_config(**self._defaults)
def __init__(self, **kwargs): super(YOLO, self).__init__() self.__dict__.update(self._defaults) # 还能这样? self.__dict__.update(kwargs) # update with user overrides self.class_names = get_classes(self.classes_path) self.anchors = get_anchors(self.anchors_path) self.colors = get_colors(self.class_names) K.set_learning_phase(0) self.inference_model = self._generate_model()
def __init__(self): anchors = utils.get_anchors(flags.FLAGS.anchors_path) num_anchors = len(anchors) anchors = np.array(anchors, dtype=np.float32) super(CocoModel, self).__init__( image_size=flags.FLAGS.image_size, image_channels=flags.FLAGS.image_channels, num_classes=flags.FLAGS.num_classes, anchors=anchors, batch_size=distribution_utils.per_device_batch_size( flags.FLAGS.batch_size, flags_core.get_num_gpus(flags.FLAGS)), num_anchors=num_anchors, learning_rate=flags.FLAGS.learning_rate, backbone=flags.FLAGS.backbone, norm=flags.FLAGS.norm, threshold=flags.FLAGS.threshold, max_num_boxes_per_image=flags.FLAGS.max_num_boxes_per_image, confidence_score=flags.FLAGS.confidence_score, data_format=flags.FLAGS.data_format, dtype=flags_core.get_tf_dtype(flags.FLAGS))
def __init__(self, **kwargs): self.__dict__.update(self._defaults) for name, value in kwargs.items(): setattr(self, name, value) self._defaults[name] = value #---------------------------------------------------# # 获得种类和先验框的数量 #---------------------------------------------------# self.class_names, self.num_classes = get_classes(self.classes_path) self.anchors, self.num_anchors = get_anchors(self.anchors_path) self.bbox_util = DecodeBox(self.anchors, self.num_classes, (self.input_shape[0], self.input_shape[1]), self.anchors_mask) #---------------------------------------------------# # 画框设置不同的颜色 #---------------------------------------------------# hsv_tuples = [(x / self.num_classes, 1., 1.) for x in range(self.num_classes)] self.colors = list(map(lambda x: colorsys.hsv_to_rgb(*x), hsv_tuples)) self.colors = list(map(lambda x: (int(x[0] * 255), int(x[1] * 255), int(x[2] * 255)), self.colors)) self.generate() show_config(**self._defaults)
# 判断当前使用的GPU数量与机器上实际的GPU数量 #------------------------------------------------------# if ngpus_per_node > 1 and ngpus_per_node > len(gpus): raise ValueError("The number of GPUs specified for training is more than the GPUs on the machine") if ngpus_per_node > 1: strategy = tf.distribute.MirroredStrategy() else: strategy = None print('Number of devices: {}'.format(ngpus_per_node)) #----------------------------------------------------# # 获取classes和anchor #----------------------------------------------------# class_names, num_classes = get_classes(classes_path) anchors, num_anchors = get_anchors(anchors_path) #----------------------------------------------------# # 判断是否多GPU载入模型和预训练权重 #----------------------------------------------------# if ngpus_per_node > 1: with strategy.scope(): #------------------------------------------------------# # 创建yolo模型 #------------------------------------------------------# model_body = yolo_body((None, None, 3), anchors_mask, num_classes, backbone, alpha, weight_decay=weight_decay) if model_path != '': #------------------------------------------------------# # 载入预训练权重 #------------------------------------------------------# print('Load weights {}.'.format(model_path))
def _main(): # parse command line arguments parser = argparse.ArgumentParser() parser.add_argument( '--train_from_checkpoint', type=str, help= "The path to where a previously trained model's weights are stored. To use the default\n\ coco weights, use the path 'model_weights/coco_pretrained_weights.ckpt'. Otherwise, the model\n\ weights will be initialized randomly. ") parser.add_argument( '--class_path', default='utils/coco_classes.txt', type=str, help= 'The path that points towards where the class names for the dataset are stored.\n\ The default path is "utils/coco_classes.txt".') parser.add_argument( '--anchors_path', default='utils/anchors.txt', type=str, help= 'The path that points towards where the anchor values for the model are stored.\n\ The default path is "utils/anchors.txt", which contains anchors trained on the coco dataset.' ) parser.add_argument( '--data_path', default='training_data/image_paths_and_box_info.txt', type=str, help= 'The path that points towards where the training data text file is stored.\n\ The default path is "training_data/image_paths_and_box_info.txt".') parser.add_argument( '--input_height', default=416, type=int, help= 'The input height of the yolov3 model. The height must be a multiple of 32.\n\ The default height is 416.') parser.add_argument( '--input_width', default=416, type=int, help= 'The input width of the yolov3 model. The width must be a mutliple of 32.\n\ The default width is 416.') parser.add_argument( '--batch_size', default=32, type=int, help= 'The training batch size, whose default value is set to 32 images per batch.' ) parser.add_argument( '--max_num_boxes_per_image', default=20, type=int, help= 'The max number of boxes that can be detected within one image. Default is 20.' ) parser.add_argument( '--num_training_epochs', default=150, type=int, help='The number of training epochs. The default is 150.') parser.add_argument( '--learning_rate', default=0.001, type=float, help='The learning rate of the model. The default is 0.001.') parser.add_argument( '--ignore_threshold', default=0.5, type=float, help= 'Impacts how the loss is calculated. Must be between zero and one, and the default is set to 0.5.' ) parser.add_argument( '--train_val_data_split', default=0.9, type=float, help= 'The split between the data that will be used for training and data that will be used\n\ for validation. Default value is 0.9.') parser.add_argument( '--train_save_path', default='model_weights/', help= "The training model's checkpoint save path. The default path is 'model_weights/'." ) parser.add_argument( '--model_name', default='model.ckpt', help= 'The name that should be given to the checkpoint file. The default name is "model.ckpt".' ) parser.add_argument( '--tensorboard_save_path', default='tensorboard/tensorboard_train/', help= 'The path where the event files to be used with tensorboard will be saved at. The default\n\ path is "tensorboard/tensorboard_train/".') parser.add_argument( '--test_model_overfit', nargs='?', default=False, type=str2bool, const=True, help= 'Whether or not to purposefully overfit the model by training it on only one image.\n\ This option is useful in testing out if the loss function is working correctly.' ) parser.add_argument( '--save_every_x_iterations', default=100, type=int, help= "How frequently the model's training weights are saved. The default value is every\n\ 100 iterations.") parser.add_argument( '--log_every_x_iterations', default=5, type=int, help= "How frequently the model's loss is logged for it to be inspected in Tensorboard.\n\ The default value is every 5 iterations.") args = vars(parser.parse_args()) args[ 'data_path'] = '/home/yl/CNN/Yolo/keras-yolo3-fine-tune/dataset/train_label.txt' # read inputs h = args['input_height'] w = args['input_width'] ignore_thresh = args['ignore_threshold'] max_num_boxes_per_image = args['max_num_boxes_per_image'] anchors = get_anchors(args['anchors_path']) lr = args['learning_rate'] num_anchors_per_detector = len(anchors) // 3 num_detectors_per_image = num_anchors_per_detector * (((h / 32) * (w / 32)) + ((h / 16) * (w / 16)) + ((h / 8) * (w / 8))) class_names = get_classes(args['class_path']) num_classes = len(class_names) tb_train_path = args['tensorboard_save_path'] + 'train/' tb_val_path = args['tensorboard_save_path'] + 'val/' training_data, validation_data, batch_size = prepare_data( args['data_path'], args['train_val_data_split'], args['batch_size'], args['test_model_overfit']) tf.reset_default_graph() # build graph with tf.variable_scope('y_true'): y_true_data = tf.placeholder( dtype=tf.float32, shape=[None, num_detectors_per_image, num_classes + 5]) with tf.variable_scope('y_true_boxes'): y_true_box_data = tf.placeholder(dtype=tf.float32, shape=[ None, max_num_boxes_per_image * num_anchors_per_detector, 4 ]) with tf.variable_scope('x_input'): X = tf.placeholder(dtype=tf.float32, shape=[None, h, w, 3]) yolo_outputs = yolo_v3(inputs=X, num_classes=len(class_names), anchors=anchors, h=h, w=w, training=True) # output loss = yolo_v3_loss(yolo_outputs, y_true_data, y_true_box_data, ignore_threshold=ignore_thresh, anchors=anchors, num_classes=num_classes, h=h, w=w, batch_size=batch_size) tf.summary.scalar('loss', loss) global_step = tf.get_variable(name='global_step', trainable=False, initializer=0, dtype=tf.int32) # returns a varlist containing only the vars of the conv layers right before the yolo layers trainable_var_list = tf.trainable_variables() last_layer_var_list = [ i for i in trainable_var_list if i.shape[-1] == (5 + num_classes) * num_anchors_per_detector ] train_op_with_frozen_variables = tf.train.AdamOptimizer( learning_rate=lr).minimize(loss, global_step=global_step, var_list=last_layer_var_list) train_op_with_all_variables = tf.train.AdamOptimizer( learning_rate=lr).minimize(loss, global_step=global_step, var_list=trainable_var_list) summ = tf.summary.merge_all() # info print('--info--') print('model weights will be saved with filename: ', args['model_name']) print('tensorboard event files located at path: ', args['tensorboard_save_path']) # build training loop with tf.Session() as sess: train_writer = tf.summary.FileWriter(tb_train_path, sess.graph) val_writer = tf.summary.FileWriter(tb_val_path) # initialize model weights either randomly or from a saved checkpoint saver = tf.train.Saver() if args['train_from_checkpoint'] is None: print('initializing variables...') sess.run(tf.global_variables_initializer()) else: print('restoring weights from checkpoint: ', args['train_from_checkpoint']) saver.restore(sess, args['train_from_checkpoint']) num_iterations = args['num_training_epochs'] * len(training_data) print('beginning to train the model...') for i in range(num_iterations): input_images, y_true, y_true_boxes = get_training_batch( training_data, anchors, num_classes, batch_size=batch_size, h=h, w=w, random=not args['test_model_overfit']) # For the first epochs, train with the frozen layers. Then, unfreeze the entire graph. if i < num_iterations // 3: sess.run(train_op_with_frozen_variables, feed_dict={ X: input_images, y_true_data: y_true, y_true_box_data: y_true_boxes }) else: sess.run(train_op_with_all_variables, feed_dict={ X: input_images, y_true_data: y_true, y_true_box_data: y_true_boxes }) if i % args['log_every_x_iterations'] == 0: # write the training loss to tensorboard lt, st = sess.run( [loss, summ], feed_dict={ X: input_images, y_true_data: y_true, y_true_box_data: y_true_boxes }) train_writer.add_summary(st, i) #write the validation loss to tensorboard if we are not in overfit mode if not args['test_model_overfit']: input_images, y_true, y_true_boxes = get_training_batch( validation_data, anchors, num_classes, batch_size=batch_size, h=h, w=w, random=not args['test_model_overfit']) lv, sv = sess.run( [loss, summ], feed_dict={ X: input_images, y_true_data: y_true, y_true_box_data: y_true_boxes }) val_writer.add_summary(sv, i) print("iteration: " + str(i) + ", training loss: " + str(round(lt, 2)) + ", validation loss: " + str(round(lv, 2))) else: print("iteration: " + str(i) + ", training loss: " + str(round(lt, 2))) if i % args['save_every_x_iterations'] == 0: print('saving model weights at path: ', args['train_save_path']) saver.save( sess, os.path.join(args['train_save_path'], args['model_name']), global_step) train_writer.close() val_writer.close()
def parse_fn(image_id, dataset, anchors_path, augmentation=None, dtype=np.float32, max_num_boxes_per_image=20, image_size=416): """Load and return ground truth data for an image (image, mask, bounding boxes).""" image = dataset.load_image(image_id) # original_shape = image.shape image, window, scale, padding, crop = utils.resize_image( image, min_dim=0, min_scale=0, max_dim=image_size, mode='square') mask, class_ids = dataset.load_mask(image_id) mask = utils.resize_mask(mask, scale, padding, crop) if augmentation: import imgaug # Augmenters that are safe to apply to masks # Some, such as Affine, have settings that make them unsafe, so always # test your augmentation on masks MASK_AUGMENTERS = [ "Sequential", "SomeOf", "OneOf", "Sometimes", "Fliplr", "Flipud", "CropAndPad", "Affine", "PiecewiseAffine" ] def hook(images, augmenter, parents, default): """Determines which augmenters to apply to masks.""" return augmenter.__class__.__name__ in MASK_AUGMENTERS # Store shapes before augmentation to compare image_shape = image.shape mask_shape = mask.shape # Make augmenters deterministic to apply similarly to images and masks det = augmentation.to_deterministic() image = det.augment_image(image) # Change mask to np.uint8 because imgaug doesn't support np.bool mask = det.augment_image(mask.astype(np.uint8), hooks=imgaug.HooksImages(activator=hook)) # Verify that shapes didn't change assert image.shape == image_shape, "Augmentation shouldn't change image size" assert mask.shape == mask_shape, "Augmentation shouldn't change mask size" # Change mask back to bool mask = mask.astype(np.bool) # Note that some boxes might be all zeros if the corresponding mask got cropped out. # and here is to filter them out _idx = np.sum(mask, axis=(0, 1)) > 0 mask = mask[:, :, _idx] class_ids = class_ids[_idx] # Bounding boxes. Note that some boxes might be all zeros # if the corresponding mask got cropped out. # bbox: [num_instances, (y1, x1, y2, x2)] bbox = utils.extract_bboxes(mask) if mask.shape[-1] > max_num_boxes_per_image: ids = np.random.choice(np.arange(mask.shape[-1]), max_num_boxes_per_image, replace=False) class_ids = class_ids[ids] bbox = bbox[ids, :] # confs = np.ones((bbox.shape[0], 1), dtype=dtype) # bbox = np.concatenate([bbox, confs], axis=-1) # Active classes # Different datasets have different classes, so track the # classes supported in the dataset of this image. # active_class_ids = np.zeros([dataset.num_classes], dtype=np.int32) # source_class_ids = dataset.source_class_ids[dataset.image_info[image_id]["source"]] # active_class_ids[source_class_ids] = 1 # image_meta = utils.compose_image_meta(image_id, original_shape, image.shape, # window, scale, active_class_ids) # image_meta.astype(dtype) # gt_mask = np.zeros((mask.shape[0], mask.shape[1], 20), mask.dtype) gt_class_ids = np.zeros(max_num_boxes_per_image, class_ids.dtype) gt_bbox = np.zeros((max_num_boxes_per_image, bbox.shape[1]), bbox.dtype) # gt_data = np.zeros((max_num_boxes_per_image, bbox.shape[1] + dataset.num_classes), dtype=dtype) if class_ids.shape[0] > 0: gt_class_ids[:class_ids.shape[0]] = class_ids # gt_mask[:, :, :mask.shape[-1]] = mask gt_bbox[:bbox.shape[0], :] = bbox gt_class_ids = np.expand_dims(gt_class_ids, axis=-1).astype(dtype) gt_bbox = np.concatenate([gt_bbox, gt_class_ids], axis=-1) anchors = utils.get_anchors(anchors_path) anchors = np.array(anchors, dtype=np.float32) boxes_yx = (gt_bbox[:, 0:2] + gt_bbox[:, 2:4]) // 2 boxes_hw = gt_bbox[:, 2:4] - gt_bbox[:, 0:2] gt_bbox[:, 0] = boxes_yx[..., 1] / image_size gt_bbox[:, 1] = boxes_yx[..., 0] / image_size gt_bbox[:, 2] = boxes_hw[..., 1] / image_size gt_bbox[:, 3] = boxes_hw[..., 0] / image_size hw = np.expand_dims(boxes_hw, -2) anchors_broad = np.expand_dims(anchors, 0) anchor_maxes = anchors_broad / 2. anchor_mins = -anchor_maxes box_maxes = hw / 2. box_mins = -box_maxes intersect_mins = np.maximum(box_mins, anchor_mins) intersect_maxes = np.minimum(box_maxes, anchor_maxes) intersect_hw = np.maximum(intersect_maxes - intersect_mins, 0.) intersect_area = intersect_hw[..., 0] * intersect_hw[..., 1] box_area = hw[..., 0] * hw[..., 1] anchor_area = anchors[..., 0] * anchors[..., 1] iou = intersect_area / (box_area + anchor_area - intersect_area) best_anchors = np.argmax(iou, axis=-1) # TODO: write a function to calculate the stride automatically. large_obj_image_size = image_size // 32 medium_obj_image_size = image_size // 16 small_obj_image_size = image_size // 8 large_obj_detectors, large_obj_boxes = get_detector_heatmap_each_scale( gt_bbox, best_anchors_=best_anchors, anchors_mask=[6, 7, 8], grid_size=(large_obj_image_size, large_obj_image_size), num_classes=dataset.num_classes) medium_obj_detectors, medium_obj_boxes = get_detector_heatmap_each_scale( gt_bbox, best_anchors_=best_anchors, anchors_mask=[3, 4, 5], grid_size=(medium_obj_image_size, medium_obj_image_size), num_classes=dataset.num_classes) small_obj_detectors, small_obj_boxes = get_detector_heatmap_each_scale( gt_bbox, best_anchors_=best_anchors, anchors_mask=[0, 1, 2], grid_size=(small_obj_image_size, small_obj_image_size), num_classes=dataset.num_classes) yolo_true_data = np.concatenate( [large_obj_detectors, medium_obj_detectors, small_obj_detectors], axis=0).reshape([-1]) yolo_true_boxes = np.concatenate( [large_obj_boxes, medium_obj_boxes, small_obj_boxes], axis=0).reshape([-1]) yolo_gt = np.concatenate([yolo_true_data, yolo_true_boxes], axis=-1) return image.astype(dtype) / 255., yolo_gt.astype(dtype)
def _main(): # parse command line arguments parser = argparse.ArgumentParser() requiredNamed = parser.add_argument_group('required named arguments') requiredNamed.add_argument( '--path_to_input_image', type=str, required=True, help= 'The path to the input image on which object detection will be performed on.\n\ This argument is required.') parser.add_argument( '--path_to_trained_model', default='model_weights/coco_pretrained_weights.ckpt', type=str, help= "The path to the location of pretrained model weights, which will be loaded into\n\ the model and then used for object detection. The default pretrained weights path is\n\ 'model_weights/coco_pretrained_weights.ckpt', which contains weights trained on\n\ the coco dataset.") parser.add_argument( '--save_as', type=str, default=None, help= 'The filename for the image on which object detection was performed. If no filename\n\ is provided, the image will be saved as "[original_name] + _yolo_v3.jpg".' ) parser.add_argument('--tensorboard_save_path', default='tensorboard/tensorboard_detect/', help="") parser.add_argument( '--class_path', default='utils/coco_classes.txt', type=str, help= 'The path that points towards where the class names for the dataset are stored.\n\ The default path is "utils/coco_classes.txt".') parser.add_argument( '--anchors_path', default='utils/anchors.txt', type=str, help= 'The path that points towards where the anchor values for the model are stored.\n\ The default path is "utils/anchors.txt", which contains anchors trained on the coco dataset.' ) parser.add_argument( '--input_height', default=416, type=int, help= 'The input height of the yolov3 model. The height must be a multiple of 32.\n\ The default height is 416.') parser.add_argument( '--input_width', default=416, type=int, help= 'The input width of the yolov3 model. The width must be a mutliple of 32.\n\ The default width is 416.') args = vars(parser.parse_args()) h = args['input_height'] w = args['input_width'] anchors = get_anchors(args['anchors_path']) classes = get_classes(args['class_path']) save_as = args['save_as'] if save_as is None: filename_w_ext = os.path.basename(args['path_to_input_image']) filename, file_extension = os.path.splitext(filename_w_ext) save_as = filename + '_yolo_v3' + file_extension image, original_im = process_image(args['path_to_input_image'], h, w) tf.reset_default_graph() # build graph with tf.variable_scope('x_input'): X = tf.placeholder(dtype=tf.float32, shape=[None, h, w, 3]) yolo_outputs = yolo_v3(inputs=X, num_classes=len(classes), anchors=anchors, h=h, w=w, training=False) # output with tf.variable_scope('obj_detections'): raw_outputs = tf.concat(yolo_outputs, axis=1) # pass image through model with tf.Session() as sess: writer = tf.summary.FileWriter(args['tensorboard_save_path'], sess.graph) writer.close() saver = tf.train.Saver() print('restoring model weights...') saver.restore(sess, save_path=args['path_to_trained_model']) print('feeding image found at filepath: ', args['path_to_input_image']) start = time.time() ro = sess.run(raw_outputs, feed_dict={X: [np.array(image, dtype=np.float32)]}) end = time.time() total_time = end - start print("total inference time was: " + str(round(total_time, 2)) + " seconds (that's " + str(round(60.0 / total_time, 2)) + " fps!)") # convert box coordinates, apply nms, and draw boxes boxes = convert_box_coordinates(ro) filtered_boxes = non_max_suppression(boxes, confidence_threshold=0.5, iou_threshold=0.4) draw_boxes(save_as, args['class_path'], filtered_boxes, original_im, image) print('image with detections saved as: ', save_as)
def train(self, classes_path, anchors_path): ''' Args: classes_path:classes路径 anchors_path:anchor路径 ''' classes_names = get_classes(classes_path) num_classes = len(classes_names) anchors = get_anchors(anchors_path) is_tiny_version = len(anchors) == 6 # default setting if is_tiny_version: model = create_tiny_model(input_shape, anchors, num_classes, freeze_body=2) else: model = create_model(input_shape, anchors, num_classes, load_pretrained=False) logging = TensorBoard(log_dir=log_dir) # checkpoint = ModelCheckpoint(log_dir + 'car_mobilenet_yolov3.ckpt', # monitor='val_loss', save_weights_only=False, period=1) checkpoint = ModelCheckpoint( log_dir + 'ep{epoch:03d}-loss{loss:.3f}-val_loss{val_loss:.3f}.h5', monitor='val_loss', save_weights_only=False, save_best_only=True, period=3) # reduce_lr = ReduceLROnPlateau(monitor='val_loss', factor=0.1, patience=3, verbose=1) reduce_lr = ReduceLROnPlateau(monitor="val_loss", factor=0.1, min_lr=1e-9, patience=5, verbose=1) # early_stopping = EarlyStopping(monitor='val_loss', min_delta=0, patience=10, verbose=1) with open(train_path) as t_f: t_lines = t_f.readlines() np.random.seed(666) np.random.shuffle(t_lines) num_val = int(len(t_lines) * val_split) num_train = len(t_lines) - num_val t_lines = t_lines[:num_train] v_lines = t_lines[num_train:] # train with frozen layers first ,to get a stable loss. # adjust num epochs to your dataset,This step is enough to obtrain a not bad model if True: model.compile( optimizer=Adam(lr=1e-3), loss={ # use custom yolo_loss Lambda layer. 'yolo_loss': lambda y_true, y_pred: y_pred }) print('Train on {} samples, val on {} samples, with batch size {}.'. format(num_train, num_val, batch_num)) model.fit_generator(data_generator_wrapper(t_lines, batch_num, input_shape, anchors, num_classes), steps_per_epoch=max(1, num_train // batch_num), validation_data=data_generator_wrapper( v_lines, batch_num, input_shape, anchors, num_classes), validation_steps=max(1, num_val // batch_num), epochs=epochs, initial_epoch=0, callbacks=[logging, checkpoint]) model.save(log_dir + 'trained_weights_stage_1.h5') # Unfreeze and continue training, to fine-tune. # Train longer if the result is not good. if True: print("Unfreeze and continue training, to fine-tune.") for i in range(len(model.layers)): model.layers[i].trainable = True model.compile(optimizer=Adam(lr=1e-4), loss={ 'yolo_loss': lambda y_true, y_pred: y_pred }) # recompile to apply the change batch_size = 16 # note that more GPU memory is required after unfreezing the body print( 'Train on {} samples, val on {} samples, with batch size {}.'. format(num_train, num_val, batch_size)) model.fit_generator( data_generator_wrapper(t_lines, batch_size, input_shape, anchors, num_classes), steps_per_epoch=max(1, num_train // batch_size), validation_data=data_generator_wrapper(v_lines, batch_size, input_shape, anchors, num_classes), validation_steps=max(1, num_val // batch_size), epochs=20, initial_epoch=0, callbacks=[logging, checkpoint, reduce_lr]) model.save(log_dir + 'trained_weights_final.h5')
def _main(): # parse command line arguments parser = argparse.ArgumentParser() parser.add_argument( '--train-from-checkpoint', type=str, help="the path to where a previously trained model's weights are stored") parser.add_argument('--class-path', default='utils/coco_classes.txt', type=str, help='the path that points to the class names for the dataset') parser.add_argument( '--anchors-path', default='utils/anchors.txt', type=str, help='the path that points to the anchor values for the models') parser.add_argument( '--data-path', default='data/image_paths_and_box_info.txt', type=str, help='the path that points to the training data text file') parser.add_argument( '--input-height', default=416, type=int, help='the input height of the yolov3 model. The height must be a multiple of 32') parser.add_argument( '--input-width', default=416, type=int, help='The input width of the yolov3 model. The width must be a multiple of 32') parser.add_argument( '--batch-size', default=32, type=int, help='the training batch size') parser.add_argument( '--max-num-boxes-per-image', default=20, type=int, help='the max number of boxes that can be detected within one image') parser.add_argument( '--num-training-epochs', default=150, type=int, help='the number of training epochs') parser.add_argument( '--learning-rate', default=0.001, type=float, help='the learning rate') parser.add_argument( '--ignore-threshold', default=0.5, type=float, help='impacts how the loss is calculated. Must be between zero and one') parser.add_argument( '--train-val-data-split', default=0.9, type=float, help='the split between the data that will be used for training and data that will be used\n\ for validation') parser.add_argument( '--train-save-path', default='model-weights/', help="the training model's checkpoint save path") parser.add_argument( '--model-name', default='model.ckpt', help='the name that should be given to the checkpoint file') parser.add_argument( '--tensorboard-save-path', default='tb-logs/tb-train/', help='the path where the event files to be used with tensorboard will be saved at') parser.add_argument( '--test-model-overfit', nargs='?', default=False, type=str2bool, const=True, help='this option is useful in testing out if the loss function is working correctly') parser.add_argument( '--save-iterations', default=100, type=int, help="how frequently the model's training weights are saved") parser.add_argument( '--log-iterations', default=5, type=int, help="how frequently the model's loss is logged for it to be inspected in Tensorboard") args = parser.parse_args() # args info print('[i] checkpoint: ', args.train_from_checkpoint) print('[i] path of class file: ', args.class_path) print('[i] path of anchors file: ', args.anchors_path) # read inputs h = args.input_height w = args.input_width ignore_thresh = args.ignore_threshold max_num_boxes_per_image = args.max_num_boxes_per_image anchors = get_anchors(args.anchors_path) lr = args.learning_rate num_anchors_per_detector = len(anchors) // 3 num_detectors_per_image = num_anchors_per_detector * ( ((h / 32) * (w / 32)) + ((h / 16) * (w / 16)) + ((h / 8) * (w / 8))) class_names = get_classes(args.class_path) num_classes = len(class_names) # tensorboard path tb_train_path = args.tensorboard_save_path + 'train/' tb_val_path = args.tensorboard_save_path + 'val/' training_data, validation_data, batch_size = prepare_data( args.data_path, args.train_val_data_split, args.batch_size, args.test_model_overfit) tf.reset_default_graph() # build graph with tf.variable_scope('y_true'): y_true_data = tf.placeholder(dtype=tf.float32, shape=[None, num_detectors_per_image, num_classes + 5]) with tf.variable_scope('y_true_boxes'): y_true_box_data = tf.placeholder(dtype=tf.float32, shape=[None, max_num_boxes_per_image * num_anchors_per_detector, 4]) with tf.variable_scope('x_input'): X = tf.placeholder(dtype=tf.float32, shape=[None, h, w, 3]) yolo_outputs = yolo_v3(inputs=X, num_classes=len(class_names), anchors=anchors, h=h, w=w, training=True) # output loss = yolo_v3_loss(yolo_outputs, y_true_data, y_true_box_data, ignore_threshold=ignore_thresh, anchors=anchors, num_classes=num_classes, h=h, w=w, batch_size=batch_size) tf.summary.scalar('loss', loss) global_step = tf.get_variable(name='global_step', trainable=False, initializer=0, dtype=tf.int32) # returns a varlist containing only the vars of the conv layers right before the yolo layers trainable_var_list = tf.trainable_variables() last_layer_var_list = [i for i in trainable_var_list if i.shape[-1] == (5 + num_classes) * num_anchors_per_detector] train_op_with_frozen_variables = tf.train.AdamOptimizer(learning_rate=lr).minimize(loss, global_step=global_step, var_list=last_layer_var_list) train_op_with_all_variables = tf.train.AdamOptimizer(learning_rate=lr).minimize(loss, global_step=global_step, var_list=trainable_var_list) summ = tf.summary.merge_all() # info print('-------info-------') print('[i] model weights will be saved with filename: ', args.model_name) print('[i] tensorboard event files located at path: ', args.tensorboard_save_path) # build training loop with tf.Session() as sess: train_writer = tf.summary.FileWriter(tb_train_path, sess.graph) val_writer = tf.summary.FileWriter(tb_val_path) # initialize model weights either randomly or from a saved checkpoint saver = tf.train.Saver() if args['train_from_checkpoint'] is None: print('[i] initializing variables...') sess.run(tf.global_variables_initializer()) else: print('[i] restoring weights from checkpoint: ', args.train_from_checkpoint) saver.restore(sess, args.train_from_checkpoint) num_iterations = args.num_epochs * len(training_data) print('[i] beginning to train the model...') for i in range(num_iterations): input_images, y_true, y_true_boxes = get_training_batch(training_data, anchors, num_classes, batch_size=batch_size, h=h, w=w, random=not args.test_model_overfit) # For the first epochs, train with the frozen layers. Then, unfreeze the entire graph. if i < num_iterations // 3: sess.run(train_op_with_frozen_variables, feed_dict={X: input_images, y_true_data: y_true, y_true_box_data: y_true_boxes}) else: sess.run(train_op_with_all_variables, feed_dict={X: input_images, y_true_data: y_true, y_true_box_data: y_true_boxes}) if i % args.log_iterations == 0: # write the training loss to tensorboard lt, st = sess.run([loss, summ], feed_dict={X: input_images, y_true_data: y_true, y_true_box_data: y_true_boxes}) train_writer.add_summary(st, i) # write the validation loss to tensorboard if we are not in overfit mode if not args.test_model_overfit: input_images, y_true, y_true_boxes = get_training_batch(validation_data, anchors, num_classes, batch_size=batch_size, h=h, w=w, random=not args.test_model_overfit) lv, sv = sess.run([loss, summ], feed_dict={X: input_images, y_true_data: y_true, y_true_box_data: y_true_boxes}) val_writer.add_summary(sv, i) print( "| iteration: " + str(i) + ", training loss: " + str(round(lt, 2)) + ", validation loss: " + str( round(lv, 2))) else: print("| iteration: " + str(i) + ", training loss: " + str(round(lt, 2))) if i % args.save_iterations == 0: print('[i] saving model weights at path: ', args.train_save_path) saver.save(sess, os.path.join(args.train_save_path, args.model_name), global_step) print('################################################################################') train_writer.close() val_writer.close()
def yolo_main(flags_obj, model_function, input_function, dataset, augmentation): """Shared main loop for yolo Models. Args: flags_obj: An object containing parsed flags. See define_yolo_flags() for details. model_function: the function that instantiates the Model and builds the ops for train/eval. This will be passed directly into the estimator. input_function: the function that processes the dataset and returns a dataset that the estimator can train on. This will be wrapped with all the relevant flags for running and passed to estimator. dataset: A dataset for training and evaluation. augmentation: Optional. An imgaug (https://github.com/aleju/imgaug) augmentation. For example, passing imgaug.augmenters.Fliplr(0.5) flips images right/left 50% of the time. """ model_helpers.apply_clean(flags_obj) # Ensures flag override logic is only executed if explicitly triggered. if flags_obj.tf_gpu_thread_mode: override_flags_and_set_envars_for_gpu_thread_pool(flags_obj) # Creates session config. allow_soft_placement = True, is required for # multi-GPU and is not harmful for other modes. session_config = tf.ConfigProto( log_device_placement=True, inter_op_parallelism_threads=flags_obj.inter_op_parallelism_threads, intra_op_parallelism_threads=flags_obj.intra_op_parallelism_threads, allow_soft_placement=True) session_config.gpu_options.allow_growth = True distribution_strategy = distribution_utils.get_distribution_strategy( flags_core.get_num_gpus(flags_obj), flags_obj.all_reduce_alg) run_config = tf.estimator.RunConfig(train_distribute=distribution_strategy, session_config=session_config, save_checkpoints_secs=60 * 60 * 24) # Initializes model with all but the dense layer from pretrained ResNet. if flags_obj.pretrained_model_checkpoint_path is not None: warm_start_settings = tf.estimator.WarmStartSettings( flags_obj.pretrained_model_checkpoint_path, vars_to_warm_start='^(?!.*dense)') else: warm_start_settings = None anchors = np.array(utils.get_anchors(flags_obj.anchors_path)) detector = tf.estimator.Estimator( model_fn=model_function, model_dir=flags_obj.model_dir, config=run_config, warm_start_from=warm_start_settings, params={ 'num_classes': flags_obj.num_classes, 'data_format': flags_obj.data_format, 'batch_size': distribution_utils.per_device_batch_size( flags_obj.batch_size, flags_core.get_num_gpus(flags_obj)), 'image_size': int(flags_obj.image_size), 'loss_scale': flags_core.get_loss_scale(flags_obj), 'dtype': flags_core.get_tf_dtype(flags_obj), 'fine_tune': flags_obj.fine_tune, 'anchors': anchors, 'num_anchors': len(anchors), 'max_num_boxes_per_image': flags_obj.max_num_boxes_per_image, 'threshold': flags_obj.threshold, 'train': dataset.num_images, 'learning_rate': flags_obj.learning_rate }) # if flags_obj.use_synthetic_data: # dataset_name = dataset_name + '-synthetic' def input_fn_train(num_epochs): return input_function( data_set=dataset, is_training=True, batch_size=distribution_utils.per_device_batch_size( flags_obj.batch_size, flags_core.get_num_gpus(flags_obj)), anchors_path=flags_obj.anchors_path, num_epochs=num_epochs, augmentation=augmentation, dtype=tf.float32, max_num_boxes_per_image=flags_obj.max_num_boxes_per_image, image_size=flags_obj.image_size, datasets_num_private_threads=flags_obj. datasets_num_private_threads, num_parallel_batches=flags_obj.datasets_num_parallel_batches) ''' def input_fn_eval(): return input_function( is_training=False, data_dir=flags_obj.data_dir, batch_size=distribution_utils.per_device_batch_size( flags_obj.batch_size, flags_core.get_num_gpus(flags_obj)), num_epochs=1, dtype=flags_core.get_tf_dtype(flags_obj)) ''' if flags_obj.eval_only or not flags_obj.train_epochs: # If --eval_only is set, perform a single loop with zero train epochs. schedule, n_loops = [0], 1 else: # Compute the number of times to loop while training. All but the last # pass will train for `epochs_between_evals` epochs, while the last will # train for the number needed to reach `training_epochs`. For instance if # train_epochs = 25 and epochs_between_evals = 10 # schedule will be set to [10, 10, 5]. That is to say, the loop will: # Train for 10 epochs and then evaluate. # Train for another 10 epochs and then evaluate. # Train for a final 5 epochs (to reach 25 epochs) and then evaluate. n_loops = math.ceil(flags_obj.train_epochs / flags_obj.epochs_between_evals) schedule = [ flags_obj.epochs_between_evals for _ in range(int(n_loops)) ] schedule[-1] = flags_obj.train_epochs - sum( schedule[:-1]) # over counting. for cycle_index, num_train_epochs in enumerate(schedule): tf.logging.info('Starting cycle: %d/%d', cycle_index, int(n_loops)) if num_train_epochs: detector.train(input_fn=lambda: input_fn_train(num_train_epochs), max_steps=flags_obj.max_train_steps) '''