def main(args): # 1、判断传入的是需要训练YoloV3Tiny还是YOLOV3正常版本 if args.tiny: model = YoloV3Tiny(args.size, training=True, classes=args.num_classes) anchors = yolo_tiny_anchors anchor_masks = yolo_tiny_anchor_masks else: model = YoloV3(args.size, training=True, classes=args.num_classes) anchors = yolo_anchors anchor_masks = yolo_anchor_masks # 2、获取传入参数的训练数据 if args.dataset: train_dataset = dataset.load_tfrecord_dataset(args.dataset, args.classes) train_dataset = train_dataset.shuffle(buffer_size=1024) train_dataset = train_dataset.batch(args.batch_size) train_dataset = train_dataset.map( lambda x, y: (dataset.transform_images(x, args.size), dataset.transform_targets(y, anchors, anchor_masks, 80))) train_dataset = train_dataset.prefetch( buffer_size=tf.data.experimental.AUTOTUNE) # 3、获取传入参数的验证集数据 if args.val_dataset: val_dataset = dataset.load_tfrecord_dataset(args.val_dataset, args.classes) val_dataset = val_dataset.batch(args.batch_size) val_dataset = val_dataset.map( lambda x, y: (dataset.transform_images(x, args.size), dataset.transform_targets(y, anchors, anchor_masks, 80))) # 4、判断是否进行迁移学习 if args.transfer != 'none': # 加载与训练模型'./data/yolov3.weights' model.load_weights(args.weights) if args.transfer == 'fine_tune': # 冻结darknet darknet = model.get_layer('yolo_darknet') freeze_all(darknet) elif args.mode == 'frozen': # 冻结所有层 freeze_all(model) else: # 重置网络后端结构 if args.tiny: init_model = YoloV3Tiny(args.size, training=True, classes=args.num_classes) else: init_model = YoloV3(args.size, training=True, classes=args.num_classes) # 如果迁移指的是darknet if args.transfer == 'darknet': # 获取网络的权重 for l in model.layers: if l.name != 'yolo_darknet' and l.name.startswith('yolo_'): l.set_weights( init_model.get_layer(l.name).get_weights()) else: freeze_all(l) elif args.transfer == 'no_output': for l in model.layers: if l.name.startswith('yolo_output'): l.set_weights( init_model.get_layer(l.name).get_weights()) else: freeze_all(l) # 5、定义优化器以及损失计算方式 optimizer = tf.keras.optimizers.Adam(lr=args.learning_rate) loss = [ YoloLoss(anchors[mask], classes=args.num_classes) for mask in anchor_masks ] # 6、训练优化过程,训练指定模式 # 用eager模式进行训练易于调试 # keras model模式简单易用 if args.mode == 'eager_tf': # 定义评估方式 avg_loss = tf.keras.metrics.Mean('loss', dtype=tf.float32) avg_val_loss = tf.keras.metrics.Mean('val_loss', dtype=tf.float32) # 迭代优化 for epoch in range(1, args.epochs + 1): for batch, (images, labels) in enumerate(train_dataset): with tf.GradientTape() as tape: # 1、计算模型输出和损失 outputs = model(images, training=True) regularization_loss = tf.reduce_sum(model.losses) pred_loss = [] for output, label, loss_fn in zip(outputs, labels, loss): # 根据输出和标签计算出损失 pred_loss.append(loss_fn(label, output)) # 计算总损失 = 平均损失 + regularization_loss total_loss = tf.reduce_sum(pred_loss) + regularization_loss # 计算梯度以及更新梯度 grads = tape.gradient(total_loss, model.trainable_variables) optimizer.apply_gradients(zip(grads, model.trainable_variables)) # 打印日志 logging.info("{}_train_{}, {}, {}".format( epoch, batch, total_loss.numpy(), list(map(lambda x: np.sum(x.numpy()), pred_loss)))) avg_loss.update_state(total_loss) # 验证数据集验证输出计算 for batch, (images, labels) in enumerate(val_dataset): outputs = model(images) # 求损失 regularization_loss = tf.reduce_sum(model.losses) pred_loss = [] # 输出结果和标签计算损失 for output, label, loss_fn in zip(outputs, labels, loss): pred_loss.append(loss_fn(label, output)) total_loss = tf.reduce_sum(pred_loss) + regularization_loss # 打印总损失 logging.info("{}_val_{}, {}, {}".format( epoch, batch, total_loss.numpy(), list(map(lambda x: np.sum(x.numpy()), pred_loss)))) avg_val_loss.update_state(total_loss) logging.info("{}, image_2: {}, val: {}".format( epoch, avg_loss.result().numpy(), avg_val_loss.result().numpy())) # 保存模型位置 avg_loss.reset_states() avg_val_loss.reset_states() model.save_weights('checkpoints/yolov3_train_{}.tf'.format(epoch)) else: model.compile(optimizer=optimizer, loss=loss) callbacks = [ ReduceLROnPlateau(verbose=1), EarlyStopping(patience=3, verbose=1), ModelCheckpoint('checkpoints/yolov3_train_{epoch}.tf', verbose=1, save_weights_only=True), TensorBoard(log_dir='logs') ] history = model.fit(train_dataset, epochs=args.epochs, callbacks=callbacks, validation_data=val_dataset)
def main(_argv): physical_devices = tf.config.experimental.list_physical_devices('GPU') for physical_device in physical_devices: tf.config.experimental.set_memory_growth(physical_device, True) if FLAGS.tiny: model = YoloV3Tiny(FLAGS.size, training=True, classes=FLAGS.num_classes) anchors = yolo_tiny_anchors anchor_masks = yolo_tiny_anchor_masks else: model = YoloV3(FLAGS.size, training=True, classes=FLAGS.num_classes) anchors = yolo_anchors anchor_masks = yolo_anchor_masks if FLAGS.dataset: train_dataset = dataset.load_tfrecord_dataset( FLAGS.dataset, FLAGS.classes, FLAGS.size) else: train_dataset = dataset.load_fake_dataset() train_dataset = train_dataset.shuffle(buffer_size=512) train_dataset = train_dataset.batch(FLAGS.batch_size) train_dataset = train_dataset.map(lambda x, y: ( dataset.transform_images(x, FLAGS.size), dataset.transform_targets(y, anchors, anchor_masks, FLAGS.size))) train_dataset = train_dataset.prefetch( buffer_size=tf.data.experimental.AUTOTUNE) if FLAGS.val_dataset: val_dataset = dataset.load_tfrecord_dataset( FLAGS.val_dataset, FLAGS.classes, FLAGS.size) else: val_dataset = dataset.load_fake_dataset() val_dataset = val_dataset.batch(FLAGS.batch_size) val_dataset = val_dataset.map(lambda x, y: ( dataset.transform_images(x, FLAGS.size), dataset.transform_targets(y, anchors, anchor_masks, FLAGS.size))) # Configure the model for transfer learning if FLAGS.transfer == 'none': pass # Nothing to do elif FLAGS.transfer in ['darknet', 'no_output']: # Darknet transfer is a special case that works # with incompatible number of classes # reset top layers if FLAGS.tiny: model_pretrained = YoloV3Tiny( FLAGS.size, training=True, classes=FLAGS.weights_num_classes or FLAGS.num_classes) else: model_pretrained = YoloV3( FLAGS.size, training=True, classes=FLAGS.weights_num_classes or FLAGS.num_classes) model_pretrained.load_weights(FLAGS.weights) if FLAGS.transfer == 'darknet': model.get_layer('yolo_darknet').set_weights( model_pretrained.get_layer('yolo_darknet').get_weights()) freeze_all(model.get_layer('yolo_darknet')) elif FLAGS.transfer == 'no_output': for l in model.layers: if not l.name.startswith('yolo_output'): l.set_weights(model_pretrained.get_layer( l.name).get_weights()) freeze_all(l) else: # All other transfer require matching classes model.load_weights(FLAGS.weights) if FLAGS.transfer == 'fine_tune': # freeze darknet and fine tune other layers darknet = model.get_layer('yolo_darknet') freeze_all(darknet) elif FLAGS.transfer == 'frozen': # freeze everything freeze_all(model) optimizer = tf.keras.optimizers.Adam(lr=FLAGS.learning_rate) loss = [YoloLoss(anchors[mask], classes=FLAGS.num_classes) for mask in anchor_masks] if FLAGS.mode == 'eager_tf': # Eager mode is great for debugging # Non eager graph mode is recommended for real training avg_loss = tf.keras.metrics.Mean('loss', dtype=tf.float32) avg_val_loss = tf.keras.metrics.Mean('val_loss', dtype=tf.float32) for epoch in range(1, FLAGS.epochs + 1): for batch, (images, labels) in enumerate(train_dataset): with tf.GradientTape() as tape: outputs = model(images, training=True) regularization_loss = tf.reduce_sum(model.losses) pred_loss = [] for output, label, loss_fn in zip(outputs, labels, loss): pred_loss.append(loss_fn(label, output)) total_loss = tf.reduce_sum(pred_loss) + regularization_loss grads = tape.gradient(total_loss, model.trainable_variables) optimizer.apply_gradients( zip(grads, model.trainable_variables)) logging.info("{}_train_{}, {}, {}".format( epoch, batch, total_loss.numpy(), list(map(lambda x: np.sum(x.numpy()), pred_loss)))) avg_loss.update_state(total_loss) for batch, (images, labels) in enumerate(val_dataset): outputs = model(images) regularization_loss = tf.reduce_sum(model.losses) pred_loss = [] for output, label, loss_fn in zip(outputs, labels, loss): pred_loss.append(loss_fn(label, output)) total_loss = tf.reduce_sum(pred_loss) + regularization_loss logging.info("{}_val_{}, {}, {}".format( epoch, batch, total_loss.numpy(), list(map(lambda x: np.sum(x.numpy()), pred_loss)))) avg_val_loss.update_state(total_loss) logging.info("{}, train: {}, val: {}".format( epoch, avg_loss.result().numpy(), avg_val_loss.result().numpy())) avg_loss.reset_states() avg_val_loss.reset_states() model.save_weights( 'checkpoints/yolov3_train_{}.tf'.format(epoch)) else: model.compile(optimizer=optimizer, loss=loss, run_eagerly=(FLAGS.mode == 'eager_fit')) callbacks = [ ReduceLROnPlateau(verbose=1), ModelCheckpoint('checkpoints/yolov3_train_{epoch}.tf', verbose=1, save_weights_only=True), TensorBoard(log_dir='logs') ] history = model.fit(train_dataset, epochs=FLAGS.epochs, callbacks=callbacks, validation_data=val_dataset)
def main(_argv): classes_count = len(get_classes(FLAGS.classes)) if FLAGS.classes_count != -1: classes_count = FLAGS.classes_count print("Using {} classes".format(classes_count)) if FLAGS.tiny: model = YoloV3Tiny(FLAGS.size, training=True, classes=classes_count) anchors = yolo_tiny_anchors anchor_masks = yolo_tiny_anchor_masks else: model = YoloV3(FLAGS.size, training=True, classes=classes_count) anchors = yolo_anchors anchor_masks = yolo_anchor_masks train_dataset = dataset.load_fake_dataset() if FLAGS.dataset: train_dataset = dataset.load_tfrecord_dataset(FLAGS.dataset, FLAGS.classes) train_dataset = train_dataset.shuffle(buffer_size=1024) # TODO: not 1024 train_dataset = train_dataset.batch(FLAGS.batch_size) train_dataset = train_dataset.map(lambda x, y: ( dataset.transform_images(x, FLAGS.size), dataset.transform_targets(y, anchors, anchor_masks, classes_count))) train_dataset = train_dataset.prefetch( buffer_size=tf.data.experimental.AUTOTUNE) val_dataset = dataset.load_fake_dataset() if FLAGS.val_dataset: val_dataset = dataset.load_tfrecord_dataset(FLAGS.val_dataset, FLAGS.classes) val_dataset = val_dataset.batch(FLAGS.batch_size) val_dataset = val_dataset.map(lambda x, y: ( dataset.transform_images(x, FLAGS.size), dataset.transform_targets(y, anchors, anchor_masks, classes_count))) if FLAGS.transfer != 'none': if FLAGS.transfer == 'fine_tune': model.load_weights(FLAGS.weights) # freeze darknet darknet = model.get_layer('yolo_darknet') freeze_all(darknet) elif FLAGS.mode == 'frozen': model.load_weights(FLAGS.weights) # freeze everything freeze_all(model) else: # reset top layers if FLAGS.tiny: # get initial weights orig_model = YoloV3Tiny(FLAGS.size, training=True, classes=80) else: orig_model = YoloV3(FLAGS.size, training=True, classes=80) if FLAGS.transfer == 'darknet': # transfering the yolo_darknet layer # here we use or orig_model with all the classes # and load up its weights model_pretrained = orig_model model_pretrained.load_weights(FLAGS.weights) model.get_layer('yolo_darknet').set_weights( model_pretrained.get_layer('yolo_darknet').get_weights()) # freeze? freeze_all(model.get_layer('yolo_darknet')) elif FLAGS.transfer == 'no_output': # here init_model is entirely empty init_model = orig_model for l in model.layers: if l.name.startswith('yolo_output'): l.set_weights( init_model.get_layer(l.name).get_weights()) else: freeze_all(l) optimizer = tf.keras.optimizers.Adam(lr=FLAGS.learning_rate) loss = [ YoloLoss(anchors[mask], classes=classes_count) for mask in anchor_masks ] if FLAGS.mode == 'eager_tf': # Eager mode is great for debugging # Non eager graph mode is recommended for real training avg_loss = tf.keras.metrics.Mean('loss', dtype=tf.float32) avg_val_loss = tf.keras.metrics.Mean('val_loss', dtype=tf.float32) for epoch in range(1, FLAGS.epochs + 1): for batch, (images, labels) in enumerate(train_dataset): with tf.GradientTape() as tape: outputs = model(images, training=True) regularization_loss = tf.reduce_sum(model.losses) pred_loss = [] for output, label, loss_fn in zip(outputs, labels, loss): pred_loss.append(loss_fn(label, output)) total_loss = tf.reduce_sum(pred_loss) + regularization_loss grads = tape.gradient(total_loss, model.trainable_variables) optimizer.apply_gradients(zip(grads, model.trainable_variables)) logging.info("{}_train_{}, {}, {}".format( epoch, batch, total_loss.numpy(), list(map(lambda x: np.sum(x.numpy()), pred_loss)))) avg_loss.update_state(total_loss) for batch, (images, labels) in enumerate(val_dataset): outputs = model(images) regularization_loss = tf.reduce_sum(model.losses) pred_loss = [] for output, label, loss_fn in zip(outputs, labels, loss): pred_loss.append(loss_fn(label, output)) total_loss = tf.reduce_sum(pred_loss) + regularization_loss logging.info("{}_val_{}, {}, {}".format( epoch, batch, total_loss.numpy(), list(map(lambda x: np.sum(x.numpy()), pred_loss)))) avg_val_loss.update_state(total_loss) logging.info("{}, train: {}, val: {}".format( epoch, avg_loss.result().numpy(), avg_val_loss.result().numpy())) avg_loss.reset_states() avg_val_loss.reset_states() model.save_weights('checkpoints/yolov3_train_{}.tf'.format(epoch)) else: model.compile(optimizer=optimizer, loss=loss, run_eagerly=(FLAGS.mode == 'eager_fit')) callbacks = [ ReduceLROnPlateau(verbose=1), EarlyStopping(patience=3, verbose=1), ModelCheckpoint('checkpoints/yolov3_train_{epoch}.tf', verbose=1, save_weights_only=True), TensorBoard(log_dir='logs') ] history = model.fit(train_dataset, epochs=FLAGS.epochs, callbacks=callbacks, validation_data=val_dataset)
def main(_argv): set_one_gpu() if FLAGS.tiny: model = YoloV3Tiny(FLAGS.size, training=True, classes=FLAGS.num_classes) anchors = yolo_tiny_anchors anchor_masks = yolo_tiny_anchor_masks else: model = YoloV3(FLAGS.size, training=True, classes=FLAGS.num_classes) anchors = yolo_anchors anchor_masks = yolo_anchor_masks # train_dataset = dataset.load_fake_dataset() dataset_name = 'data/' + FLAGS.dataset + '.train.record' val_dataset_name = 'data/' + FLAGS.dataset + '.val.record' train_dataset = dataset.load_tfrecord_dataset(dataset_name, FLAGS.classes) train_dataset = train_dataset.shuffle(buffer_size=1024) # TODO: not 1024 train_dataset = train_dataset.batch(FLAGS.batch_size) train_dataset = train_dataset.map( lambda x, y: (dataset.transform_images(x, FLAGS.size), dataset.transform_targets(y, anchors, anchor_masks, 80))) train_dataset = train_dataset.prefetch( buffer_size=tf.data.experimental.AUTOTUNE) tf_name = FLAGS.name if not tf_name: tf_name = 'train' + FLAGS.gpu best_tf_name = "checkpoints/%s_best.tf" % tf_name last_tf_name = "checkpoints/%s_last.tf" % tf_name # val_dataset = dataset.load_fake_dataset() val_dataset = dataset.load_tfrecord_dataset(val_dataset_name, FLAGS.classes) val_dataset = val_dataset.batch(FLAGS.batch_size) val_dataset = val_dataset.map( lambda x, y: (dataset.transform_images(x, FLAGS.size), dataset.transform_targets(y, anchors, anchor_masks, 80))) if FLAGS.transfer != 'none': model.load_weights(FLAGS.weights) if FLAGS.transfer == 'fine_tune': # freeze darknet darknet = model.get_layer('yolo_darknet') freeze_all(darknet) elif FLAGS.transfer == 'frozen': # freeze everything freeze_all(model) else: # reset top layers if FLAGS.tiny: # get initial weights init_model = YoloV3Tiny(FLAGS.size, training=True, classes=FLAGS.num_classes) else: init_model = YoloV3(FLAGS.size, training=True, classes=FLAGS.num_classes) if FLAGS.transfer == 'darknet': for l in model.layers: if l.name != 'yolo_darknet' and l.name.startswith('yolo_'): l.set_weights( init_model.get_layer(l.name).get_weights()) else: freeze_all(l) elif FLAGS.transfer == 'no_output': for l in model.layers: if l.name.startswith('yolo_output'): l.set_weights( init_model.get_layer(l.name).get_weights()) else: freeze_all(l) optimizer = tf.keras.optimizers.Adam(lr=FLAGS.learning_rate) loss = [ YoloLoss(anchors[mask], classes=FLAGS.num_classes) for mask in anchor_masks ] best_val_loss = 0 history = None if FLAGS.mode == 'eager_tf': # Eager mode is great for debugging # Non eager graph mode is recommended for real training avg_loss = tf.keras.metrics.Mean('loss', dtype=tf.float32) avg_val_loss = tf.keras.metrics.Mean('val_loss', dtype=tf.float32) for epoch in range(1, FLAGS.epochs + 1): for batch, (images, labels) in enumerate(train_dataset): with tf.GradientTape() as tape: outputs = model(images, training=True) regularization_loss = tf.reduce_sum(model.losses) pred_loss = [] for output, label, loss_fn in zip(outputs, labels, loss): pred_loss.append(loss_fn(label, output)) total_loss = tf.reduce_sum(pred_loss) + regularization_loss grads = tape.gradient(total_loss, model.trainable_variables) optimizer.apply_gradients(zip(grads, model.trainable_variables)) # logging.info("{}_train_{}, {}, {}".format( # epoch, batch, total_loss.numpy(), # list(map(lambda x: np.sum(x.numpy()), pred_loss)))) avg_loss.update_state(total_loss) for batch, (images, labels) in enumerate(val_dataset): outputs = model(images) regularization_loss = tf.reduce_sum(model.losses) pred_loss = [] for output, label, loss_fn in zip(outputs, labels, loss): pred_loss.append(loss_fn(label, output)) total_loss = tf.reduce_sum(pred_loss) + regularization_loss # logging.info("{}_val_{}, {}, {}".format( # epoch, batch, total_loss.numpy(), # list(map(lambda x: np.sum(x.numpy()), pred_loss)))) avg_val_loss.update_state(total_loss) val_lost = avg_val_loss.result().numpy() logging.info("{}, train: {}, val: {}".format( epoch, avg_loss.result().numpy(), val_lost)) avg_loss.reset_states() avg_val_loss.reset_states() model.save_weights(last_tf_name) if best_val_loss == 0 or best_val_loss > val_lost: best_val_loss = val_lost logging.info("saving best val loss: %s" % best_tf_name) model.save_weights(best_tf_name) else: model.compile(optimizer=optimizer, loss=loss, run_eagerly=(FLAGS.mode == 'eager_fit')) callbacks = [ ReduceLROnPlateau(verbose=1), EarlyStopping(patience=3, verbose=1), ModelCheckpoint(best_tf_name, verbose=1, save_weights_only=True), TensorBoard(log_dir='logs') ] history = model.fit(train_dataset, epochs=FLAGS.epochs, callbacks=callbacks, validation_data=val_dataset) if history is not None: print(history.history['val_loss']) best_val_loss = min(history.history['val_loss']) model.save_weights(best_tf_name) print("Best weights are saved as %s" % best_tf_name) tiny = 'tiny_' if FLAGS.tiny else '' out_name = "%s_d%s_%sm%s_bs%d_s%s_e%d_val%d" % \ (tf_name, FLAGS.dataset, tiny, FLAGS.transfer, FLAGS.batch_size, FLAGS.size, FLAGS.epochs, best_val_loss) mfn = "data/model/%s/" % out_name final_tf_name = "%s.tf" % out_name copy_tf("%s_best.tf" % tf_name, final_tf_name) print("Final checkpoint file saved as: %s" % final_tf_name) model.load_weights(best_tf_name) tf.saved_model.save(model, mfn) print("Model file saved to: %s" % mfn)
def main(): train_path = '/Users/justinbutler/Desktop/school/Calgary/ML_Work/Datasets/Shapes/tfrecord_single/coco_train.record-00000-of-00001' valid_path = '/Users/justinbutler/Desktop/school/Calgary/ML_Work/Datasets/Shapes/tfrecord_single/coco_val.record-00000-of-00001' weights_path = '/Users/justinbutler/Desktop/school/Calgary/ML_Work/yolov3-tf2/checkpoints/yolov3.tf' # Path to text? file containing all classes, 1 per line classes = '/Users/justinbutler/Desktop/school/Calgary/ML_Work/yolov3-tf2/shapes/shapes.names' # Usually fit # mode = 'fit' # Can be 'fit', 'eager_fit', 'eager_tf', 'valid' mode = 'fit' ''' 'fit: model.fit, ' 'eager_fit: model.fit(run_eagerly=True), ' 'eager_tf: custom GradientTape' ''' # Usually darknet transfer = 'none' ''' 'none: Training from scratch, ' 'darknet: Transfer darknet, ' 'no_output: Transfer all but output, ' 'frozen: Transfer and freeze all, ' 'fine_tune: Transfer all and freeze darknet only'), 'pre': Use a pre-trained model for validation ''' image_size = 416 num_epochs = 1 batch_size = 8 learning_rate = 1e-3 num_classes = 4 # num class for `weights` file if different, useful in transfer learning with different number of classes weight_num_classes = 80 iou_threshold = 0.5 # saved_weights_path = '/Users/justinbutler/Desktop/school/Calgary/ML_Work/yolov3-tf2/weights/' saved_weights_path = '/home/justin/ml_models/yolov3-tf2/weights/shapes_{}.tf'.format( num_epochs) anchors = yolo_anchors anchor_masks = yolo_anchor_masks # Training dataset #dataset_train = tf.data.TFRecordDataset(train_path) #dataset_val = tf.data.TFRecordDataset(valid_path) dataset_train = load_tfrecord_dataset(train_path, classes, image_size) dataset_train = dataset_train.shuffle(buffer_size=512) dataset_train = dataset_train.batch(batch_size) #dataset_train = dataset_train.map(lambda x, y: ( # transform_images(x, image_size), # transform_targets(y, anchors, anchor_masks, image_size))) #dataset_train = dataset_train.prefetch( # buffer_size=tf.data.experimental.AUTOTUNE) dataset_val = load_tfrecord_dataset(valid_path, classes, image_size) dataset_val = dataset_val.shuffle(buffer_size=512) dataset_val = dataset_val.batch(batch_size) #dataset_val = dataset_val.map(lambda x, y: ( # transform_images(x, image_size), # transform_targets(y, anchors, anchor_masks, image_size))) # Create model in training mode yolo = models.YoloV3(image_size, training=True, classes=num_classes) model_pretrained = YoloV3(image_size, training=True, classes=weight_num_classes or num_classes) model_pretrained.load_weights(weights_path) # Which weights to start with? print('Loading Weights...') #yolo.load_weights(weights_path) yolo.get_layer('yolo_darknet').set_weights( model_pretrained.get_layer('yolo_darknet').get_weights()) freeze_all(yolo.get_layer('yolo_darknet')) optimizer = tf.keras.optimizers.Adam(lr=learning_rate) loss = [ YoloLoss(anchors[mask], classes=num_classes) for mask in anchor_masks ] # Passing loss as a list might sometimes fail? dict might be better? yolo.compile(optimizer=optimizer, loss=loss, run_eagerly=(mode == 'eager_fit')) callbacks = [ ReduceLROnPlateau(verbose=1), EarlyStopping(patience=3, verbose=1), ModelCheckpoint('checkpoints/yolov3_train_{epoch}.tf', verbose=1, save_weights_only=True), TensorBoard(log_dir='logs') ] history = yolo.fit(dataset_train, epochs=num_epochs, callbacks=callbacks, validation_data=dataset_val) yolo.save_weights(saved_weights_path) # Detect/ROC model = YoloV3(image_size, training=False, classes=num_classes) model.load_weights(saved_weights_path).expect_partial() batch_size = 1 val_dataset = load_tfrecord_dataset(valid_path, classes, image_size) val_dataset = val_dataset.batch(batch_size) val_dataset = val_dataset.map( lambda x, y: (transform_images(x, image_size), transform_targets(y, anchors, anchor_masks, image_size))) images = [] for img, labs in val_dataset: img = np.squeeze(img) images.append(img) predictions = [] evaluator = Evaluator(iou_thresh=iou_threshold) # labels - (N, grid, grid, anchors, [x, y, w, h, obj, class]) boxes, scores, classes, num_detections = model.predict(val_dataset) # boxes -> (num_imgs, num_detections (200), box coords (4)) # scores -> (num_imgs, num_detections) # classes -> (num_imgs, num_detections) # num_detections -> num_imgs # Aim for labels shape (per batch): [num_imgs, 3x[num_boxes x [x1,y1,x2,y2,score,class]] # full_labels = [label for _, label in val_dataset] # Shape : [Num images, 3 scales, grid, grid, anchor, 6 ] filtered_labels = [] for _, label in val_dataset: img_labels = [] # Label has shape [3 scales x[1, grid, grid, 3, 6]] for scale in label: # Shape [1, grid, grid, 3, 6] scale = np.asarray(scale) grid = scale.shape[1] scale2 = np.reshape(scale, (3, grid * grid, 6)) # Shape: [3, grix*grid, 6] for anchor in scale2: filtered_anchors = [] for box in anchor: if box[4] > 0: filtered_anchors.append(np.asarray(box)) img_labels.append(filtered_anchors) img_labels = np.asarray(img_labels) filtered_labels.append(img_labels) print(len(filtered_labels)) print(len(filtered_labels[0])) print(len(filtered_labels[0][2])) # i is the num_images index # predictions = [np.hstack([boxes[i][x], scores[i][x], classes[i][x]]) for i in range(len(num_detections)) for x in range(len(scores[i])) if scores[i][x] > 0] for img in range(len(num_detections)): row = [] for sc in range(len(scores[img])): if scores[img][sc] > 0: row.append( np.hstack([ boxes[img][sc] * image_size, scores[img][sc], classes[img][sc] ])) predictions.append(np.asarray(row)) predictions = np.asarray( predictions) # numpy array of shape [num_imgs x num_preds x 6] if len(predictions) == 0: # No predictions made print('No predictions made - exiting.') exit() # Predictions shape: [num_imgs x num_preds x[box coords(4), conf, classes]] # Box coords should be in format x1 y1 x2 y2 # Labels shape: [num_imgs, 3x[num_boxes x [x1,y1,x2,y2,score,class]] evaluator(predictions, filtered_labels, images) # Check gts box coords '''
def main(_argv): # Setting up the accelerator strategy = setup_accelerator(FLAGS.accelerator) if 'TPU' in FLAGS.accelerator: with strategy.scope(): if FLAGS.tiny: model = YoloV3Tiny(FLAGS.size, training=True, classes=FLAGS.num_classes, mixed_precision=FLAGS.mixed_precision) anchors = adjust_yolo_anchors(yolo_tiny_anchors, FLAGS.size) anchor_masks = yolo_tiny_anchor_masks else: model = YoloV3(FLAGS.size, training=True, classes=FLAGS.num_classes, mixed_precision=FLAGS.mixed_precision) anchors = adjust_yolo_anchors(yolo_anchors, FLAGS.size) anchor_masks = yolo_anchor_masks else: if FLAGS.tiny: model = YoloV3Tiny(FLAGS.size, training=True, classes=FLAGS.num_classes, mixed_precision=FLAGS.mixed_precision) anchors = adjust_yolo_anchors(yolo_tiny_anchors, FLAGS.size) anchor_masks = yolo_tiny_anchor_masks else: model = YoloV3(FLAGS.size, training=True, classes=FLAGS.num_classes, mixed_precision=FLAGS.mixed_precision) anchors = adjust_yolo_anchors(yolo_anchors, FLAGS.size) anchor_masks = yolo_anchor_masks if FLAGS.dataset: train_dataset = dataset.load_tfrecord_dataset( FLAGS.dataset, FLAGS.classes, FLAGS.size) train_dataset = train_dataset.shuffle(buffer_size=512).batch(FLAGS.batch_size) train_dataset = train_dataset.map(lambda x, y: ( dataset.transform_images(x, FLAGS.size), dataset.transform_targets(y, anchors, anchor_masks, FLAGS.size))) train_dataset = train_dataset.prefetch( buffer_size=tf.data.experimental.AUTOTUNE) if FLAGS.val_dataset: val_dataset = dataset.load_tfrecord_dataset( FLAGS.val_dataset, FLAGS.classes, FLAGS.size) val_dataset = val_dataset.batch(FLAGS.batch_size) val_dataset = val_dataset.map(lambda x, y: ( dataset.transform_images(x, FLAGS.size), dataset.transform_targets(y, anchors, anchor_masks, FLAGS.size))) # Configure the model for transfer learning if FLAGS.transfer == 'none': pass # Nothing to do elif FLAGS.transfer in ['no_darknet', 'no_output']: # Darknet transfer is a special case that works # with incompatible number of classes # reset top layers if FLAGS.tiny: model_pretrained = YoloV3Tiny( FLAGS.size, training=True, classes=FLAGS.pretrained_weights_num_classes or FLAGS.num_classes, mixed_precision=FLAGS.pretrained_mixed_precision ) else: model_pretrained = YoloV3( FLAGS.size, training=True, classes=FLAGS.pretrained_weights_num_classes or FLAGS.num_classes, mixed_precision=FLAGS.pretrained_mixed_precision ) model_pretrained.load_weights(FLAGS.weights) if FLAGS.transfer == 'no_darknet': model.get_layer('yolo_darknet').set_weights( model_pretrained.get_layer('yolo_darknet').get_weights()) freeze_all(model.get_layer('yolo_darknet')) elif FLAGS.transfer == 'no_output': for l in model.layers: if not l.name.startswith('yolo_output'): l.set_weights(model_pretrained.get_layer(l.name).get_weights()) freeze_all(l) else: # All other transfer require matching classes model.load_weights(FLAGS.weights) if FLAGS.transfer == 'fine_tune': # freeze darknet and fine tune other layers darknet = model.get_layer('yolo_darknet') freeze_all(darknet) elif FLAGS.transfer == 'all': # fine_tune the whole model. for i, l in enumerate(model.layers): print(i, "layer: ", l.name) freeze_all(l, frozen=False) elif FLAGS.transfer == 'frozen': # freeze everything freeze_all(model) optimizer = tf.keras.optimizers.Adam(lr=FLAGS.learning_rate) loss = [YoloLoss(anchors[mask], classes=FLAGS.num_classes, ignore_thresh=FLAGS.ignore_thresh, dtype=FLAGS.mixed_precision) for mask in anchor_masks] logging.info("\nloss: {}".format(loss)) for l in loss: print(l) logging.info("global policy: {}".format(tf_mixed_precision.global_policy())) logging.info("global policy loss: {}\n".format(tf_mixed_precision.global_policy().loss_scale)) # Print model summary and plot it to .png print_all_layers(model, p_details=FLAGS.pdetails) if FLAGS.tiny: tf.keras.utils.plot_model(model, to_file='yolov3_tiny.png', show_shapes=True, show_layer_names=False) else: tf.keras.utils.plot_model(model, to_file='yolov3.png', show_shapes=True, show_layer_names=False) if FLAGS.mode == 'eager_tf': # Eager mode is great for debugging # Non eager graph mode is recommended for real training #optimizer = tf_mixed_precision.LossScaleOptimizer(optimizer, loss_scale='dynamic') avg_loss = tf.keras.metrics.Mean('loss', dtype=tf.float32) avg_val_loss = tf.keras.metrics.Mean('val_loss', dtype=tf.float32) y = [] y_val = [] for epoch in range(1, FLAGS.epochs + 1): for batch, (images, labels) in enumerate(train_dataset): with tf.GradientTape() as tape: outputs = model(images, training=True) regularization_loss = tf.reduce_sum(model.losses) pred_loss = [] for output, label, loss_fn in zip(outputs, labels, loss): pred_loss.append(loss_fn(label, output)) total_loss = tf.reduce_sum(pred_loss) + regularization_loss grads = tape.gradient(total_loss, model.trainable_variables) optimizer.apply_gradients( zip(grads, model.trainable_variables)) logging.info("{}_train_{}, {}, {}".format( epoch, batch, total_loss.numpy(), list(map(lambda x: np.sum(x.numpy()), pred_loss)))) avg_loss.update_state(total_loss) for batch, (images, labels) in enumerate(val_dataset): outputs = model(images) regularization_loss = tf.reduce_sum(model.losses) pred_loss = [] for output, label, loss_fn in zip(outputs, labels, loss): pred_loss.append(loss_fn(label, output)) total_loss = tf.reduce_sum(pred_loss) + regularization_loss logging.info("{}_val_{}, {}, {}".format( epoch, batch, total_loss.numpy(), list(map(lambda x: np.sum(x.numpy()), pred_loss)))) avg_val_loss.update_state(total_loss) logging.info(GREEN+"{}, train: {}, val: {}\33".format( epoch, avg_loss.result().numpy(), avg_val_loss.result().numpy())+DEFAULT) y.append(avg_loss.result().numpy()) y_val.append(avg_val_loss.result().numpy()) avg_loss.reset_states() avg_val_loss.reset_states() if epoch%4 == 0: save_epoch = epoch model.save_weights( 'checkpoints/yolov3_train_{}.tf'.format(save_epoch)) logging.info(GREEN+'checkpoints/yolov3_train_{}.tf weights saved'.format(save_epoch)+DEFAULT) else: metrics = None if FLAGS.metrics: thresholds = [x for x in np.linspace(0,1,11)] precision = [] #recall = [] for mask in anchor_masks: precision.append([experimentalYoloPrecision(anchors=anchors[mask], num_classes=FLAGS.num_classes, thresholds=thresholds, name='precision')]) #recall.append([experimentalYoloRecall(anchors=anchors[mask], num_classes=FLAGS.num_classes, thresholds=thresholds, name='recall')]) #recall = [experimentalYoloRecall(anchors=anchors[mask], num_classes=FLAGS.num_classes, thresholds=thresholds, name='recall') for mask in anchor_masks] metrics = [] for p in zip(precision): metrics.append(p) if 'TPU' in FLAGS.accelerator: with strategy.scope(): model.compile(optimizer=optimizer, loss=loss, metrics = metrics, run_eagerly=(FLAGS.mode=='eager_fit')) else: model.compile(optimizer=optimizer, loss=loss, metrics = metrics, run_eagerly=(FLAGS.mode=='eager_fit')) callbacks = [ ModelCheckpoint(FLAGS.output_weights, verbose=1, save_best_only=True, save_weights_only=True), TensorBoard(log_dir=FLAGS.log_output_fit+'-'+datetime.datetime.now().strftime("%Y%m%d-%H%M%S"), histogram_freq=1) ] print("\n--- START FITTING ---\n") training_history = model.fit(train_dataset, epochs=FLAGS.epochs, callbacks=callbacks, validation_data=val_dataset, workers=2) logging.info(" > Average loss: \t{}".format(np.average(training_history.history['loss']))) logging.info(" > Average val loss: \t{}".format(np.average(training_history.history['val_loss']))) if FLAGS.metrics: for i, k in enumerate(list(training_history.history.keys())): if 'precision' in k: logging.info(" > {0}, mAP={1:.4}%".format(k, np.average(training_history.history[k])*100)) if 'recall' in k: logging.info(" > {0}, mAR={1:.4}%".format(k, np.average(training_history.history[k])*100))
def setup_model(): if FLAGS.tiny: model = YoloV3Tiny(FLAGS.size, training=True, classes=FLAGS.num_classes) anchors = yolo_tiny_anchors anchor_masks = yolo_tiny_anchor_masks else: model = YoloV3(FLAGS.size, training=True, classes=FLAGS.num_classes) anchors = yolo_anchors anchor_masks = yolo_anchor_masks # Configure the model for transfer learning if FLAGS.transfer == 'none': pass # Nothing to do elif FLAGS.transfer in ['darknet', 'no_output']: # Darknet transfer is a special case that works # with incompatible number of classes # reset top layers if FLAGS.tiny: model_pretrained = YoloV3Tiny(FLAGS.size, training=True, classes=FLAGS.weights_num_classes or FLAGS.num_classes) else: model_pretrained = YoloV3(FLAGS.size, training=True, classes=FLAGS.weights_num_classes or FLAGS.num_classes) model_pretrained.load_weights(FLAGS.weights) if FLAGS.transfer == 'darknet': model.get_layer('yolo_darknet').set_weights( model_pretrained.get_layer('yolo_darknet').get_weights()) freeze_all(model.get_layer('yolo_darknet')) elif FLAGS.transfer == 'no_output': for l in model.layers: if not l.name.startswith('yolo_output'): l.set_weights( model_pretrained.get_layer(l.name).get_weights()) freeze_all(l) else: # All other transfer require matching classes model.load_weights(FLAGS.weights) if FLAGS.transfer == 'fine_tune': # freeze darknet and fine tune other layers darknet = model.get_layer('yolo_darknet') freeze_all(darknet) elif FLAGS.transfer == 'frozen': # freeze everything freeze_all(model) optimizer = tf.keras.optimizers.Adam(lr=FLAGS.learning_rate) loss = [ YoloLoss(anchors[mask], classes=FLAGS.num_classes) for mask in anchor_masks ] model.compile(optimizer=optimizer, loss=loss, run_eagerly=(FLAGS.mode == 'eager_fit')) return model, optimizer, loss, anchors, anchor_masks
def main(_argv): physical_devices = tf.config.experimental.list_physical_devices('GPU') for physical_device in physical_devices: tf.config.experimental.set_memory_growth(physical_device, True) check_weighs_exist(tiny=FLAGS.tiny) if FLAGS.tiny: model = YoloV3Tiny( FLAGS.size, training=True, classes=FLAGS.num_classes ) model.summary() plot_model(model, to_file='yoloV3Tiny-model-plot.png', show_shapes=True, show_layer_names=True) anchors = yolo_tiny_anchors anchor_masks = yolo_tiny_anchor_masks else: model = YoloV3( FLAGS.size, training=True, classes=FLAGS.num_classes ) model.summary() plot_model(model, to_file='yoloV3-model-plot.png', show_shapes=True, show_layer_names=True) anchors = yolo_anchors anchor_masks = yolo_anchor_masks # Load the dataset train_dataset = dataset.load_fake_dataset() if FLAGS.dataset: train_dataset = dataset.load_tfrecord_dataset( file_pattern=FLAGS.dataset, class_file=FLAGS.classes, size=FLAGS.size ) # Shuffle the dataset train_dataset = train_dataset.shuffle(buffer_size=FLAGS.buffer_size, reshuffle_each_iteration=True) train_dataset_length = [i for i, _ in enumerate(train_dataset)][-1] + 1 print(f"Dataset for training consists of {train_dataset_length} images.") train_dataset = train_dataset.batch(FLAGS.batch_size) train_dataset = train_dataset.map(lambda x, y: (dataset.transform_images(x, FLAGS.size), dataset.transform_targets(y, anchors, anchor_masks, FLAGS.size))).repeat() train_dataset = train_dataset.prefetch(buffer_size=tf.data.experimental.AUTOTUNE) val_dataset = dataset.load_fake_dataset() if FLAGS.val_dataset: val_dataset = dataset.load_tfrecord_dataset( FLAGS.val_dataset, FLAGS.classes, FLAGS.size) val_dataset_length = [i for i, _ in enumerate(val_dataset)][-1] + 1 print(f"Dataset for validation consists of {val_dataset_length} images.") val_dataset = val_dataset.batch(FLAGS.batch_size) val_dataset = val_dataset.map(lambda x, y: (dataset.transform_images(x, FLAGS.size), dataset.transform_targets(y, anchors, anchor_masks, FLAGS.size))).repeat() # Configure the model for transfer learning if FLAGS.transfer == 'none': pass # Nothing to do elif FLAGS.transfer in ['darknet', 'no_output']: # Darknet transfer is a special case that works # with incompatible number of classes # reset top layers if FLAGS.tiny: model_pretrained = YoloV3Tiny( size=FLAGS.size, training=True, classes=FLAGS.weights_num_classes or FLAGS.num_classes) model_pretrained.load_weights(FLAGS.weights_tf_format_tiny) else: model_pretrained = YoloV3( size=FLAGS.size, training=True, classes=FLAGS.weights_num_classes or FLAGS.num_classes) model_pretrained.load_weights(FLAGS.weights_tf_format) if FLAGS.transfer == 'darknet': # Set yolo darknet layer weights to the loaded pretrained model weights model.get_layer('yolo_darknet').set_weights( model_pretrained.get_layer('yolo_darknet').get_weights()) # Freeze these layers freeze_all(model.get_layer('yolo_darknet')) elif FLAGS.transfer == 'no_output': for i in model.layers: if not i.name.startswith('yolo_output'): i.set_weights(model_pretrained.get_layer( i.name).get_weights()) freeze_all(i) else: # All other transfer require matching classes if FLAGS.tiny: model.load_weights(FLAGS.weights_tf_format_tiny) else: model.load_weights(FLAGS.weights_tf_format) if FLAGS.transfer == 'fine_tune': # freeze darknet and fine tune other layers darknet = model.get_layer('yolo_darknet') freeze_all(darknet) elif FLAGS.transfer == 'frozen': # freeze everything freeze_all(model) # Use the Adam optimizer with the specified learning rate optimizer = tf.keras.optimizers.Adam(lr=FLAGS.learning_rate) # YoloLoss function loss = [YoloLoss(anchors[mask], classes=FLAGS.num_classes) for mask in anchor_masks] if FLAGS.mode == 'eager_tf': print(f"Mode is: {FLAGS.mode}") # Eager mode is great for debugging # Non eager graph mode is recommended for real training avg_loss = tf.keras.metrics.Mean('loss', dtype=tf.float32) avg_val_loss = tf.keras.metrics.Mean('val_loss', dtype=tf.float32) for epoch in range(1, FLAGS.epochs + 1): for batch, (images, labels) in enumerate(train_dataset): with tf.GradientTape() as tape: outputs = model(images, training=True) regularization_loss = tf.reduce_sum(model.losses) pred_loss = [] for output, label, loss_fn in zip(outputs, labels, loss): pred_loss.append(loss_fn(label, output)) total_loss = tf.reduce_sum(pred_loss) + regularization_loss grads = tape.gradient(total_loss, model.trainable_variables) optimizer.apply_gradients(zip(grads, model.trainable_variables)) logging.info(f"epoch_{epoch}_train_batch_{batch}," f"{total_loss.numpy()}," f"{list(map(lambda x: np.sum(x.numpy()), pred_loss))}") avg_loss.update_state(total_loss) for batch, (images, labels) in enumerate(val_dataset): outputs = model(images) regularization_loss = tf.reduce_sum(model.losses) pred_loss = [] for output, label, loss_fn in zip(outputs, labels, loss): pred_loss.append(loss_fn(label, output)) total_loss = tf.reduce_sum(pred_loss) + regularization_loss logging.info("{}_val_{}, {}, {}".format( epoch, batch, total_loss.numpy(), list(map(lambda x: np.sum(x.numpy()), pred_loss))) ) avg_val_loss.update_state(total_loss) logging.info("{}, train: {}, val: {}".format( epoch, avg_loss.result().numpy(), avg_val_loss.result().numpy())) avg_loss.reset_states() avg_val_loss.reset_states() model.save_weights(f'checkpoints/{data_set}_tiny_{FLAGS.tiny}_im_size_{FLAGS.size}.tf') else: print(f"Compiling the model") model.compile( optimizer=optimizer, loss=loss, run_eagerly=(FLAGS.mode == 'eager_fit'), metrics=['accuracy']) callbacks = [ EarlyStopping(monitor='val_loss', patience=125, verbose=1), ReduceLROnPlateau(monitor='val_loss', verbose=1, factor=0.90, min_lr=0, patience=20, mode="auto"), ModelCheckpoint( str(f'checkpoints/{data_set}_tiny_{FLAGS.tiny}_im_size_{FLAGS.size}.tf'), verbose=1, save_weights_only=True, save_best_only=True, mode="auto", ), TensorBoard(log_dir='logs'), CSVLogger(f'checkpoints/logs/{data_set}_tiny_{FLAGS.tiny}_im_size_{FLAGS.size}', separator=',') ] history = model.fit(train_dataset, epochs=FLAGS.epochs, verbose=2, callbacks=callbacks, validation_data=val_dataset, steps_per_epoch=np.ceil(train_dataset_length / FLAGS.batch_size), validation_steps=np.ceil(val_dataset_length / FLAGS.batch_size))
def main(_argv): physical_devices = tf.config.experimental.list_physical_devices('GPU') for physical_device in physical_devices: tf.config.experimental.set_memory_growth(physical_device, True) if FLAGS.tiny: model = YoloV3Tiny(FLAGS.size, training=True, classes=FLAGS.num_classes) anchors = yolo_tiny_anchors anchor_masks = yolo_tiny_anchor_masks else: model = YoloV3(FLAGS.size, training=True, classes=FLAGS.num_classes) anchors = yolo_anchors anchor_masks = yolo_anchor_masks post_process_outputs = post_process_block(model.outputs, classes=FLAGS.num_classes) post_process_model = Model(model.inputs, post_process_outputs) train_dataset = dataset.load_fake_dataset() if FLAGS.dataset: train_dataset = dataset.load_tfrecord_dataset(FLAGS.dataset, FLAGS.classes, FLAGS.size) train_dataset = train_dataset.shuffle(buffer_size=512) train_dataset = train_dataset.batch(FLAGS.batch_size) train_dataset = train_dataset.map( lambda x, y: (dataset.transform_images(x, FLAGS.size), y)) # dataset.transform_targets(y, anchors, anchor_masks, FLAGS.size))) train_dataset = train_dataset.prefetch( buffer_size=tf.data.experimental.AUTOTUNE) val_dataset = dataset.load_fake_dataset() if FLAGS.val_dataset: val_dataset = dataset.load_tfrecord_dataset(FLAGS.val_dataset, FLAGS.classes, FLAGS.size) val_dataset = val_dataset.batch(FLAGS.batch_size) val_dataset = val_dataset.map(lambda x, y: (dataset.transform_images(x, FLAGS.size), y)) # dataset.transform_targets(y, anchors, anchor_masks, FLAGS.size))) # Configure the model for transfer learning if FLAGS.transfer == 'none': pass # Nothing to do elif FLAGS.transfer in ['darknet', 'no_output']: # Darknet transfer is a special case that works # with incompatible number of classes # reset top layers if FLAGS.tiny: model_pretrained = YoloV3Tiny(FLAGS.size, training=True, classes=FLAGS.weights_num_classes or FLAGS.num_classes) else: model_pretrained = YoloV3(FLAGS.size, training=True, classes=FLAGS.weights_num_classes or FLAGS.num_classes) model_pretrained.load_weights(FLAGS.weights) if FLAGS.transfer == 'darknet': model.get_layer('yolo_darknet').set_weights( model_pretrained.get_layer('yolo_darknet').get_weights()) freeze_all(model.get_layer('yolo_darknet')) elif FLAGS.transfer == 'no_output': for l in model.layers: if not l.name.startswith('yolo_output'): l.set_weights( model_pretrained.get_layer(l.name).get_weights()) freeze_all(l) else: # All other transfer require matching classes model.load_weights(FLAGS.weights) if FLAGS.transfer == 'fine_tune': # freeze darknet and fine tune other layers darknet = model.get_layer('yolo_darknet') freeze_all(darknet) elif FLAGS.transfer == 'frozen': # freeze everything freeze_all(model) optimizer = tf.keras.optimizers.Adam(lr=FLAGS.learning_rate) loss = [ YoloLoss(anchors[mask], classes=FLAGS.num_classes) for mask in anchor_masks ] # (batch_size, grid, grid, anchors, (x, y, w, h, obj, ...cls)) # model.outputs shape: [[N, 13, 13, 3, 85], [N, 26, 26, 3, 85], [N, 52, 52, 3, 85]] # labels shape: ([N, 13, 13, 3, 6], [N, 26, 26, 3, 6], [N, 52, 52, 3, 6]) if FLAGS.mode == 'eager_tf': # Eager mode is great for debugging # Non eager graph mode is recommended for real training avg_loss = tf.keras.metrics.Mean('loss', dtype=tf.float32) avg_val_loss = tf.keras.metrics.Mean('val_loss', dtype=tf.float32) for epoch in range(1, FLAGS.epochs + 1): for batch, (images, labels) in enumerate(train_dataset): with tf.GradientTape() as tape: outputs = model(images, training=True) regularization_loss = tf.reduce_sum(model.losses) pred_loss = [] transf_labels = dataset.transform_targets( labels, anchors, anchor_masks, FLAGS.size) for output, label, loss_fn in zip(outputs, transf_labels, loss): pred_loss.append(loss_fn(label, output)) total_loss = tf.reduce_sum(pred_loss, axis=None) + regularization_loss grads = tape.gradient(total_loss, model.trainable_variables) optimizer.apply_gradients(zip(grads, model.trainable_variables)) log_batch(logging, epoch, batch, total_loss, pred_loss) avg_loss.update_state(total_loss) if batch >= 100: break true_pos_total = np.zeros(FLAGS.num_classes) false_pos_total = np.zeros(FLAGS.num_classes) n_pos_total = np.zeros(FLAGS.num_classes) for batch, (images, labels) in enumerate(val_dataset): # get losses outputs = model(images) regularization_loss = tf.reduce_sum(model.losses) pred_loss = [] transf_labels = dataset.transform_targets( labels, anchors, anchor_masks, FLAGS.size) for output, label, loss_fn in zip(outputs, transf_labels, loss): pred_loss.append(loss_fn(label, output)) total_loss = tf.reduce_sum(pred_loss) + regularization_loss log_batch(logging, epoch, batch, total_loss, pred_loss) avg_val_loss.update_state(total_loss) # get true positives, false positives, and positive labels preds = post_process_model(images) true_pos, false_pos, n_pos = batch_true_false_positives( preds.numpy(), labels.numpy(), FLAGS.num_classes) true_pos_total += true_pos false_pos_total += false_pos n_pos_total += n_pos if batch >= 20: break # precision-recall by class precision, recall = batch_precision_recall(true_pos_total, false_pos_total, n_pos_total) for c in range(FLAGS.num_classes): print('Class {} - Prec: {}, Rec: {}'.format( c, precision[c], recall[c])) # total precision-recall print('Total - Prec: {}, Rec: {}'.format( calc_precision(np.sum(true_pos_total), np.sum(false_pos_total)), calc_recall(np.sum(true_pos_total), np.sum(n_pos_total)))) import pdb pdb.set_trace() # log losses logging.info("{}, train: {}, val: {}".format( epoch, avg_loss.result().numpy(), avg_val_loss.result().numpy())) # reset loop and save weights avg_loss.reset_states() avg_val_loss.reset_states() model.save_weights( os.path.join(FLAGS.checkpoint_dir, 'yolov3_train_{}.tf'\ .format(epoch))) else: model.compile(optimizer=optimizer, loss=loss, run_eagerly=(FLAGS.mode == 'eager_fit')) callbacks = [ ReduceLROnPlateau(verbose=1), EarlyStopping(patience=3, verbose=1), ModelCheckpoint(os.path.join(FLAGS.checkpoint_dir, 'yolov3_train_{epoch}.tf'), verbose=1, save_weights_only=True), TensorBoard(log_dir=FLAGS.log_dir) ] history = model.fit(train_dataset, epochs=FLAGS.epochs, callbacks=callbacks, validation_data=val_dataset)
def main(args): image_size = 416 # 416 num_epochs = args.epochs batch_size = args.batch_size learning_rate = 1e-3 num_classes = args.num_classes # num class for `weights` file if different, useful in transfer learning with different number of classes weight_num_classes = args.num_weight_class valid_path = args.valid_dataset weights_path = args.weights # Path to text? file containing all classes, 1 per line classes = args.classes anchors = yolo_anchors anchor_masks = yolo_anchor_masks val_dataset = dataset.load_tfrecord_dataset(valid_path, classes, image_size) val_dataset = val_dataset.batch(batch_size) val_dataset = val_dataset.map(lambda x, y: ( dataset.transform_images(x, image_size), dataset.transform_targets(y, anchors, anchor_masks, image_size))) model = YoloV3(image_size, training=True, classes=num_classes) # Darknet transfer is a special case that works # with incompatible number of classes # reset top layers model_pretrained = YoloV3(image_size, training=True, classes=weight_num_classes or num_classes) model_pretrained.load_weights(weights_path) if transfer == 'darknet': model.get_layer('yolo_darknet').set_weights( model_pretrained.get_layer('yolo_darknet').get_weights()) freeze_all(model.get_layer('yolo_darknet')) predictions = [] evaluator = Evaluator(iou_thresh=args.iou) # labels - (N, grid, grid, anchors, [x, y, w, h, obj, class]) boxes, scores, classes, num_detections = model.predict(val_dataset) # boxes -> (num_imgs, num_detections, box coords) # Full labels shape -> [num_batches, grid scale, imgs] # Full labels shape -> [num_batches, [grid, grid, anchors, [x,y,w,h,obj,class]]] full_labels = np.asarray([label for _, label in val_dataset]) # Shape -> [num_batches, num_imgs_in_batch, 3] # Shape -> [num_batches, num_imgs, 3x[grid,grid,anchors,[x,y,w,h,score,class]]] full_labels_trans = full_labels.transpose(0, 2, 1) full_labels_flat = [] for batch in full_labels_trans: for img in batch: row = [] for scale in img: row.append(scale) full_labels_flat.append(row) # Shape -> [num_imgs x 3] full_labels_flat = np.asarray(full_labels_flat) # Remove any labels consisting of all 0's filt_labels = [] # for img in range(len(full_labels_flat)): for img in full_labels_flat: test = [] # for scale in full_labels_flat[img]: for scale in img: lab_list = [] for g1 in scale: for g2 in g1: for anchor in g2: if anchor[0] > 0: temp = [ anchor[0] * image_size, anchor[1] * image_size, anchor[2] * image_size, anchor[3] * image_size, anchor[4], anchor[5] ] temp = [float(x) for x in temp] lab_list.append(np.asarray(temp)) test.append(np.asarray(lab_list)) filt_labels.append(np.asarray(test)) filt_labels = np.asarray( filt_labels ) # Numpy array of shape [num_imgs, 3x[num_boxesx[x1,y1,x2,y2,score,class]]] # filt_labels = filt_labels[:, :4] * image_size # i is the num_images index # predictions = [np.hstack([boxes[i][x], scores[i][x], classes[i][x]]) for i in range(len(num_detections)) for x in range(len(scores[i])) if scores[i][x] > 0] for img in range(len(num_detections)): row = [] for sc in range(len(scores[img])): if scores[img][sc] > 0: row.append( np.hstack([ boxes[img][sc] * image_size, scores[img][sc], classes[img][sc] ])) predictions.append(np.asarray(row)) predictions = np.asarray( predictions) # numpy array of shape [num_imgs x num_preds x 6] if len(predictions) == 0: # No predictions made print('No predictions made - exiting.') exit() # predictions[:, :, 0:4] = predictions[:, :, 0:4] * image_size # Predictions format - [num_imgs x num_preds x [box coords x4, score, classes]] # Box coords should be in format x1 y1 x2 y2 evaluator(predictions, filt_labels, images) # Check gts box coords confidence_thresholds = np.linspace(0.1, 1, 15) confidence_thresholds = [0.5] all_tp_rates = [] all_fp_rates = [] # Compute ROCs for above range of thresholds # Compute one for each class vs. the other classes for index, conf in enumerate(confidence_thresholds): tp_of_img = [] fp_of_img = [] all_classes = [] tp_rates = {} fp_rates = {} boxes, scores, classes, num_detections = model.predict(val_dataset) # Full labels shape -> [num_batches, grid scale, imgs] # Full labels shape -> [num_batches, [grid, grid, anchors, [x,y,w,h,obj,class]]] full_labels = np.asarray([label for _, label in val_dataset]) # Shape -> [num_batches, num_imgs_in_batch, 3] # Shape -> [num_batches, num_imgs, 3x[grid,grid,anchors,[x,y,w,h,score,class]]] full_labels_trans = full_labels.transpose(0, 2, 1) full_labels_flat = [] for batch in full_labels_trans: for img in batch: row = [] for scale in img: row.append(scale) full_labels_flat.append(row) # Shape -> [num_imgs x 3] full_labels_flat = np.asarray(full_labels_flat) # Remove any labels consisting of all 0's filt_labels = [] # for img in range(len(full_labels_flat)): for img in full_labels_flat: test = [] # for scale in full_labels_flat[img]: for scale in img: lab_list = [] for g1 in scale: for g2 in g1: for anchor in g2: if anchor[0] > 0: temp = [ anchor[0] * image_size, anchor[1] * image_size, anchor[2] * image_size, anchor[3] * image_size, anchor[4], anchor[5] ] temp = [float(x) for x in temp] lab_list.append(np.asarray(temp)) test.append(np.asarray(lab_list)) filt_labels.append(np.asarray(test)) filt_labels = np.asarray( filt_labels ) # Numpy array of shape [num_imgs, 3x[num_boxesx[x1,y1,x2,y2,score,class]]] # filt_labels = filt_labels[:, :4] * image_size # i is the num_images index # predictions = [np.hstack([boxes[i][x], scores[i][x], classes[i][x]]) for i in range(len(num_detections)) for x in range(len(scores[i])) if scores[i][x] > 0] for img in range(len(num_detections)): row = [] for sc in range(len(scores[img])): if scores[img][sc] > 0: row.append( np.hstack([ boxes[img][sc] * image_size, scores[img][sc], classes[img][sc] ])) predictions.append(np.asarray(row)) predictions = np.asarray( predictions) # numpy array of shape [num_imgs x num_preds x 6] if len(predictions) == 0: # No predictions made print('No predictions made - exiting.') exit() # predictions[:, :, 0:4] = predictions[:, :, 0:4] * image_size # Predictions format - [num_imgs x num_preds x [box coords x4, score, classes]] # Box coords should be in format x1 y1 x2 y2 evaluator(predictions, filt_labels, images) # Check gts box coords classes = list(set(r['class_ids'])) # All unique class ids for c in classes: if c not in all_classes: all_classes.append(c) complete_classes = dataset_val.class_ids[1:] # Need TPR and FPR rates for each class versus the other classes # Recall == TPR tpr = utils.compute_ap_indiv_class(gt_bbox, gt_class_id, gt_mask, r["rois"], r["class_ids"], r["scores"], r['masks'], complete_classes) total_fpr = utils.compute_fpr_indiv_class(gt_bbox, gt_class_id, gt_mask, r["rois"], r["class_ids"], r["scores"], r['masks'], complete_classes) # print(f'For Image: TPR: {tpr} -- FPR: {total_fpr}') tp_of_img.append(tpr) fp_of_img.append(total_fpr) all_classes = dataset_val.class_ids[1:] # Need to get average TPR and FPR for number of images used for c in all_classes: tp_s = 0 for item in tp_of_img: if c in item.keys(): tp_s += item[c] else: tp_s += 0 tp_rates[c] = tp_s / len(image_ids) # tp_rates[c] = tp_s # print(tp_rates) for c in all_classes: fp_s = 0 for item in fp_of_img: if c in item.keys(): fp_s += item[c] else: fp_s += 0 fp_rates[c] = fp_s / len(image_ids) # fp_rates[c] = fp_s all_fp_rates.append(fp_rates) all_tp_rates.append(tp_rates) print(f'TP Rates: {all_tp_rates}') print(f'FP Rates: {all_fp_rates}')
def main(_argv): try: tpu = tf.distribute.cluster_resolver.TPUClusterResolver() print('Running on TPU ', tpu.cluster_spec().as_dict()['worker']) except: tpu = None if tpu: tf.config.experimental_connect_to_cluster(tpu) tf.tpu.experimental.initialize_tpu_system(tpu) strategy = tf.distribute.experimental.TPUStrategy(tpu) else: strategy = tf.distribute.get_strategy() #print("REPLICAS: ", strategy.num_replicas_in_sync) FLAGS.batch_size = FLAGS.batch_size * strategy.num_replicas_in_sync with strategy.scope(): if FLAGS.tiny: model = YoloV3Tiny(FLAGS.size, training=True, classes=FLAGS.num_classes) anchors = yolo_tiny_anchors anchor_masks = yolo_tiny_anchor_masks else: model = YoloV3(FLAGS.size, training=True, classes=FLAGS.num_classes) anchors = yolo_anchors anchor_masks = yolo_anchor_masks train_dataset = dataset.load_fake_dataset() if FLAGS.dataset: train_dataset = dataset.load_tfrecord_dataset(FLAGS.dataset, FLAGS.classes, FLAGS.size) train_dataset = train_dataset.shuffle(buffer_size=FLAGS.buffer_size) train_dataset = train_dataset.batch(FLAGS.batch_size, drop_remainder=True) train_dataset = train_dataset.map(lambda x, y: ( dataset.transform_images(x, FLAGS.size), dataset.transform_targets(y, anchors, anchor_masks, FLAGS.size))) train_dataset = train_dataset.prefetch( buffer_size=tf.data.experimental.AUTOTUNE) val_dataset = dataset.load_fake_dataset() if FLAGS.val_dataset: val_dataset = dataset.load_tfrecord_dataset(FLAGS.val_dataset, FLAGS.classes, FLAGS.size) val_dataset = val_dataset.batch(FLAGS.batch_size, drop_remainder=True) val_dataset = val_dataset.map(lambda x, y: ( dataset.transform_images(x, FLAGS.size), dataset.transform_targets(y, anchors, anchor_masks, FLAGS.size))) # Configure the model for transfer learning if FLAGS.transfer == 'none': pass # Nothing to do elif FLAGS.transfer in ['darknet', 'no_output']: # Darknet transfer is a special case that works # with incompatible number of classes # reset top layers if FLAGS.tiny: model_pretrained = YoloV3Tiny(FLAGS.size, training=True, classes=FLAGS.weights_num_classes or FLAGS.num_classes) else: model_pretrained = YoloV3(FLAGS.size, training=True, classes=FLAGS.weights_num_classes or FLAGS.num_classes) model_pretrained.load_weights(FLAGS.weights) if FLAGS.transfer == 'darknet': model.get_layer('yolo_darknet').set_weights( model_pretrained.get_layer('yolo_darknet').get_weights()) freeze_all(model.get_layer('yolo_darknet')) elif FLAGS.transfer == 'no_output': for l in model.layers: if not l.name.startswith('yolo_output'): l.set_weights( model_pretrained.get_layer(l.name).get_weights()) freeze_all(l) else: # All other transfer require matching classes model.load_weights(FLAGS.weights) if FLAGS.transfer == 'fine_tune': # freeze darknet and fine tune other layers darknet = model.get_layer('yolo_darknet') freeze_all(darknet) elif FLAGS.transfer == 'frozen': # freeze everything freeze_all(model) if FLAGS.optimizer == 'Adam': optimizer = tf.keras.optimizers.Adam(lr=FLAGS.learning_rate) elif FLAGS.optimizer == 'nAdam': optimizer = tf.keras.optimizers.Nadam(lr=FLAGS.learning_rate) elif FLAGS.optimizer == 'Adagrad': optimizer = tf.keras.optimizers.Adagrad(lr=FLAGS.learning_rate) elif FLAGS.optimizer == 'RMSprop': optimizer = tf.keras.optimizers.RMSprop(lr=FLAGS.learning_rate, rho=0.9) loss = [ YoloLoss(anchors[mask], classes=FLAGS.num_classes) for mask in anchor_masks ] if FLAGS.mode == 'eager_tf': # Eager mode is great for debugging # Non eager graph mode is recommended for real training avg_loss = tf.keras.metrics.Mean('loss', dtype=tf.float32) avg_val_loss = tf.keras.metrics.Mean('val_loss', dtype=tf.float32) for epoch in range(1, FLAGS.epochs + 1): for batch, (images, labels) in enumerate(train_dataset): with tf.GradientTape() as tape: outputs = model(images, training=True) regularization_loss = tf.reduce_sum(model.losses) pred_loss = [] for output, label, loss_fn in zip(outputs, labels, loss): pred_loss.append(loss_fn(label, output)) total_loss = tf.reduce_sum(pred_loss) + regularization_loss grads = tape.gradient(total_loss, model.trainable_variables) optimizer.apply_gradients(zip(grads, model.trainable_variables)) logging.info("{}_train_{}, {}, {}".format( epoch, batch, total_loss.numpy(), list(map(lambda x: np.sum(x.numpy()), pred_loss)))) avg_loss.update_state(total_loss) for batch, (images, labels) in enumerate(val_dataset): outputs = model(images) regularization_loss = tf.reduce_sum(model.losses) pred_loss = [] for output, label, loss_fn in zip(outputs, labels, loss): pred_loss.append(loss_fn(label, output)) total_loss = tf.reduce_sum(pred_loss) + regularization_loss logging.info("{}_val_{}, {}, {}".format( epoch, batch, total_loss.numpy(), list(map(lambda x: np.sum(x.numpy()), pred_loss)))) avg_val_loss.update_state(total_loss) logging.info("{}, train: {}, val: {}".format( epoch, avg_loss.result().numpy(), avg_val_loss.result().numpy())) avg_loss.reset_states() avg_val_loss.reset_states() model.save_weights('checkpoints/yolov3_train_{}.tf'.format(epoch)) else: model.compile(optimizer=optimizer, loss=loss, run_eagerly=(FLAGS.mode == 'eager_fit')) if tpu: callbacks = [ ReduceLROnPlateau(verbose=1), EarlyStopping(patience=3, verbose=1), ModelCheckpoint( 'yolov3_train_{epoch}.h5', save_weights_only=True, verbose=1, period=FLAGS.period ) #, monitor='val_loss', mode='min', save_best_only=True), #1000 ] else: callbacks = [ ReduceLROnPlateau(verbose=1), #EarlyStopping(patience=3, verbose=1), ModelCheckpoint('./checkpoints/yolov3_train_{epoch}.tf', verbose=1, save_weights_only=True, period=FLAGS.period), #1000 TensorBoard(log_dir='logs') ] history = model.fit(train_dataset, epochs=FLAGS.epochs, callbacks=callbacks, validation_data=val_dataset)
def main(_argv): # Change flag values if FLAGS.height is None: FLAGS.height = FLAGS.size if FLAGS.width is None: FLAGS.width = FLAGS.size size = (FLAGS.height, FLAGS.width) physical_devices = tf.config.experimental.list_physical_devices('GPU') for physical_device in physical_devices: tf.config.experimental.set_memory_growth(physical_device, True) if FLAGS.tiny: model = YoloV3Tiny(size, training=True, classes=FLAGS.num_classes, recurrent=FLAGS.recurrent) anchors = yolo_tiny_anchors anchor_masks = yolo_tiny_anchor_masks else: model = YoloV3(size, training=True, classes=FLAGS.num_classes, recurrent=FLAGS.recurrent) anchors = yolo_anchors anchor_masks = yolo_anchor_masks if FLAGS.dataset: train_dataset = dataset.load_tfrecord_dataset(FLAGS.dataset, FLAGS.classes, size) else: train_dataset = dataset.load_fake_dataset() train_dataset = train_dataset.shuffle(buffer_size=8) train_dataset = train_dataset.batch(FLAGS.batch_size) train_dataset = train_dataset.map(lambda x, y: (dataset.transform_images( x, size), dataset.transform_targets(y, anchors, anchor_masks, size))) if FLAGS.recurrent: train_dataset = train_dataset.map( lambda x, y: (dataset.get_recurrect_inputs( x, y, anchors, anchor_masks, FLAGS.num_classes), y)) train_dataset = train_dataset.prefetch( buffer_size=tf.data.experimental.AUTOTUNE) if FLAGS.val_dataset: val_dataset = dataset.load_tfrecord_dataset(FLAGS.val_dataset, FLAGS.classes, size) else: val_dataset = dataset.load_fake_dataset() val_dataset = val_dataset.batch(FLAGS.batch_size) val_dataset = val_dataset.map(lambda x, y: (dataset.transform_images( x, size), dataset.transform_targets(y, anchors, anchor_masks, size))) if FLAGS.recurrent: val_dataset = val_dataset.map( lambda x, y: (dataset.get_recurrect_inputs( x, y, anchors, anchor_masks, FLAGS.num_classes), y)) # Configure the model for transfer learning if FLAGS.transfer != 'none': # if we need all weights, no need to create another model if FLAGS.transfer == 'all': model.load_weights(FLAGS.weights) # else, we need only some of the weights # create appropriate model_pretrained, load all weights and copy the ones we need else: if FLAGS.tiny: model_pretrained = YoloV3Tiny(size, training=True, classes=FLAGS.weights_num_classes or FLAGS.num_classes, recurrent=FLAGS.recurrent) else: model_pretrained = YoloV3(size, training=True, classes=FLAGS.weights_num_classes or FLAGS.num_classes, recurrent=FLAGS.recurrent) # load pretrained weights model_pretrained.load_weights(FLAGS.weights) # transfer darknet darknet = model.get_layer('yolo_darknet') darknet.set_weights( model_pretrained.get_layer('yolo_darknet').get_weights()) # transfer 'yolo_conv_i' layer weights if FLAGS.transfer in [ 'yolo_conv', 'yolo_output_conv', 'yolo_output' ]: for l in model.layers: if l.name.startswith('yolo_conv'): model.get_layer(l.name).set_weights( model_pretrained.get_layer(l.name).get_weights()) # transfer 'yolo_output_i' first conv2d layer if FLAGS.transfer == 'yolo_output_conv': # transfer tiny output conv2d for l in model.layers: if l.name.startswith('yolo_output'): # get and set the weights of the appropriate layers model.get_layer(l.name).layers[1].set_weights( model_pretrained.get_layer( l.name).layers[1].get_weights()) # should I freeze batch_norm as well? # transfer 'yolo_output_i' layer weights if FLAGS.transfer == 'yolo_output': for l in model.layers: if l.name.startswith('yolo_output'): model.get_layer(l.name).set_weights( model_pretrained.get_layer(l.name).get_weights()) # no transfer learning else: pass # freeze layers, if requested if FLAGS.freeze != 'none': if FLAGS.freeze == 'all': freeze_all(model) if FLAGS.freeze in [ 'yolo_darknet' 'yolo_conv', 'yolo_output_conv', 'yolo_output' ]: darknet = model.get_layer('yolo_darknet') freeze_all(darknet) if FLAGS.freeze in ['yolo_conv', 'yolo_output_conv', 'yolo_output']: for l in model.layers: if l.name.startswith('yolo_conv'): freeze_all(l) if FLAGS.freeze == 'yolo_output_conv': if FLAGS.tiny: # freeze the appropriate layers freeze_all(model.layers[4].layers[1]) freeze_all(model.layers[5].layers[1]) else: # freeze the appropriate layers freeze_all(model.layers[5].layers[1]) freeze_all(model.layers[6].layers[1]) freeze_all(model.layers[7].layers[1]) if FLAGS.transfer == 'yolo_output': for l in model.layers: if l.name.startswith('yolo_output'): freeze_all(l) # freeze nothing else: pass optimizer = tf.keras.optimizers.Adam(lr=FLAGS.learning_rate) loss = [ YoloLoss(anchors[mask], classes=FLAGS.num_classes) for mask in anchor_masks ] if FLAGS.mode == 'eager_tf': # Eager mode is great for debugging # Non eager graph mode is recommended for real training avg_loss = tf.keras.metrics.Mean('loss', dtype=tf.float32) avg_val_loss = tf.keras.metrics.Mean('val_loss', dtype=tf.float32) for epoch in range(1, FLAGS.epochs + 1): for batch, (images, labels) in enumerate(train_dataset): with tf.GradientTape() as tape: outputs = model(images, training=True) regularization_loss = tf.reduce_sum(model.losses) pred_loss = [] for output, label, loss_fn in zip(outputs, labels, loss): pred_loss.append(loss_fn(label, output)) total_loss = tf.reduce_sum(pred_loss) + regularization_loss grads = tape.gradient(total_loss, model.trainable_variables) optimizer.apply_gradients(zip(grads, model.trainable_variables)) logging.info("{}_train_{}, {}, {}".format( epoch, batch, total_loss.numpy(), list(map(lambda x: np.sum(x.numpy()), pred_loss)))) avg_loss.update_state(total_loss) for batch, (images, labels) in enumerate(val_dataset): outputs = model(images) regularization_loss = tf.reduce_sum(model.losses) pred_loss = [] for output, label, loss_fn in zip(outputs, labels, loss): pred_loss.append(loss_fn(label, output)) total_loss = tf.reduce_sum(pred_loss) + regularization_loss logging.info("{}_val_{}, {}, {}".format( epoch, batch, total_loss.numpy(), list(map(lambda x: np.sum(x.numpy()), pred_loss)))) avg_val_loss.update_state(total_loss) logging.info("{}, train: {}, val: {}".format( epoch, avg_loss.result().numpy(), avg_val_loss.result().numpy())) avg_loss.reset_states() avg_val_loss.reset_states() model.save_weights('checkpoints/yolov3_train_{}.tf'.format(epoch)) else: model.compile(optimizer=optimizer, loss=loss, run_eagerly=(FLAGS.mode == 'eager_fit')) callbacks = [ ReduceLROnPlateau(verbose=1), EarlyStopping(patience=3, verbose=1), ModelCheckpoint('checkpoints/yolov3_train_{epoch}.tf', verbose=1, save_weights_only=True), TensorBoard(log_dir='logs') ] history = model.fit(train_dataset, epochs=FLAGS.epochs, callbacks=callbacks, validation_data=val_dataset)
def main(_argv): physical_devices = tf.config.experimental.list_physical_devices('GPU') if len(physical_devices) > 0: # 设置仅在需要时申请显存空间 tf.config.experimental.set_memory_growth(physical_devices[0], True) # 判断训练tiny版本的YOLO还是完整版的YOLO if FLAGS.tiny: model = YoloV3Tiny(FLAGS.size, training=True, classes=FLAGS.num_classes) anchors = yolo_tiny_anchors anchor_masks = yolo_tiny_anchor_masks else: model = YoloV3(FLAGS.size, training=True, classes=FLAGS.num_classes) anchors = yolo_anchors anchor_masks = yolo_anchor_masks # 如果未指定数据集则加载一张图片作为数据集=>fake_dataset train_dataset = dataset.load_fake_dataset() # 判断数据集路径是否为空 if FLAGS.dataset: # 从TFRecode文件加载数据集 train_dataset:(x_train, y_train) train_dataset = dataset.load_tfrecord_dataset(FLAGS.dataset, FLAGS.classes, FLAGS.size) # 生成批训练数据 # 打乱数据顺序 train_dataset = train_dataset.shuffle(buffer_size=512) train_dataset = train_dataset.batch(FLAGS.batch_size) # y.shape:train_dataset.as_numpy_iterator().next()[1].shape # =>(batch_size, yolo_max_boxes, 5) 5=>(xmin, ymin, xmax, ymax, classlabel) train_dataset = train_dataset.map(lambda x, y: ( # 图像数据归一化[0,1] dataset.transform_images(x, FLAGS.size), # 根据先验框anchor确定bbox属于哪一层特征图(13*13, 26*26, 52*52) # 并计算出bbox的中心点在特征图上的位置 dataset.transform_targets(y, anchors, anchor_masks, FLAGS.size))) # 数据预读取,提高延迟和吞吐量 # tf.data.experimental.AUTOTUNE:根据可用CPU动态设置并行调用的数量 train_dataset = train_dataset.prefetch( buffer_size=tf.data.experimental.AUTOTUNE) # 加载伪验证集,防止没有添加验证集路径时报错 val_dataset = dataset.load_fake_dataset() # 加载验证集 if FLAGS.val_dataset: val_dataset = dataset.load_tfrecord_dataset(FLAGS.val_dataset, FLAGS.classes, FLAGS.size) val_dataset = val_dataset.batch(FLAGS.batch_size) val_dataset = val_dataset.map(lambda x, y: ( dataset.transform_images(x, FLAGS.size), dataset.transform_targets(y, anchors, anchor_masks, FLAGS.size))) # Configure the model for transfer learning # 训练模式选择 # 随机初始化权重,从0开始训练整个网络 if FLAGS.transfer == 'none': pass # Nothing to do # 迁移训练的两种方式 elif FLAGS.transfer in ['darknet', 'no_output']: # Darknet transfer is a special case that works # with incompatible number of classes # reset top layers if FLAGS.tiny: model_pretrained = YoloV3Tiny(FLAGS.size, training=True, classes=FLAGS.weights_num_classes or FLAGS.num_classes) else: # 模型网络结构 model_pretrained = YoloV3(FLAGS.size, training=True, classes=FLAGS.weights_num_classes or FLAGS.num_classes) # 加载预训练权重 model_pretrained.load_weights(FLAGS.weights) # 设置darknet网络权重并冻结网络,即主干网络不参与训练,其余参数随机初始化 if FLAGS.transfer == 'darknet': model.get_layer('yolo_darknet').set_weights( model_pretrained.get_layer('yolo_darknet').get_weights()) freeze_all(model.get_layer('yolo_darknet')) # 设置YOLO输出层以外的网络的权重并冻结, 即仅训练YOLO的输出层且参数随机初始化 elif FLAGS.transfer == 'no_output': for l in model.layers: if not l.name.startswith('yolo_output'): l.set_weights( model_pretrained.get_layer(l.name).get_weights()) freeze_all(l) # 迁移学习fine_tune和frozen模式要求训练的类别数和预训练权重一致(80类) else: # All other transfer require matching classes # 加载网络所有预训练权重参数 model.load_weights(FLAGS.weights) # 冻结darknet(骨干网络)权重, 其余参数在预训练权重的基础上训练 if FLAGS.transfer == 'fine_tune': # freeze darknet and fine tune other layers darknet = model.get_layer('yolo_darknet') freeze_all(darknet) # 冻结所有参数,训练不起作用. elif FLAGS.transfer == 'frozen': # freeze everything freeze_all(model) # 定义优化器:Adam optimizer = tf.keras.optimizers.Adam(lr=FLAGS.learning_rate) loss = [ YoloLoss(anchors[mask], classes=FLAGS.num_classes) for mask in anchor_masks ] # 调试模型:速度慢: Eager: op 在调用后会立即运行 if FLAGS.mode == 'eager_tf': # Eager mode is great for debugging # Non eager graph mode is recommended for real training # 训练集上的平均loss/验证集上的平均loss avg_loss = tf.keras.metrics.Mean('loss', dtype=tf.float32) avg_val_loss = tf.keras.metrics.Mean('val_loss', dtype=tf.float32) # 迭代每个epoch for epoch in range(1, FLAGS.epochs + 1): for batch, (images, labels) in enumerate(train_dataset): # 梯度带:自动计算变量梯度 with tf.GradientTape() as tape: # model(): eager模式下选择此方式,不需要编译直接运行, 速度快. # model.predict()第一次运行时需要先编译图模式 outputs = model(images, training=True) # 计算张量各维度的元素之和. regularization_loss = tf.reduce_sum(model.losses) pred_loss = [] for output, label, loss_fn in zip(outputs, labels, loss): pred_loss.append(loss_fn(label, output)) total_loss = tf.reduce_sum(pred_loss) + regularization_loss # 梯度 grads = tape.gradient(total_loss, model.trainable_variables) # 执行最优化器 optimizer.apply_gradients(zip(grads, model.trainable_variables)) # 记录日志文件 logging.info("{}_train_{}, {}, {}".format( epoch, batch, total_loss.numpy(), list(map(lambda x: np.sum(x.numpy()), pred_loss)))) # 更新平均loss avg_loss.update_state(total_loss) # 在验证集上验证 for batch, (images, labels) in enumerate(val_dataset): outputs = model(images) regularization_loss = tf.reduce_sum(model.losses) pred_loss = [] for output, label, loss_fn in zip(outputs, labels, loss): pred_loss.append(loss_fn(label, output)) total_loss = tf.reduce_sum(pred_loss) + regularization_loss logging.info("{}_val_{}, {}, {}".format( epoch, batch, total_loss.numpy(), list(map(lambda x: np.sum(x.numpy()), pred_loss)))) avg_val_loss.update_state(total_loss) # .result():返回累计结果 logging.info("{}, train: {}, val: {}".format( epoch, avg_loss.result().numpy(), avg_val_loss.result().numpy())) # reset_states:清除累计值 avg_loss.reset_states() avg_val_loss.reset_states() # 每个epoch保存一次模型权重 model.save_weights('checkpoints/yolov3_train_{}.tf'.format(epoch)) # 训练模式 else: # 编译模型 model.compile(optimizer=optimizer, loss=loss, metrics=['accuracy'], run_eagerly=(FLAGS.mode == 'eager_fit')) # 回调函数 callbacks = [ # lr衰减 ReduceLROnPlateau(verbose=1), # lr不变时停止训练 EarlyStopping(patience=3, verbose=1), # 保存模型 ModelCheckpoint('checkpoints/yolov3_train_{epoch}.tf', verbose=1, save_weights_only=True), # 训练结果可视化 TensorBoard(log_dir='logs', write_images=True, update_freq='batch') ] # 进行迭代训练 history = model.fit(train_dataset, epochs=FLAGS.epochs, callbacks=callbacks, validation_data=val_dataset)