INPUT_SIZE = 416 BATCH_SIZE = 1 EPOCHS = 20 LR = 0.001 SHUFFLE_SIZE = 1 weights_path = "/home/yang/test/yolov3.weights" sess = tf.Session() classes = utils.read_coco_names('./data/coco.names') num_classes = len(classes) file_pattern = "./data/train_data/quick_train_data/tfrecords/quick_train_data*.tfrecords" anchors = utils.get_anchors('./data/yolo_anchors.txt') dataset = tf.data.TFRecordDataset(filenames=tf.gfile.Glob(file_pattern)) dataset = dataset.map(utils.parser(anchors, num_classes).parser_example, num_parallel_calls=10) dataset = dataset.repeat().shuffle(SHUFFLE_SIZE).batch(BATCH_SIZE).prefetch( BATCH_SIZE) iterator = dataset.make_one_shot_iterator() example = iterator.get_next() images, *y_true = example model = yolov3.yolov3(num_classes) with tf.variable_scope('yolov3'): y_pred = model.forward(images, is_training=False) loss = model.compute_loss(y_pred, y_true) y_pred = model.predict(y_pred) load_ops = utils.load_weights(tf.global_variables(scope='yolov3'), weights_path) sess.run(load_ops)
INPUT_SIZE = 416 BATCH_SIZE = 16 EPOCHS = 5000000 LR = 0.0001 SHUFFLE_SIZE = 10000 sess = tf.Session() classes = utils.read_coco_names('./data/coco.names') num_classes = len(classes) # file_pattern = "../COCO/tfrecords/coco*.tfrecords" file_pattern = "./data/train_data/quick_train_data/tfrecords/quick_train_data*.tfrecords" anchors = utils.get_anchors('./data/yolo_anchors.txt') is_training = tf.placeholder(dtype=tf.bool, name="phase_train") dataset = tf.data.TFRecordDataset(filenames = tf.gfile.Glob(file_pattern)) dataset = dataset.map(utils.parser(anchors, num_classes).parser_example, num_parallel_calls = 10) dataset = dataset.repeat().shuffle(SHUFFLE_SIZE).batch(BATCH_SIZE).prefetch(BATCH_SIZE) iterator = dataset.make_one_shot_iterator() example = iterator.get_next() images, *y_true = example model = yolov3.yolov3(num_classes) with tf.variable_scope('yolov3'): y_pred = model.forward(images, is_training=is_training) loss = model.compute_loss(y_pred, y_true) y_pred = model.predict(y_pred) optimizer = tf.train.MomentumOptimizer(LR, momentum=0.9) train_op = optimizer.minimize(loss[0]) saver = tf.train.Saver(max_to_keep=2)