def train(self):
for epoch in range(self.first_stage_epochs + self.second_stage_epochs):
if epoch < self.first_stage_epochs:
if not self.isfreeze:
self.isfreeze = True
for name in self.freeze_layers:
freeze = self.model.get_layer(name)
freeze_all(freeze)
elif epoch >= self.first_stage_epochs:
if self.isfreeze:
self.isfreeze = False
for name in self.freeze_layers:
freeze = self.model.get_layer(name)
unfreeze_all(freeze)
for image_data, target in self.trainset:
self.train_step(image_data, target)
for image_data, target in self.testset:
self.test_step(image_data, target)
self.model.save_weights("{}/checkpoints/yolov4".format(
self.base_path))
def main(_argv):
physical_devices = tf.config.experimental.list_physical_devices('GPU')
# if len(physical_devices) > 0:
if physical_devices:
# tf.config.experimental.set_memory_growth(physical_devices[0], True)
tf.config.experimental.set_visible_devices(physical_devices[0], "GPU")
trainset = Dataset(FLAGS, is_training=True)
testset = Dataset(FLAGS, is_training=False)
logdir = "./data/log"
isfreeze = False
steps_per_epoch = len(trainset)
first_stage_epochs = cfg.TRAIN.FISRT_STAGE_EPOCHS
second_stage_epochs = cfg.TRAIN.SECOND_STAGE_EPOCHS
global_steps = tf.Variable(1, trainable=False, dtype=tf.int64)
warmup_steps = cfg.TRAIN.WARMUP_EPOCHS * steps_per_epoch
total_steps = (first_stage_epochs + second_stage_epochs) * steps_per_epoch
# train_steps = (first_stage_epochs + second_stage_epochs) * steps_per_period
input_layer = tf.keras.layers.Input([cfg.TRAIN.INPUT_SIZE, cfg.TRAIN.INPUT_SIZE, 3])
STRIDES, ANCHORS, NUM_CLASS, XYSCALE = utils.load_config(FLAGS)
IOU_LOSS_THRESH = cfg.YOLO.IOU_LOSS_THRESH
freeze_layers = utils.load_freeze_layer(FLAGS.model, FLAGS.tiny)
feature_maps = YOLO(input_layer, NUM_CLASS, FLAGS.model, FLAGS.tiny)
if FLAGS.tiny:
bbox_tensors = []
for i, fm in enumerate(feature_maps):
if i == 0:
bbox_tensor = decode_train(fm, cfg.TRAIN.INPUT_SIZE // 16, NUM_CLASS, STRIDES, ANCHORS, i, XYSCALE)
else:
bbox_tensor = decode_train(fm, cfg.TRAIN.INPUT_SIZE // 32, NUM_CLASS, STRIDES, ANCHORS, i, XYSCALE)
bbox_tensors.append(fm)
bbox_tensors.append(bbox_tensor)
else:
bbox_tensors = []
for i, fm in enumerate(feature_maps):
if i == 0:
bbox_tensor = decode_train(fm, cfg.TRAIN.INPUT_SIZE // 8, NUM_CLASS, STRIDES, ANCHORS, i, XYSCALE)
elif i == 1:
bbox_tensor = decode_train(fm, cfg.TRAIN.INPUT_SIZE // 16, NUM_CLASS, STRIDES, ANCHORS, i, XYSCALE)
else:
bbox_tensor = decode_train(fm, cfg.TRAIN.INPUT_SIZE // 32, NUM_CLASS, STRIDES, ANCHORS, i, XYSCALE)
bbox_tensors.append(fm)
bbox_tensors.append(bbox_tensor)
model = tf.keras.Model(input_layer, bbox_tensors)
model.summary()
if FLAGS.weights == None:
print("Training from scratch")
else:
if FLAGS.weights.split(".")[len(FLAGS.weights.split(".")) - 1] == "weights":
utils.load_weights(model, FLAGS.weights, FLAGS.model, FLAGS.tiny)
else:
model.load_weights(FLAGS.weights)
print('Restoring weights from: %s ... ' % FLAGS.weights)
optimizer = tf.keras.optimizers.Adam()
if os.path.exists(logdir): shutil.rmtree(logdir)
writer = tf.summary.create_file_writer(logdir)
# define training step function
# @tf.function
def train_step(image_data, target):
with tf.GradientTape() as tape:
pred_result = model(image_data, training=True)
giou_loss = conf_loss = prob_loss = 0
# optimizing process
for i in range(len(freeze_layers)):
conv, pred = pred_result[i * 2], pred_result[i * 2 + 1]
loss_items = compute_loss(pred, conv, target[i][0], target[i][1], STRIDES=STRIDES, NUM_CLASS=NUM_CLASS, IOU_LOSS_THRESH=IOU_LOSS_THRESH, i=i)
giou_loss += loss_items[0]
conf_loss += loss_items[1]
prob_loss += loss_items[2]
total_loss = giou_loss + conf_loss + prob_loss
gradients = tape.gradient(total_loss, model.trainable_variables)
optimizer.apply_gradients(zip(gradients, model.trainable_variables))
tf.print("=> STEP %4d/%4d lr: %.6f giou_loss: %4.2f conf_loss: %4.2f "
"prob_loss: %4.2f total_loss: %4.2f" % (global_steps, total_steps, optimizer.lr.numpy(),
giou_loss, conf_loss,
prob_loss, total_loss))
# update learning rate
global_steps.assign_add(1)
if global_steps < warmup_steps:
lr = global_steps / warmup_steps * cfg.TRAIN.LR_INIT
else:
lr = cfg.TRAIN.LR_END + 0.5 * (cfg.TRAIN.LR_INIT - cfg.TRAIN.LR_END) * (
(1 + tf.cos((global_steps - warmup_steps) / (total_steps - warmup_steps) * np.pi))
)
optimizer.lr.assign(lr.numpy())
# writing summary data
with writer.as_default():
tf.summary.scalar("lr", optimizer.lr, step=global_steps)
tf.summary.scalar("loss/total_loss", total_loss, step=global_steps)
tf.summary.scalar("loss/giou_loss", giou_loss, step=global_steps)
tf.summary.scalar("loss/conf_loss", conf_loss, step=global_steps)
tf.summary.scalar("loss/prob_loss", prob_loss, step=global_steps)
writer.flush()
def test_step(image_data, target):
with tf.GradientTape() as tape:
pred_result = model(image_data, training=True)
giou_loss = conf_loss = prob_loss = 0
# optimizing process
for i in range(len(freeze_layers)):
conv, pred = pred_result[i * 2], pred_result[i * 2 + 1]
loss_items = compute_loss(pred, conv, target[i][0], target[i][1], STRIDES=STRIDES, NUM_CLASS=NUM_CLASS, IOU_LOSS_THRESH=IOU_LOSS_THRESH, i=i)
giou_loss += loss_items[0]
conf_loss += loss_items[1]
prob_loss += loss_items[2]
total_loss = giou_loss + conf_loss + prob_loss
tf.print("=> TEST STEP %4d giou_loss: %4.2f conf_loss: %4.2f "
"prob_loss: %4.2f total_loss: %4.2f" % (global_steps, giou_loss, conf_loss,
prob_loss, total_loss))
for epoch in range(first_stage_epochs + second_stage_epochs):
if epoch < first_stage_epochs:
if not isfreeze:
isfreeze = True
for name in freeze_layers:
freeze = model.get_layer(name)
freeze_all(freeze)
elif epoch >= first_stage_epochs:
if isfreeze:
isfreeze = False
for name in freeze_layers:
freeze = model.get_layer(name)
unfreeze_all(freeze)
for image_data, target in trainset:
train_step(image_data, target)
for image_data, target in testset:
test_step(image_data, target)
model.save_weights("./checkpoints/yolov4")
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)
model = Yolo(FLAGS.size,
training=True,
classes=FLAGS.num_classes,
aux=True)
anchors, mask_len = get_anchors(FLAGS.size)
train_dataset = dataset.load_tfrecord_dataset(FLAGS.dataset, FLAGS.size)
if FLAGS.check:
logging.info('Current dataset check')
with open(FLAGS.classes) as classes_f:
class_names = [c.strip() for c in classes_f.readlines()]
logging.info('Classes loaded')
check_dataset(train_dataset, class_names)
train_dataset = train_dataset.repeat()
train_dataset = train_dataset.shuffle(buffer_size=200)
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, mask_len, FLAGS.size)))
train_dataset = train_dataset.prefetch(
buffer_size=tf.data.experimental.AUTOTUNE)
if FLAGS.validation_dataset:
validation_dataset = dataset.load_tfrecord_dataset(
FLAGS.validation_dataset, FLAGS.classes, FLAGS.size)
else:
validation_dataset = dataset.load_dummy_validation()
validation_dataset = validation_dataset.batch(FLAGS.batch_size)
validation_dataset = validation_dataset.map(lambda x, y: (
dataset.transform_images(x, FLAGS.size),
dataset.transform_targets(y, anchors, mask_len, FLAGS.size)))
if FLAGS.weights:
if FLAGS.pretrained:
model_pretrained = Yolo(FLAGS.size,
training=True,
classes=FLAGS.weights_num_classes)
model_pretrained.load_weights(FLAGS.weights)
model.get_layer('darknet_recursive').set_weights(
model_pretrained.get_layer('darknet_recursive').get_weights())
else:
model.load_weights(FLAGS.weights)
if FLAGS.freeze:
freeze_all(model.get_layer('darknet_recursive'))
optimizer = tf.keras.optimizers.Adam(lr=FLAGS.learning_rate)
loss = [
Loss(anchors[:mask_len], classes=FLAGS.num_classes),
Loss(anchors[mask_len:], classes=FLAGS.num_classes)
]
model.compile(optimizer=optimizer,
loss=loss,
run_eagerly=True,
metrics=['accuracy'])
callbacks = [
ReduceLROnPlateau(verbose=1, patience=10),
EarlyStopping(patience=3, verbose=1),
ModelCheckpoint('checkpoints/yolov3_face.tf',
verbose=1,
save_weights_only=True),
TensorBoard(log_dir='logs')
]
history = model.fit(train_dataset,
epochs=FLAGS.epochs,
callbacks=callbacks,
validation_data=validation_dataset,
steps_per_epoch=10)
def main(_argv):
input_channel = 3
patience = 30
steps_in_epoch = 0
epoch_loss = np.inf
prev_minloss = np.inf
trainset = Dataset(FLAGS,input_channel, is_training=True)
# testset = Dataset(FLAGS, input_channel, is_training=False)
logdir = "./data/log"
isfreeze = False
steps_per_epoch = len(trainset)
print("steps_per_epoch:{}".format(steps_per_epoch))
first_stage_epochs = cfg.TRAIN.FISRT_STAGE_EPOCHS
second_stage_epochs = cfg.TRAIN.SECOND_STAGE_EPOCHS
global_steps = tf.Variable(1, trainable=False, dtype=tf.int64)
warmup_steps = cfg.TRAIN.WARMUP_EPOCHS * steps_per_epoch
total_steps = (first_stage_epochs + second_stage_epochs) * steps_per_epoch
# train_steps = (first_stage_epochs + second_stage_epochs) * steps_per_period
loss_tracker = []
losses_in_epoch = 0.
input_layer = tf.keras.layers.Input([cfg.TRAIN.INPUT_SIZE, cfg.TRAIN.INPUT_SIZE, 3])
STRIDES, ANCHORS, NUM_CLASS, XYSCALE = utils.load_config(FLAGS)
IOU_LOSS_THRESH = cfg.YOLO.IOU_LOSS_THRESH
freeze_layers = utils.load_freeze_layer(FLAGS.model, FLAGS.num_detection_layer)
# feature_maps = YOLO(input_layer, NUM_CLASS, FLAGS.model, FLAGS.tiny)
feature_maps = YOLOv4_more_tiny(input_layer, NUM_CLASS)
bbox_tensors = []
for i, fm in enumerate(feature_maps): # fm shape: (None, featw, feath, filters)
bbox_tensor = decode_train(fm, cfg.TRAIN.INPUT_SIZE // 16, NUM_CLASS, STRIDES, ANCHORS, i, XYSCALE)
bbox_tensors.append(fm)
bbox_tensors.append(bbox_tensor)
if cfg.YOLO.NUM_YOLOLAYERS == 3: # yolov4
bbox_tensors = []
for i, fm in enumerate(feature_maps):
if i == 0:
bbox_tensor = decode_train(fm, cfg.TRAIN.INPUT_SIZE // 8, NUM_CLASS, STRIDES, ANCHORS, i, XYSCALE)
elif i == 1:
bbox_tensor = decode_train(fm, cfg.TRAIN.INPUT_SIZE // 16, NUM_CLASS, STRIDES, ANCHORS, i, XYSCALE)
else:
bbox_tensor = decode_train(fm, cfg.TRAIN.INPUT_SIZE // 32, NUM_CLASS, STRIDES, ANCHORS, i, XYSCALE)
bbox_tensors.append(fm)
bbox_tensors.append(bbox_tensor)
elif cfg.YOLO.NUM_YOLOLAYERS == 2: # yolo tiny
bbox_tensors = []
for i, fm in enumerate(feature_maps):
if i == 0:
bbox_tensor = decode_train(fm, cfg.TRAIN.INPUT_SIZE // 16, NUM_CLASS, STRIDES, ANCHORS, i, XYSCALE)
else:
bbox_tensor = decode_train(fm, cfg.TRAIN.INPUT_SIZE // 32, NUM_CLASS, STRIDES, ANCHORS, i, XYSCALE)
bbox_tensors.append(fm)
bbox_tensors.append(bbox_tensor)
elif cfg.YOLO.NUM_YOLOLAYERS == 1: # custom yolo
bbox_tensors = []
bbox_tensor = decode_train(feature_maps[0], cfg.TRAIN.INPUT_SIZE // cfg.YOLO.STRIDES_CUSTOM[0], NUM_CLASS, STRIDES, ANCHORS, 0, XYSCALE)
bbox_tensors.append(feature_maps[0])
bbox_tensors.append(bbox_tensor)
model = tf.keras.Model(input_layer, bbox_tensors)
model.summary()
if FLAGS.weights == None:
print("Training from scratch")
else:
if FLAGS.weights.split(".")[len(FLAGS.weights.split(".")) - 1] == "weights":
utils.load_weights(model, FLAGS.weights, FLAGS.model, FLAGS.num_detection_layer)
else:
model.load_weights(FLAGS.weights)
print('Restoring weights from: %s ... ' % FLAGS.weights)
optimizer = tf.keras.optimizers.Adam()
if os.path.exists(logdir): shutil.rmtree(logdir)
writer = tf.summary.create_file_writer(logdir)
# define training step function
# @tf.function
def train_step(image_data, target):
with tf.GradientTape() as tape:
pred_result = model(image_data, training=True)
ciou_loss = conf_loss = prob_loss = 0
# optimizing process
for i in range(len(freeze_layers)):
conv, pred = pred_result[i * 2], pred_result[i * 2 + 1] # target[i][1]:(32, 150, 4)
# print("conv shape:{} pred shape:{}".format(tf.keras.backend.int_shape(conv), tf.keras.backend.int_shape(pred)))
# print("target[i][0]:{} target[i][1]:{}".format(np.array(target[i][0]).shape, np.array(target[i][1]).shape))
loss_items = compute_loss(pred, conv, target[i][0], target[i][1], STRIDES=STRIDES, NUM_CLASS=NUM_CLASS, IOU_LOSS_THRESH=IOU_LOSS_THRESH, i=i)
ciou_loss += loss_items[0]
conf_loss += loss_items[1]
prob_loss += loss_items[2]
total_loss = ciou_loss + conf_loss + prob_loss
gradients = tape.gradient(total_loss, model.trainable_variables)
optimizer.apply_gradients(zip(gradients, model.trainable_variables))
tf.print("=> STEP %4d/%4d lr: %.6f ciou_loss: %4.2f conf_loss: %4.2f "
"prob_loss: %4.2f total_loss: %4.2f" % (global_steps, total_steps, optimizer.lr.numpy(),
ciou_loss, conf_loss,
prob_loss, total_loss))
# update learning rate
global_steps.assign_add(1)
if global_steps < warmup_steps:
lr = global_steps / warmup_steps * cfg.TRAIN.LR_INIT
else:
lr = cfg.TRAIN.LR_END + 0.5 * (cfg.TRAIN.LR_INIT - cfg.TRAIN.LR_END) * (
(1 + tf.cos((global_steps - warmup_steps) / (total_steps - warmup_steps) * np.pi))
)
optimizer.lr.assign(lr.numpy())
# writing summary data
with writer.as_default():
tf.summary.scalar("lr", optimizer.lr, step=global_steps)
tf.summary.scalar("loss/total_loss", total_loss, step=global_steps)
tf.summary.scalar("loss/ciou_loss", ciou_loss, step=global_steps)
tf.summary.scalar("loss/conf_loss", conf_loss, step=global_steps)
tf.summary.scalar("loss/prob_loss", prob_loss, step=global_steps)
writer.flush()
# return total_loss
@tf.function
def test_step(image_data, target):
with tf.GradientTape() as tape:
pred_result = model(image_data, training=True)
ciou_loss = conf_loss = prob_loss = 0
preds = None
gt_infos = None
gt_bboxes = None
# optimizing process
for i in range(len(freeze_layers)):
conv, pred = pred_result[i * 2], pred_result[i * 2 + 1]
loss_items = compute_loss(pred, conv, target[i][0], target[i][1], STRIDES=STRIDES, NUM_CLASS=NUM_CLASS, IOU_LOSS_THRESH=IOU_LOSS_THRESH, i=i)
ciou_loss += loss_items[0]
conf_loss += loss_items[1]
prob_loss += loss_items[2]
if FLAGS.show_map:
batch_, _, _, _, ch_ = tf.keras.backend.int_shape(pred)
if preds == None:
preds = tf.reshape(pred, (batch_, -1, ch_))
else:
preds = tf.concat([preds, tf.reshape(preds, (batch_, -1, ch_))], axis=1)
if gt_infos == None:
gt_infos = tf.reshape(target[i][0], (batch_, -1, ch_))
gt_bboxes = target[i][1]
else:
gt_infos = tf.concat([gt_infos, tf.reshape(target[i][0], (batch_, -1, ch_))], axis=1)
gt_bboxes = tf.concat([gt_bboxes, tf.reshape(target[i][1], (batch_, -1, 4))], axis=1)
if FLAGS.show_map:
map = compute_map(preds, gt_bboxes, gt_infos, NUM_CLASS)
total_map = map
total_loss = ciou_loss + conf_loss + prob_loss
tf.print("=> TEST STEP %4d ciou_loss: %4.2f conf_loss: %4.2f "
"prob_loss: %4.2f total_loss: %4.2f" % (global_steps, ciou_loss, conf_loss,
prob_loss, total_loss))
tf.print("=> TEST STEP %4d map: %4.2f" % (total_map))
# tf.print("=> TEST STEP %4d giou_loss: %4.2f conf_loss: %4.2f "
# "prob_loss: %4.2f total_loss: %4.2f" % (global_steps, giou_loss, conf_loss,
# prob_loss, total_loss))
for epoch in range(first_stage_epochs + second_stage_epochs):
start_ch = time.time()
train_loss = 0.
if epoch < first_stage_epochs:
if not isfreeze:
isfreeze = True
for name in freeze_layers:
try:
freeze = model.get_layer(name)
freeze_all(freeze)
except ValueError:
pass
elif epoch >= first_stage_epochs:
if isfreeze:
isfreeze = False
for name in freeze_layers:
try: # try구문 추가
freeze = model.get_layer(name)
freeze_all(freeze)
except ValueError:
pass
for image_data, target in trainset:
# bboxes_list = target[0][1]
# label_data = target[0][0]
#
# for batch_idx, bboxes in enumerate(bboxes_list):
# # print("bboxes:{}".format(bboxes))
# class_inds = []
# check = np.array(image_data[batch_idx]).reshape(cfg.TRAIN.INPUT_SIZE, cfg.TRAIN.INPUT_SIZE, 3)
# # print(r"D:\tf_data_not\preprocessed_{}.npy".format(batch_idx))
# # np.save(r"D:\tf_data_not\image_{}.npy".format(batch_idx), check)
# # np.save(r"D:\tf_data_not\preprocessed_{}.npy".format(batch_idx), label_data[batch_idx])
# # np.save(r"D:\tf_data_not\bboxes_{}.npy".format(batch_idx), bboxes)
# label = label_data[batch_idx]
# label_class = label[...,5:]
# # print("label shape:{}".format(np.array(label).shape))
#
# for line_label in label_class:
# if np.sum(line_label) > 0:
# print(line_label)
# class_ = np.array(label[..., 5:]).flatten()
# class_ = np.where(class_>0.1, class_, 0)
# for class_idx, class_label in enumerate(class_):
# if class_label > 0:
# class_inds.append(class_idx % cfg.YOLO.NUM_CLASSES)
#
# for bbox_idx, bbox in enumerate(bboxes):
# half_h = bbox[3] / 2
# half_w = bbox[2] / 2
#
# if np.sum(bbox) > 0:
# # print("class:{}".format(class_inds))
# cv2.rectangle(check, (int(bbox[0] - half_w), int(bbox[1] - half_h)),
# (int(bbox[0] + half_w), int(bbox[1] + half_h)), color=(0, 255, 0), thickness=3)
# cv2.putText(check, text="{}".format(class_inds[bbox_idx]),
# org=(int(bbox[0] + half_w)-10, int(bbox[1] + half_h)-30), thickness=2, color=(0, 255, 0),
# fontFace=cv2.FONT_HERSHEY_SIMPLEX, fontScale=0.7)
#
# cv2.imshow("check_aug", check)
# cv2.waitKey(0)
# cv2.destroyAllWindows()
# cv2.imwrite(os.path.join(r"D:\Public\JHS\SIMPLE_DATA_INPUT_CHECK", "{}.jpg".format(batch_idx)), check*255)
train_step(image_data, target)
# for image_data, target in testset:
# test_step(image_data, target)
### for loss graph
"""
if steps_in_epoch >= steps_per_epoch:
loss_tracker.append(losses_in_epoch / steps_per_epoch)
plt.plot(loss_tracker)
plt.xlabel('epoch')
plt.ylabel('loss')
# plt.show()
plt.savefig("D:/checkpoint/loss.png")
"""
# print("steps_in_epoch:{}, steps_per_epoch:{}".format(global_steps, steps_per_epoch))
# print("prev_minloss:{}, epoch_loss:{}".format(prev_minloss, epoch_loss))
# # early stopping
# if len(losses_in_epoch) >= steps_per_epoch:
#
# epoch_loss = losses_in_epoch / steps_per_epoch
# if prev_minloss > epoch_loss: # save best weight
# prev_minloss = epoch_loss
#
# if epoch > 800: # 최소학습 에폭
# model.save("D:\ckpt_best")
# print("{} epoch save best weights".format(epoch))
#
# if len(loss_tracker) > patience:
# print("check loss_tracker len:{}".format(len(loss_tracker)))
# if loss_tracker[0] > np.min(loss_tracker[1:]):
# loss_tracker.pop(0)
# else:
# print("total loss didn't decreased during {} epochs. train stop".format(patience))
# return
# steps_in_epoch = 0
# epoch_loss = train_loss
# else:
# epoch_loss += train_loss
if (epoch+1) % 500 == 0:
model.save(r"D:\notf_ckpt-epoch{}".format(epoch))
print("{} epoch model saved".format(epoch))
model.save(r"D:\notf_ckpt-last")
def main():
strategy = tf.distribute.MirroredStrategy()
num_gpus = strategy.num_replicas_in_sync
with strategy.scope():
train_batch_size = cfg.TRAIN.BATCH_SIZE
val_batch_size = cfg.VAL.BATCH_SIZE
trainset = Dataset('train')
valset = Dataset('val')
train_generator = tf.data.Dataset.from_generator(
lambda: trainset,
(tf.float32, ((tf.float32, tf.float32), (tf.float32, tf.float32),
(tf.float32, tf.float32)))).batch(train_batch_size)
val_generator = tf.data.Dataset.from_generator(
lambda: valset,
(tf.float32, ((tf.float32, tf.float32), (tf.float32, tf.float32),
(tf.float32, tf.float32)))).batch(val_batch_size)
dist_trainset = strategy.experimental_distribute_dataset(
train_generator)
dist_valset = strategy.experimental_distribute_dataset(val_generator)
isfreeze = False
steps_per_epoch = math.ceil(len(trainset) / train_batch_size)
val_steps_per_epoch = math.ceil(len(valset) / val_batch_size)
first_stage_epochs = cfg.TRAIN.FISRT_STAGE_EPOCHS
second_stage_epochs = cfg.TRAIN.SECOND_STAGE_EPOCHS
global_steps = tf.Variable(1, trainable=False, dtype=tf.int64)
val_global_steps = tf.Variable(1, trainable=False, dtype=tf.int64)
warmup_steps = cfg.TRAIN.WARMUP_EPOCHS * steps_per_epoch
total_steps = (first_stage_epochs +
second_stage_epochs) * steps_per_epoch
NUM_CLASS = len(utils.read_class_names(cfg.YOLO.CLASSES))
STRIDES = np.array(cfg.YOLO.STRIDES)
IOU_LOSS_THRESH = cfg.YOLO.IOU_LOSS_THRESH
XYSCALE = cfg.YOLO.XYSCALE
ANCHORS = utils.get_anchors(cfg.YOLO.ANCHORS)
model = YOLOv4(NUM_CLASS, STRIDES, ANCHORS, XYSCALE)
optimizer = tf.keras.optimizers.Adam()
ckpt = tf.train.Checkpoint(optimizer=optimizer, model=model)
ckpt_manager = tf.train.CheckpointManager(ckpt,
args.logdir,
max_to_keep=3)
writer = tf.summary.create_file_writer(args.logdir)
start_iteration = int(ckpt_manager.latest_checkpoint.split('-')
[-1]) if args.weights else 0
if args.weights:
print('Restoring weights from: %s ... ' %
ckpt_manager.latest_checkpoint)
ckpt.restore(ckpt_manager.latest_checkpoint)
else:
print('Training from scratch')
print('Start iteration:', start_iteration)
@tf.function
def train_step(image_data, target):
with tf.GradientTape() as tape:
pred_result = model(image_data, training=True)
giou_loss = conf_loss = prob_loss = 0
# optimizing process
for i in range(3):
conv, pred = pred_result[i * 2], pred_result[i * 2 + 1]
loss_items = utils.compute_loss(
pred,
conv,
target[i][0],
target[i][1],
STRIDES=STRIDES,
NUM_CLASS=NUM_CLASS,
IOU_LOSS_THRESH=IOU_LOSS_THRESH,
i=i)
giou_loss += loss_items[0]
conf_loss += loss_items[1]
prob_loss += loss_items[2]
total_loss = giou_loss + conf_loss + prob_loss
giou_loss = tf.reduce_sum(giou_loss) * (1.0 / train_batch_size)
conf_loss = tf.reduce_sum(conf_loss) * (1.0 / train_batch_size)
prob_loss = tf.reduce_sum(prob_loss) * (1.0 / train_batch_size)
total_loss = tf.reduce_sum(total_loss) * (1.0 / train_batch_size)
gradients = tape.gradient(total_loss, model.trainable_variables)
optimizer.apply_gradients(zip(gradients,
model.trainable_variables))
return total_loss, giou_loss, conf_loss, prob_loss
def distribute_train_step(image_data, target):
with strategy.scope():
tloss_temp, gloss_temp, closs_temp, ploss_temp = strategy.experimental_run_v2(
train_step, args=(image_data, target))
tloss = strategy.reduce(tf.distribute.ReduceOp.MEAN,
tloss_temp,
axis=None)
gloss = strategy.reduce(tf.distribute.ReduceOp.MEAN,
gloss_temp,
axis=None)
closs = strategy.reduce(tf.distribute.ReduceOp.MEAN,
closs_temp,
axis=None)
ploss = strategy.reduce(tf.distribute.ReduceOp.MEAN,
ploss_temp,
axis=None)
return tloss, gloss, closs, ploss
@tf.function
def val_step(image_data, target):
with tf.GradientTape() as tape:
pred_result = model(image_data, training=True)
giou_loss = conf_loss = prob_loss = 0
# optimizing process
for i in range(3):
conv, pred = pred_result[i * 2], pred_result[i * 2 + 1]
loss_items = utils.compute_loss(
pred,
conv,
target[i][0],
target[i][1],
STRIDES=STRIDES,
NUM_CLASS=NUM_CLASS,
IOU_LOSS_THRESH=IOU_LOSS_THRESH,
i=i)
giou_loss += loss_items[0]
conf_loss += loss_items[1]
prob_loss += loss_items[2]
total_loss = giou_loss + conf_loss + prob_loss
giou_loss = tf.reduce_sum(giou_loss) * (1.0 / val_batch_size)
conf_loss = tf.reduce_sum(conf_loss) * (1.0 / val_batch_size)
prob_loss = tf.reduce_sum(prob_loss) * (1.0 / val_batch_size)
total_loss = tf.reduce_sum(total_loss) * (1.0 / val_batch_size)
return total_loss, giou_loss, conf_loss, prob_loss
def distribute_val_step(image_data, target):
with strategy.scope():
tloss_temp, gloss_temp, closs_temp, ploss_temp = strategy.experimental_run_v2(
val_step, args=(image_data, target))
tloss = strategy.reduce(tf.distribute.ReduceOp.MEAN,
tloss_temp,
axis=None)
gloss = strategy.reduce(tf.distribute.ReduceOp.MEAN,
gloss_temp,
axis=None)
closs = strategy.reduce(tf.distribute.ReduceOp.MEAN,
closs_temp,
axis=None)
ploss = strategy.reduce(tf.distribute.ReduceOp.MEAN,
ploss_temp,
axis=None)
return tloss, gloss, closs, ploss
def train(epoch, max_epoch):
for image_data, target in dist_trainset:
total_loss, giou_loss, conf_loss, prob_loss = distribute_train_step(
image_data, target)
tf.print(
"=> EPOCH [%3d/%3d] ITER [%5d/%5d] lr: %.6f giou_loss: %4.2f conf_loss: %4.2f "
"prob_loss: %4.2f total_loss: %4.2f" %
(epoch, max_epoch, global_steps -
(epoch - 1) * steps_per_epoch, steps_per_epoch,
optimizer.lr.numpy(), giou_loss, conf_loss, prob_loss,
total_loss))
# writing summary data
with writer.as_default():
tf.summary.scalar("lr", optimizer.lr, step=global_steps)
tf.summary.scalar("loss/total_loss",
total_loss,
step=global_steps)
tf.summary.scalar("loss/giou_loss",
giou_loss,
step=global_steps)
tf.summary.scalar("loss/conf_loss",
conf_loss,
step=global_steps)
tf.summary.scalar("loss/prob_loss",
prob_loss,
step=global_steps)
writer.flush()
# update learning rate
global_steps.assign_add(1)
if global_steps < warmup_steps:
lr = global_steps / warmup_steps * cfg.TRAIN.LR_INIT
else:
lr = cfg.TRAIN.LR_END + 0.5 * (
cfg.TRAIN.LR_INIT - cfg.TRAIN.LR_END) * ((1 + tf.cos(
(global_steps - warmup_steps) /
(total_steps - warmup_steps) * np.pi)))
optimizer.lr.assign(lr.numpy())
def val(epoch, max_epoch):
for image_data, target in dist_valset:
total_loss, giou_loss, conf_loss, prob_loss = distribute_val_step(
image_data, target)
tf.print(
"=> VALIDATION EPOCH [%3d/%3d] ITER [%5d/%5d] giou_loss: %4.2f conf_loss: %4.2f "
"prob_loss: %4.2f total_loss: %4.2f" %
(epoch, max_epoch, val_global_steps -
(epoch - 1) * val_steps_per_epoch, val_steps_per_epoch,
giou_loss, conf_loss, prob_loss, total_loss))
# writing summary data
with writer.as_default():
tf.summary.scalar("val_loss/total_loss",
total_loss,
step=val_global_steps)
tf.summary.scalar("val_loss/giou_loss",
giou_loss,
step=val_global_steps)
tf.summary.scalar("val_loss/conf_loss",
conf_loss,
step=val_global_steps)
tf.summary.scalar("val_loss/prob_loss",
prob_loss,
step=val_global_steps)
writer.flush()
val_global_steps.assign_add(1)
for epoch in range(first_stage_epochs + second_stage_epochs):
if epoch < first_stage_epochs:
if not isfreeze:
isfreeze = True
for name in ['conv_93', 'conv_101', 'conv_109']:
freeze = model.get_layer(name)
utils.freeze_all(freeze)
elif epoch >= first_stage_epochs:
if isfreeze:
isfreeze = False
for name in ['conv_93', 'conv_101', 'conv_109']:
freeze = model.get_layer(name)
utils.unfreeze_all(freeze)
train(epoch + 1, first_stage_epochs + second_stage_epochs)
val(epoch + 1, first_stage_epochs + second_stage_epochs)
ckpt_manager.save(checkpoint_number=epoch)
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)
trainset = Dataset('train')
testset = Dataset('test')
logdir = "./data/log"
isfreeze = False
steps_per_epoch = len(trainset)
first_stage_epochs = cfg.TRAIN.FISRT_STAGE_EPOCHS
second_stage_epochs = cfg.TRAIN.SECOND_STAGE_EPOCHS
#global_steps =
warmup_steps = cfg.TRAIN.WARMUP_EPOCHS * steps_per_epoch
total_steps = (first_stage_epochs + second_stage_epochs) * steps_per_epoch
# train_steps = (first_stage_epochs + second_stage_epochs) * steps_per_period
NUM_CLASS = len(utils.read_class_names(cfg.YOLO.CLASSES))
STRIDES = np.array(cfg.YOLO.STRIDES)
IOU_LOSS_THRESH = cfg.YOLO.IOU_LOSS_THRESH
XYSCALE = cfg.YOLO.XYSCALE
ANCHORS = utils.get_anchors(cfg.YOLO.ANCHORS)
input_layer = tf.keras.layers.Input([cfg.TRAIN.INPUT_SIZE, cfg.TRAIN.INPUT_SIZE, 3])
#YOLOV4
feature_maps = YOLOv4(input_layer, NUM_CLASS)
bbox_tensors = []
for i, fm in enumerate(feature_maps):
bbox_tensor = decode_train(fm, NUM_CLASS, STRIDES, ANCHORS, i)
bbox_tensors.append(fm)
bbox_tensors.append(bbox_tensor)
model = tf.keras.Model(input_layer, bbox_tensors)
if FLAGS.weights == "none":
print("Training from scratch")
else:
if FLAGS.weights.split(".")[len(FLAGS.weights.split(".")) - 1] == "weights":
if FLAGS.tiny:
utils.load_weights_tiny(model, FLAGS.weights)
else:
if FLAGS.model == 'yolov3':
utils.load_weights_v3(model, FLAGS.weights)
else:
utils.load_weights(model, FLAGS.weights)
else:
model.load_weights(FLAGS.weights)
print('Restoring weights from: %s ... ' % FLAGS.weights)
if os.path.exists(logdir):
shutil.rmtree(logdir)
optimizer = tf.keras.optimizers.Adam()
writer = tf.summary.create_file_writer(logdir)
ckpt = tf.train.Checkpoint(step=tf.Variable(1, trainable=False, dtype=tf.int64), optimizer=optimizer, net=model)
manager = tf.train.CheckpointManager(ckpt, './checkpoints', max_to_keep=10)
def train_step(image_data, target):
with tf.GradientTape() as tape:
pred_result = model(image_data, training=True)
giou_loss = conf_loss = prob_loss = 0
# optimizing process
for i in range(3):
conv, pred = pred_result[i * 2], pred_result[i * 2 + 1]
loss_items = compute_loss(pred, conv, target[i][0], target[i][1], STRIDES=STRIDES, NUM_CLASS=NUM_CLASS, IOU_LOSS_THRESH=IOU_LOSS_THRESH, i=i)
giou_loss += loss_items[0]
conf_loss += loss_items[1]
prob_loss += loss_items[2]
total_loss = giou_loss + conf_loss + prob_loss
gradients = tape.gradient(total_loss, model.trainable_variables)
optimizer.apply_gradients(zip(gradients, model.trainable_variables))
tf.print("=> STEP %4d lr: %.6f giou_loss: %4.2f conf_loss: %4.2f "
"prob_loss: %4.2f total_loss: %4.2f" % (ckpt.step, optimizer.lr.numpy(),
giou_loss, conf_loss,
prob_loss, total_loss))
# update learning rate
ckpt.step.assign_add(1)
if ckpt.step < warmup_steps:
lr = ckpt.step / warmup_steps * cfg.TRAIN.LR_INIT
else:
lr = cfg.TRAIN.LR_END + 0.5 * (cfg.TRAIN.LR_INIT - cfg.TRAIN.LR_END) * (
(1 + tf.cos((ckpt.step - warmup_steps) / (total_steps - warmup_steps) * np.pi))
)
optimizer.lr.assign(lr.numpy())
# writing summary data
with writer.as_default():
tf.summary.scalar("lr", optimizer.lr, step=ckpt.step)
tf.summary.scalar("loss/total_loss", total_loss, step=ckpt.step)
tf.summary.scalar("loss/giou_loss", giou_loss, step=ckpt.step)
tf.summary.scalar("loss/conf_loss", conf_loss, step=ckpt.step)
tf.summary.scalar("loss/prob_loss", prob_loss, step=ckpt.step,)
writer.flush()
def test_step(image_data, target):
with tf.GradientTape() as tape:
pred_result = model(image_data, training=True)
giou_loss = conf_loss = prob_loss = 0
# optimizing process
for i in range(3):
conv, pred = pred_result[i * 2], pred_result[i * 2 + 1]
loss_items = compute_loss(pred, conv, target[i][0], target[i][1], STRIDES=STRIDES, NUM_CLASS=NUM_CLASS, IOU_LOSS_THRESH=IOU_LOSS_THRESH, i=i)
giou_loss += loss_items[0]
conf_loss += loss_items[1]
prob_loss += loss_items[2]
total_loss = giou_loss + conf_loss + prob_loss
tf.print("=> TEST STEP %4d giou_loss: %4.2f conf_loss: %4.2f "
"prob_loss: %4.2f total_loss: %4.2f" % (ckpt.step, giou_loss, conf_loss,
prob_loss, total_loss))
ckpt.restore(manager.latest_checkpoint)
if manager.latest_checkpoint:
print("Restored from {}".format(manager.latest_checkpoint))
else:
print("Initializing from scratch.")
for epoch in range(first_stage_epochs + second_stage_epochs):
ckpt.step.assign_add(1)
print("----------- EPOCH :: ", str(epoch), " :: -----------")
if epoch < first_stage_epochs:
print("epoch < first_stage_epochs")
if not isfreeze:
isfreeze = True
for name in ['conv2d_93', 'conv2d_101', 'conv2d_109']:
freeze = model.get_layer(name)
freeze_all(freeze)
elif epoch >= first_stage_epochs:
print("epoch < first_stage_epochs")
if isfreeze:
isfreeze = False
for name in ['conv2d_93', 'conv2d_101', 'conv2d_109']:
freeze = model.get_layer(name)
unfreeze_all(freeze)
for image_data, target in trainset:
train_step(image_data, target)
if int(ckpt.step) % 10 == 0:
save_path = manager.save()
print("Saved checkpoint for step {}: {}".format(int(ckpt.step), save_path))
def main(_argv):
gpus = tf.config.experimental.list_physical_devices('GPU')
if gpus:
# 텐서플로가 첫 번째 GPU만 사용하도록 제한
try:
tf.config.experimental.set_visible_devices(gpus[7], 'GPU')
except RuntimeError as e:
# 프로그램 시작시에 접근 가능한 장치가 설정되어야만 합니다
print(e)
trainset = Dataset(FLAGS, is_training=True)
#print(next(iter(trainset)))
#_train = next(iter(trainset)) #_train[0] = (3, 416, 416, 3), len(train[1] =3,
#len(train[1][0])) =2,
# _train[1][0][0].shape : (3, 26, 26, 3, 25)
# _train[1][0][1].shape : (3, 150, 4)
# len(_train[1][1]) =2,
# _train[1][1][0].shape : (3, 26, 26, 3, 25)
# _train[1][1][1].shape : (3, 150, 4)
testset = Dataset(FLAGS, is_training=False)
logdir = "./data/log"
isfreeze = False
steps_per_epoch = len(trainset)
first_stage_epochs = cfg.TRAIN.FISRT_STAGE_EPOCHS
second_stage_epochs = cfg.TRAIN.SECOND_STAGE_EPOCHS
global_steps = tf.Variable(1, trainable=False, dtype=tf.int64)
warmup_steps = cfg.TRAIN.WARMUP_EPOCHS * steps_per_epoch
total_steps = (first_stage_epochs + second_stage_epochs) * steps_per_epoch
# train_steps = (first_stage_epochs + second_stage_epochs) * steps_per_period
input_layer = tf.keras.layers.Input([cfg.TRAIN.INPUT_SIZE, cfg.TRAIN.INPUT_SIZE, 3])
STRIDES, ANCHORS, NUM_CLASS, XYSCALE = utils.load_config(FLAGS)
IOU_LOSS_THRESH = cfg.YOLO.IOU_LOSS_THRESH
freeze_layers = utils.load_freeze_layer(FLAGS.model, FLAGS.tiny)
#yolov4 : ['conv2d_93', 'conv2d_101', 'conv2d_109']
#yolov4_tiny: ['conv2d_17', 'conv2d_20']
feature_maps = YOLO(input_layer, NUM_CLASS, FLAGS.model, FLAGS.tiny)
# '''
# import core.backbone as backbone
# route_1, route_2, conv = backbone.cspdarknet53(input_layer)
# route_1
# Out[97]: <tf.Tensor 'Mul_355:0' shape=(None, 52, 52, 256) dtype=float32>
# route_2
# Out[98]: <tf.Tensor 'Mul_376:0' shape=(None, 26, 26, 512) dtype=float32>
# conv
# Out[99]: <tf.Tensor 'LeakyRelu_164:0' shape=(None, 13, 13, 512) dtype=float32>
# '''
#yolov4 : [<tf.Tensor 'conv2d_93/BiasAdd:0' shape=(None, 52, 52, 75) dtype=float32>,
#<tf.Tensor 'conv2d_101/BiasAdd:0' shape=(None, 26, 26, 75) dtype=float32>,
#<tf.Tensor 'conv2d_109/BiasAdd:0' shape=(None, 13, 13, 75) dtype=float32>]
#yolov4_tiny : [<tf.Tensor 'conv2d_130/BiasAdd:0' shape=(None, 26, 26, 75) dtype=float32>,
#<tf.Tensor 'conv2d_127/BiasAdd:0' shape=(None, 13, 13, 75) dtype=float32>]
if FLAGS.tiny:
bbox_tensors = []
for i, fm in enumerate(feature_maps):
if i == 0:
bbox_tensor = decode_train(fm, cfg.TRAIN.INPUT_SIZE // 16, NUM_CLASS, STRIDES, ANCHORS, i, XYSCALE)
else:
bbox_tensor = decode_train(fm, cfg.TRAIN.INPUT_SIZE // 32, NUM_CLASS, STRIDES, ANCHORS, i, XYSCALE)
bbox_tensors.append(fm)
bbox_tensors.append(bbox_tensor)
else:
bbox_tensors = []
for i, fm in enumerate(feature_maps):
if i == 0:
bbox_tensor = decode_train(fm, cfg.TRAIN.INPUT_SIZE // 8, NUM_CLASS, STRIDES, ANCHORS, i, XYSCALE)
elif i == 1:
bbox_tensor = decode_train(fm, cfg.TRAIN.INPUT_SIZE // 16, NUM_CLASS, STRIDES, ANCHORS, i, XYSCALE)
else:
bbox_tensor = decode_train(fm, cfg.TRAIN.INPUT_SIZE // 32, NUM_CLASS, STRIDES, ANCHORS, i, XYSCALE)
bbox_tensors.append(fm)
bbox_tensors.append(bbox_tensor)
model = tf.keras.Model(input_layer, bbox_tensors)
model.summary()
# from contextlib import redirect_stdout
# with open('modelsummary.txt', 'w') as f:
# with redirect_stdout(f):
# model.summary()
if FLAGS.weights == None:
print("Training from scratch")
else:
if FLAGS.weights.split(".")[len(FLAGS.weights.split(".")) - 1] == "weights":
utils.load_weights(model, FLAGS.weights, FLAGS.model, FLAGS.tiny)
else:
model.load_weights(FLAGS.weights)
print('Restoring weights from: %s ... ' % FLAGS.weights)
# wf = open(FLAGS.weights, 'rb')
# major, minor, revision, seen, _ = np.fromfile(wf, dtype=np.int32, count=5)
# wf.close()
# layer_size = 110
# output_pos = [93, 101, 109]
# j = 0
# for i in range(layer_size):
# conv_layer_name = 'conv2d_%d' %i if i > 0 else 'conv2d'
# bn_layer_name = 'batch_normalization_%d' %j if j > 0 else 'batch_normalization'
# conv_layer = model.get_layer(conv_layer_name)
# filters = conv_layer.filters
# k_size = conv_layer.kernel_size[0]
# in_dim = conv_layer.input_shape[-1]
optimizer = tf.keras.optimizers.Adam()
if os.path.exists(logdir): shutil.rmtree(logdir)
writer = tf.summary.create_file_writer(logdir)
# define training step function
# @tf.function
def train_step(image_data, target):
with tf.GradientTape() as tape:
pred_result = model(image_data, training=True)
giou_loss = conf_loss = prob_loss = 0
# optimizing process
for i in range(len(freeze_layers)):
conv, pred = pred_result[i * 2], pred_result[i * 2 + 1]
# print(pred)
loss_items = compute_loss(pred, conv, target[i][0], target[i][1], STRIDES=STRIDES, NUM_CLASS=NUM_CLASS, IOU_LOSS_THRESH=IOU_LOSS_THRESH, i=i)
giou_loss += loss_items[0]
conf_loss += loss_items[1]
prob_loss += loss_items[2]
total_loss = giou_loss + conf_loss + prob_loss
gradients = tape.gradient(total_loss, model.trainable_variables)
optimizer.apply_gradients(zip(gradients, model.trainable_variables))
tf.print("=> STEP %4d/%4d lr: %.6f giou_loss: %4.2f conf_loss: %4.2f "
"prob_loss: %4.2f total_loss: %4.2f" % (global_steps, total_steps, optimizer.lr.numpy(),
giou_loss, conf_loss,
prob_loss, total_loss))
# update learning rate
global_steps.assign_add(1)
if global_steps < warmup_steps:
lr = global_steps / warmup_steps * cfg.TRAIN.LR_INIT
else:
lr = cfg.TRAIN.LR_END + 0.5 * (cfg.TRAIN.LR_INIT - cfg.TRAIN.LR_END) * (
(1 + tf.cos((global_steps - warmup_steps) / (total_steps - warmup_steps) * np.pi))
)
optimizer.lr.assign(lr.numpy())
# writing summary data
with writer.as_default():
tf.summary.scalar("lr", optimizer.lr, step=global_steps)
tf.summary.scalar("loss/total_loss", total_loss, step=global_steps)
tf.summary.scalar("loss/giou_loss", giou_loss, step=global_steps)
tf.summary.scalar("loss/conf_loss", conf_loss, step=global_steps)
tf.summary.scalar("loss/prob_loss", prob_loss, step=global_steps)
writer.flush()
def test_step(image_data, target):
with tf.GradientTape() as tape:
pred_result = model(image_data, training=True)
giou_loss = conf_loss = prob_loss = 0
# optimizing process
for i in range(len(freeze_layers)):
conv, pred = pred_result[i * 2], pred_result[i * 2 + 1]
loss_items = compute_loss(pred, conv, target[i][0], target[i][1], STRIDES=STRIDES, NUM_CLASS=NUM_CLASS, IOU_LOSS_THRESH=IOU_LOSS_THRESH, i=i)
giou_loss += loss_items[0]
conf_loss += loss_items[1]
prob_loss += loss_items[2]
total_loss = giou_loss + conf_loss + prob_loss
tf.print("=> TEST STEP %4d giou_loss: %4.2f conf_loss: %4.2f "
"prob_loss: %4.2f total_loss: %4.2f" % (global_steps, giou_loss, conf_loss,
prob_loss, total_loss))
for epoch in range(first_stage_epochs + second_stage_epochs):
if epoch < first_stage_epochs:
if not isfreeze:
isfreeze = True
for name in freeze_layers:
freeze = model.get_layer(name)
freeze_all(freeze)
elif epoch >= first_stage_epochs:
if isfreeze:
isfreeze = False
for name in freeze_layers:
freeze = model.get_layer(name)
unfreeze_all(freeze)
for image_data, target in trainset:
train_step(image_data, target)
#for image_data, target in testset:
# test_step(image_data, target)
model.save_weights("./checkpoints/yolov4")
def main():
INPUT_SIZE = FLAGS.size
STRIDES, ANCHORS, NUM_CLASS, XYSCALE = utils.load_config(FLAGS)
CLASSES = utils.read_class_names(cfg.YOLO.CLASSES)
IOU_LOSS_THRESH = cfg.YOLO.IOU_LOSS_THRESH
predicted_dir_path = './AP/predicted'
loss_dir_path = './AP/loss'
text_result_path = './AP/detect'
trainset = Dataset(FLAGS, is_training=True)
adv_lambda = 0.001
steps_per_epoch = len(trainset)
first_stage_epochs = cfg.TRAIN.FISRT_STAGE_EPOCHS
second_stage_epochs = cfg.TRAIN.SECOND_STAGE_EPOCHS
global_steps = tf.Variable(1, trainable=False, dtype=tf.int64)
warmup_steps = cfg.TRAIN.WARMUP_EPOCHS * steps_per_epoch
total_steps = (first_stage_epochs + second_stage_epochs) * steps_per_epoch
input_layer = tf.keras.layers.Input([cfg.TRAIN.INPUT_SIZE, cfg.TRAIN.INPUT_SIZE, 3])
feature_maps = YOLO(FLAGS.scale_v5, input_layer, NUM_CLASS)
bbox_tensors = []
for i, fm in enumerate(feature_maps):
if i == 0:
bbox_tensor = decode_train(fm, cfg.TRAIN.INPUT_SIZE // 8, NUM_CLASS, STRIDES, ANCHORS, i, XYSCALE)
bbox_tensors.append(fm)
bbox_tensors.append(bbox_tensor)
elif i == 1:
bbox_tensor = decode_train(fm, cfg.TRAIN.INPUT_SIZE // 16, NUM_CLASS, STRIDES, ANCHORS, i, XYSCALE)
bbox_tensors.append(fm)
bbox_tensors.append(bbox_tensor)
elif i==2:
bbox_tensor = decode_train(fm, cfg.TRAIN.INPUT_SIZE // 32, NUM_CLASS, STRIDES, ANCHORS, i, XYSCALE)
bbox_tensors.append(fm)
bbox_tensors.append(bbox_tensor)
else:
result_G = fm
bbox_tensors.append(result_G)
D2_model = tf.keras.Model(input_layer, bbox_tensors)
full_model = NLayerDiscriminator(ndf=64, n_layers=5)
full_model.build([None, cfg.TRAIN.INPUT_SIZE, cfg.TRAIN.INPUT_SIZE, 3])
optimizer_D2 = tf.keras.optimizers.Adam()
optimizer_Dis = tf.keras.optimizers.SGD()
criterionG, criterionD = get_loss()
if cfg.TRAIN.DoubleGAN:
patch_model = NLayerDiscriminator(ndf=64, n_layers=3)
patch_model.build([None, cfg.TRAIN.INPUT_SIZE, cfg.TRAIN.INPUT_SIZE, 3])
adv_trainer = DoubleGAN(patch_model, full_model, criterionD)
else:
adv_trainer = SingleGAN(full_model, criterionD)
if cfg.TRAIN.GRADNORM:
optimizer_P = tf.keras.optimizers.SGD()
D2_model.summary()
T_freeze_layers = utils.load_True_freeze_layer(FLAGS.scale_v5)
# metrics draw now
total_loss_metric = tf.metrics.Mean()
loss1_metric = tf.metrics.Mean()
loss2_metric = tf.metrics.Mean()
total_loss_result= []
loss1_result= []
loss2_result= []
Weightloss1 = tf.Variable(1.0)
Weightloss2 = tf.Variable(1.0)
params = [Weightloss1, Weightloss2]
alph = 0.16
## @tf.function
def train_step(image_data, target):
start_time = time.time()
# with tf.GradientTape() as tape1,tf.GradientTape() as tape2:
##Experiments have found that this performance is better than "persistent=True".
with tf.GradientTape() as tape1,tf.GradientTape() as tape2,tf.GradientTape() as tape3,tf.GradientTape() as tape4,tf.GradientTape() as tape5:
pred_result = D2_model(image_data[0], training=True)
G_im = pred_result[-1]
loss_D = loss_content = loss_adv = loss_G = giou_loss = conf_loss = prob_loss = 0
#update Discriminator
loss_D = 1000 * adv_lambda * adv_trainer.loss_d(G_im, image_data[1])
gradients_Dis = tape1.gradient(loss_D, adv_trainer.get_params())
optimizer_Dis.apply_gradients(zip(gradients_Dis, adv_trainer.get_params()))
#update D2Net
loss_content = criterionG(G_im, image_data[1])
loss_adv = adv_trainer.loss_g(G_im, image_data[1])
loss_G = 1000*(loss_content + adv_lambda * loss_adv)
for i in range(ANCHORS.shape[0]):
conv, pred = pred_result[i * 2], pred_result[i * 2 + 1]
loss_items = compute_loss(pred, conv, target[i][0], target[i][1], STRIDES=STRIDES, NUM_CLASS=NUM_CLASS, IOU_LOSS_THRESH=IOU_LOSS_THRESH, i=i)
giou_loss += loss_items[0]
conf_loss += loss_items[1]
prob_loss += loss_items[2]
yolo_loss = giou_loss + conf_loss + prob_loss
l1 = params[0]*yolo_loss
l2 = params[1]*loss_G
total_loss = (l1 + l2)/2
gradients_D2 = tape2.gradient(total_loss, D2_model.trainable_variables)
optimizer_D2.apply_gradients(zip(gradients_D2, D2_model.trainable_variables))
###Gradnorm###
L0 = 183
LP = D2_model.trainable_variables[162] #D_conv2d_54
G1R = tape3.gradient(l1, LP)
G1 = tf.norm(G1R, ord=2)
G2R = tape4.gradient(l2, LP)
G2 = tf.norm(G2R, ord=2)
G_avg = (G1+G2)/2
# Calculating relative losses
lhat1 = (l1)/L0
lhat2 = (l2)/L0
lhat_avg = (lhat1 + lhat2)/2
inv_rate1 = lhat1/lhat_avg
inv_rate2 = lhat2/lhat_avg
C1 = G_avg*(inv_rate1)**alph
C2 = G_avg*(inv_rate2)**alph
C1 = tf.stop_gradient(tf.identity(C1))
C2 = tf.stop_gradient(tf.identity(C2))
# Gradnorm loss
loss_gradnorm = tf.math.reduce_sum(tf.math.abs(G1-C1)) + tf.math.reduce_sum(tf.math.abs(G2-C2))
grad_grad = tape5.gradient(loss_gradnorm, params)
optimizer_P.apply_gradients(grads_and_vars=zip(grad_grad, params))
total_loss_metric.update_state(values=total_loss)
loss1_metric.update_state(values=yolo_loss)
loss2_metric.update_state(values=loss_G)
time_per_step = time.time() - start_time
print("Step: {}/{},lr: {:.6f}, {:.2f}s/step, total_loss: {:.5f}, "
"yolo loss: {:.5f}, G_loss: {:.5f}, loss_adv: {:.5f}, D_loss: {:.5f}".format(
global_steps.numpy(),
total_steps,
optimizer_D2.lr.numpy(),
time_per_step,
total_loss,
yolo_loss,
loss_G,
adv_lambda*loss_adv,
loss_D
))
# update learning rate
global_steps.assign_add(1)
if global_steps < warmup_steps:
lr = global_steps / warmup_steps * cfg.TRAIN.LR_INIT
else:
lr = cfg.TRAIN.LR_END + 0.5 * (cfg.TRAIN.LR_INIT - cfg.TRAIN.LR_END) * (
(1 + tf.cos((global_steps - warmup_steps) / (total_steps - warmup_steps) * np.pi))
)
optimizer_D2.lr.assign(lr.numpy())
optimizer_Dis.lr.assign(10*lr.numpy())
loss_list_step = [optimizer_D2.lr.numpy(),total_loss,yolo_loss,
loss_G,loss_D,giou_loss,conf_loss, prob_loss,loss_content,adv_lambda * loss_adv]
return np.array(loss_list_step)
## @tf.function
def test_epoch(D2_model,dectect_epoch_path):
with open(cfg.TEST.ANNOT_PATH, 'r') as annotation_file:
for num, line in enumerate(annotation_file):
annotation = line.strip().split()
image_path = annotation[0]
image_name = image_path.split('/')[-1]
predict_result_path = os.path.join(predicted_epoch_path, str(image_name) + '.txt')
original_image = cv2.imread(image_path)
image = cv2.cvtColor(original_image, cv2.COLOR_BGR2RGB)
# Predict Process
## image_letter, ratio, (dw, dh) = utils.letterbox(image)
image_letter = utils.test_image_preprocess(np.copy(image), [INPUT_SIZE, INPUT_SIZE])
image_data = image_letter[np.newaxis, ...].astype(np.float32)
batch_data = tf.constant(image_data)
bbox_tensors = []
prob_tensors = []
pred_result = D2_model(batch_data,training=False)
G_im = pred_result[-1][0]
for i in range(ANCHORS.shape[0]):
fm = pred_result[i * 2]
if i == 0:
output_tensors = decode(fm, FLAGS.size // 8, NUM_CLASS, STRIDES, ANCHORS, i, XYSCALE)
elif i == 1:
output_tensors = decode(fm, FLAGS.size // 16, NUM_CLASS, STRIDES, ANCHORS, 1, XYSCALE)
elif i==2:
output_tensors = decode(fm, FLAGS.size // 32, NUM_CLASS, STRIDES, ANCHORS, 2, XYSCALE)
bbox_tensors.append(output_tensors[0])
prob_tensors.append(output_tensors[1])
pred_bbox = tf.concat(bbox_tensors, axis=1)
pred_prob = tf.concat(prob_tensors, axis=1)
boxes, pred_conf = filter_boxes(pred_bbox, pred_prob, score_threshold=FLAGS.score_thres, input_shape=tf.constant([FLAGS.size, FLAGS.size]))
pred_bbox = tf.concat([boxes, pred_conf], axis=-1)
boxes = pred_bbox[:, :, 0:4]
pred_conf = pred_bbox[:, :, 4:]
boxes, scores, classes, valid_detections = tf.image.combined_non_max_suppression(
boxes=tf.reshape(boxes, (tf.shape(boxes)[0], -1, 1, 4)),
scores=tf.reshape(
pred_conf, (tf.shape(pred_conf)[0], -1, tf.shape(pred_conf)[-1])),
max_output_size_per_class=1,
max_total_size=1,
iou_threshold=FLAGS.iou,
score_threshold=FLAGS.score
)
boxes, scores, classes, valid_detections = [boxes.numpy(), scores.numpy(), classes.numpy(), valid_detections.numpy()]
if num % 1 ==0:
G_im = pred_result[-1][0]
G_im = G_im * 255
G_im = np.array(G_im).astype(np.int32)
image_result = utils.draw_bbox(np.copy(G_im), [boxes, scores, classes, valid_detections])
image_result = image_result[:,:,::-1]
filepath = dectect_epoch_path+"/"+ str(image_name)
cv2.imwrite(filepath, image_result, [int(cv2.IMWRITE_JPEG_QUALITY), 100])
################################################################
################################################################
################################################################
if os.path.exists(loss_dir_path): shutil.rmtree(loss_dir_path)
os.mkdir(loss_dir_path)
for epoch in range(first_stage_epochs + second_stage_epochs):
if epoch >= first_stage_epochs:
for name in T_freeze_layers:
try:
freeze = D2_model.get_layer(name)
freeze_all(freeze)
print("Successfully freeze {}...".format(name))
except:
print("{} not exist...".format(name))
loss_epoch = np.zeros((steps_per_epoch,10),dtype=np.float32)
for index, (image_data, target) in enumerate(trainset):
loss_step = train_step(image_data, target)
loss_epoch[index] = loss_step
mask = loss_epoch[:,0] >0
loss_mean = np.mean(tf.boolean_mask(loss_epoch,mask),0)
loss_list_step = {"D2:lr":loss_mean[0],"total_loss":loss_mean[1],"loss/yolo_loss":loss_mean[2],
"G_loss":loss_mean[3],"D_loss":loss_mean[4],"loss/giou_loss":loss_mean[5],"loss/conf_loss":loss_mean[6],
"loss/prob_loss":loss_mean[7],"loss_content":loss_mean[8],"adv_lambda * loss_adv":loss_mean[9]}
loss_epoch_path = os.path.join(loss_dir_path, "epoch-{}".format(epoch) + '.txt')
with open(loss_epoch_path, 'w') as f:
for vm in loss_list_step.values():
loss_mess = ' '.join([str(vm)]) + '\n'
f.write(loss_mess)
print("No {} epoch params are {} and {}:".format(epoch,params[0].numpy(),params[1].numpy()))
total_loss_result.append(total_loss_metric.result())
loss1_result.append(loss1_metric.result())
loss2_result.append(loss2_metric.result())
total_loss_metric.reset_states()
loss1_metric.reset_states()
loss2_metric.reset_states()
if epoch % FLAGS.save_frequency == 0:
D2_model.save_weights(filepath=FLAGS.save_model_dir+"epoch-{}.h5".format(epoch), save_format="h5")
full_model.save_weights(filepath=FLAGS.save_model_dir+"Dis"+"epoch-{}".format(epoch), save_format="h5")
print("No {} epoch saved successfully...".format(epoch))
#Evaluation model
dectect_epoch_path = text_result_path + "-epoch-{}".format(epoch)
if os.path.exists(dectect_epoch_path): shutil.rmtree(dectect_epoch_path)
os.mkdir(dectect_epoch_path)
test_epoch(D2_model,dectect_epoch_path)
print("Evaluation completed...")
#####draw#####
total_loss = np.array(total_loss_result)
Yolo_loss = np.array(loss1_result)
G_loss = np.array(loss2_result)
epochs_range = np.arange(0,epoch+1,1)
plt.figure(dpi=1000,num=1,figsize=(6, 3))
plt.plot(epochs_range, total_loss, marker='*',linestyle='-',linewidth=1, markersize=2,label='total_loss')
plt.plot(epochs_range, Yolo_loss,marker='o', linestyle='-',linewidth=1, markersize=2,label='Yolo_loss')
plt.plot(epochs_range, G_loss, label='Deblur_loss')
plt.xlabel('Epoch')
plt.ylabel('Loss')
plt.legend(loc='upper right')
plt.savefig('Tranin_Loss_result.png',bbox_inches="tight",dpi=1000)
plt.show()
def main(_argv):
input_channel = 3
patience = 30
steps_in_epoch = 0
epoch_loss = 0
prev_minloss = np.inf
trainset = DatasetTF(FLAGS, input_channel, is_training=True)
testset = DatasetTF(FLAGS, input_channel, is_training=False)
# testset = Dataset(FLAGS, input_channel, is_training=False)
logdir = "./data/log_tf"
isfreeze = False
steps_per_epoch = len(trainset)
## for plot graph , early stopping
loss_tracker = []
# print("steps_per_epoch:{}".format(steps_per_epoch))
first_stage_epochs = cfg.TRAIN.FISRT_STAGE_EPOCHS
second_stage_epochs = cfg.TRAIN.SECOND_STAGE_EPOCHS
global_steps = tf.Variable(1, trainable=False, dtype=tf.int64)
warmup_steps = cfg.TRAIN.WARMUP_EPOCHS * steps_per_epoch
total_steps = (first_stage_epochs + second_stage_epochs) * steps_per_epoch
# train_steps = (first_stage_epochs + second_stage_epochs) * steps_per_period
input_layer = tf.keras.layers.Input(
[cfg.TRAIN.INPUT_SIZE, cfg.TRAIN.INPUT_SIZE, 3])
STRIDES, ANCHORS, NUM_CLASS, XYSCALE = utils.load_config(FLAGS)
IOU_LOSS_THRESH = cfg.YOLO.IOU_LOSS_THRESH
freeze_layers = utils.load_freeze_layer(FLAGS.model,
FLAGS.num_detection_layer)
# feature_maps = YOLO(input_layer, NUM_CLASS, FLAGS.model, FLAGS.tiny)
feature_maps = YOLOv4_more_tiny(input_layer, NUM_CLASS)
bbox_tensors = []
for i, fm in enumerate(feature_maps): # fm shape: (None, 40, 40, 21)
bbox_tensor = decode_train(fm, cfg.TRAIN.INPUT_SIZE // 16, NUM_CLASS,
STRIDES, ANCHORS, i, XYSCALE)
bbox_tensors.append(fm)
bbox_tensors.append(bbox_tensor)
if cfg.YOLO.NUM_YOLOLAYERS == 3: # yolov4
bbox_tensors = []
for i, fm in enumerate(feature_maps):
if i == 0:
bbox_tensor = decode_train(fm, cfg.TRAIN.INPUT_SIZE // 8,
NUM_CLASS, STRIDES, ANCHORS, i,
XYSCALE)
elif i == 1:
bbox_tensor = decode_train(fm, cfg.TRAIN.INPUT_SIZE // 16,
NUM_CLASS, STRIDES, ANCHORS, i,
XYSCALE)
else:
bbox_tensor = decode_train(fm, cfg.TRAIN.INPUT_SIZE // 32,
NUM_CLASS, STRIDES, ANCHORS, i,
XYSCALE)
bbox_tensors.append(fm)
bbox_tensors.append(bbox_tensor)
elif cfg.YOLO.NUM_YOLOLAYERS == 2: # yolo tiny
bbox_tensors = []
for i, fm in enumerate(feature_maps):
if i == 0:
bbox_tensor = decode_train(fm, cfg.TRAIN.INPUT_SIZE // 16,
NUM_CLASS, STRIDES, ANCHORS, i,
XYSCALE)
else:
bbox_tensor = decode_train(fm, cfg.TRAIN.INPUT_SIZE // 32,
NUM_CLASS, STRIDES, ANCHORS, i,
XYSCALE)
bbox_tensors.append(fm)
bbox_tensors.append(bbox_tensor)
elif cfg.YOLO.NUM_YOLOLAYERS == 1: # custom yolo
bbox_tensors = []
bbox_tensor = decode_train(
feature_maps[0],
cfg.TRAIN.INPUT_SIZE // cfg.YOLO.STRIDES_CUSTOM[0], NUM_CLASS,
STRIDES, ANCHORS, 0, XYSCALE)
bbox_tensors.append(feature_maps[0])
bbox_tensors.append(bbox_tensor)
model = tf.keras.Model(input_layer, bbox_tensors)
model.summary()
try:
if FLAGS.weights.split(".")[len(FLAGS.weights.split(".")) -
1] == "weights":
utils.load_weights(model, FLAGS.weights, FLAGS.model,
FLAGS.num_detection_layer)
else:
model.load_weights(FLAGS.weights)
print('Restoring weights from: %s ... ' % FLAGS.weights)
except Exception:
print("Training from scratch")
FLAGS.weights = None
optimizer = tf.keras.optimizers.Adam()
if os.path.exists(logdir): shutil.rmtree(logdir)
writer = tf.summary.create_file_writer(logdir)
# define training step function
@tf.function
def train_step(image_data, target):
with tf.GradientTape() as tape:
pred_result = model(image_data, training=True)
ciou_loss = conf_loss = prob_loss = 0
# optimizing process
for i in range(len(freeze_layers)):
conv, pred = pred_result[i * 2], pred_result[i * 2 + 1]
loss_items = compute_loss(
pred,
conv,
target[i][0],
target[i][1],
STRIDES=STRIDES,
NUM_CLASS=NUM_CLASS, # label, bbox
IOU_LOSS_THRESH=IOU_LOSS_THRESH,
i=i)
ciou_loss += loss_items[0]
conf_loss += loss_items[1]
prob_loss += loss_items[2]
total_loss = ciou_loss + conf_loss + prob_loss
gradients = tape.gradient(total_loss, model.trainable_variables)
optimizer.apply_gradients(zip(gradients,
model.trainable_variables))
tf.print("=> STEP ", global_steps, "/", total_steps, " lr: ",
optimizer.lr, " ciou_loss: ", ciou_loss,
" conf_loss: ", conf_loss, " prob_loss: ", prob_loss,
" total_loss: ", total_loss)
# update learning rate
global_steps.assign_add(1)
if global_steps < warmup_steps:
lr = global_steps / warmup_steps * cfg.TRAIN.LR_INIT
else:
lr = cfg.TRAIN.LR_END + 0.5 * (
cfg.TRAIN.LR_INIT - cfg.TRAIN.LR_END) * ((1 + tf.cos(
(global_steps - warmup_steps) /
(total_steps - warmup_steps) * np.pi)))
optimizer.lr.assign(tf.cast(lr, tf.float32))
# writing summary data
with writer.as_default():
tf.summary.scalar("lr", optimizer.lr, step=global_steps)
tf.summary.scalar("loss/total_loss",
total_loss,
step=global_steps)
tf.summary.scalar("loss/ciou_loss",
ciou_loss,
step=global_steps)
tf.summary.scalar("loss/conf_loss",
conf_loss,
step=global_steps)
tf.summary.scalar("loss/prob_loss",
prob_loss,
step=global_steps)
writer.flush()
return total_loss
@tf.function
def test_step(image_data, target):
with tf.GradientTape() as tape:
pred_result = model(image_data, training=True)
ciou_loss = conf_loss = prob_loss = 0
# optimizing process
for i in range(len(freeze_layers)):
conv, pred = pred_result[i * 2], pred_result[i * 2 + 1]
loss_items = compute_loss(pred,
conv,
target[i][0],
target[i][1],
STRIDES=STRIDES,
NUM_CLASS=NUM_CLASS,
IOU_LOSS_THRESH=IOU_LOSS_THRESH,
i=i)
# metric_items = compute_map(pred, conv, target[i][0], target[i][1], STRIDES=STRIDES, NUM_CLASS=NUM_CLASS, i=i)
# loss_items = compute_loss(pred, conv, target[1][0], target[1][1], STRIDES=STRIDES, NUM_CLASS=NUM_CLASS, IOU_LOSS_THRESH=IOU_LOSS_THRESH, i=i)
ciou_loss += loss_items[0]
conf_loss += loss_items[1]
prob_loss += loss_items[2]
total_loss = ciou_loss + conf_loss + prob_loss
tf.print("==> TEST STEP: ", global_steps)
tf.print(" ciou_loss: ", ciou_loss)
tf.print(" conf_loss: ", conf_loss)
tf.print(" prob_loss: ", prob_loss)
tf.print("==> total_loss", total_loss)
# tf.print("=> TEST STEP %4d giou_loss: %4.2f conf_loss: %4.2f "
# "prob_loss: %4.2f total_loss: %4.2f" % (global_steps, giou_loss, conf_loss,
# prob_loss, total_loss))
it = iter(trainset.get_dataset())
it_test = iter(testset.get_dataset())
start = time.time()
for epoch in range(first_stage_epochs + second_stage_epochs):
train_loss = 0.
if epoch < first_stage_epochs:
if not isfreeze:
isfreeze = True
for name in freeze_layers:
try: # try구문 추가
freeze = model.get_layer(name)
freeze_all(freeze)
except ValueError:
pass
elif epoch >= first_stage_epochs:
if isfreeze:
isfreeze = False
for name in freeze_layers:
try: # try구문 추가
freeze = model.get_layer(name)
freeze_all(freeze)
except ValueError:
pass
# freeze = model.get_layer(name)
# unfreeze_all(freeze)
for i in range(trainset.steps_for_train):
data = it.get_next()
### check input train images
# image_data = data[0]
# label_data = np.array(data[1])
#
# bboxes_list = np.array(label_data[0][1])
# label_list = np.array(label_data[0][0])
# for batch_idx, bboxes in enumerate(bboxes_list):
# class_inds = []
# check = np.array(image_data[batch_idx]).reshape(cfg.TRAIN.INPUT_SIZE, cfg.TRAIN.INPUT_SIZE, 3)
# # print(r"D:\tf_data_\preprocessed_{}.npy".format(batch_idx))
# # np.save(r"D:\tf_data_\image_{}.npy".format(batch_idx), check)
# # np.save(r"D:\tf_data_\preprocessed_{}.npy".format(batch_idx), label_list[batch_idx])
# # np.save(r"D:\tf_data_\bboxes_{}.npy".format(batch_idx), bboxes_list[batch_idx])
# label = np.array(label_list[batch_idx])
#
# class_ = label[...,5:]
# class_ = np.array(class_).flatten()
# class_ = np.where(class_ > 0.1, class_, 0)
#
# for class_idx, class_label in enumerate(class_):
# # print("class_label:{}".format(class_label))
# if class_label > 0:
# class_inds.append(class_idx % cfg.YOLO.NUM_CLASSES)
#
# bboxes = np.array(bboxes)
# for bbox_idx, bbox in enumerate(bboxes):
# half_h = bbox[3] / 2
# half_w = bbox[2] / 2
#
# if half_h + half_w > 0:
# cv2.rectangle(check, (int(bbox[0] - half_w), int(bbox[1] - half_h)),
# (int(bbox[0] + half_w), int(bbox[1] + half_h)), color=(0, 255, 0), thickness=3)
# cv2.putText(check, text="{}".format(class_inds),
# org=(int(bbox[0] + half_w)-10, int(bbox[1] + half_h)-30), thickness=2, color=(0, 255, 0),
# fontFace=cv2.FONT_HERSHEY_SIMPLEX, fontScale=0.7)
#
# cv2.imshow("check_aug", check)
# cv2.waitKey(0)
# cv2.destroyAllWindows()
# cv2.imwrite(os.path.join(r"D:\Public\JHS\SIMPLE_DATA_INPUT_CHECK", "{}.jpg".format(batch_idx)), check*255)
train_loss = train_step(data[0], data[1])
# epoch_loss += train_loss
# for i in range(testset.steps_for_train):
# data = it_test.get_next()
# test_step(data[0], data[1])
# CHECK TEST
# image_data = data[0]
# label_data = data[1]
# bboxes_list = label_data[0][1]
#
# for batch_idx, bboxes in enumerate(bboxes_list):
# check = np.array(image_data[batch_idx]).reshape(640, 640, 3)
# for bbox in bboxes:
# half_h = bbox[3] / 2
# half_w = bbox[2] / 2
# # bbox_upper_left = (bbox[0])
# if np.sum(bbox) > 0:
# cv2.rectangle(check, (bbox[0] - half_w, bbox[1] - half_h),(bbox[0] + half_w,bbox[1]+half_h), color=(0,255,0), thickness=3)
# cv2.imshow("check_aug TEST", check)
# cv2.waitKey(0)
# cv2.destroyAllWindows()
# for image_data, target in testset:
# test_step(image_data, target)
# for loss graph
"""
if steps_in_epoch >= steps_per_epoch:
loss_tracker.append(losses_in_epoch / steps_per_epoch)
## 밑으로
# plt.plot(loss_tracker)
# plt.xlabel('epoch')
# plt.ylabel('loss')
# # plt.show()
# plt.savefig("D:/checkpoint/loss.png")
"""
# early stopping
"""
epoch_loss = epoch_loss / steps_per_epoch
loss_tracker.append(epoch_loss)
if prev_minloss > epoch_loss: # save best weight
prev_minloss = epoch_loss
if epoch > 2000: # 최소학습 에폭
model.save("D:\ckpt_best")
print("{} epoch save best weights".format(epoch))
if len(loss_tracker) > patience:
print("check loss_tracker len:{}".format(len(loss_tracker)))
if loss_tracker[0] > loss_tracker[-1]:
loss_tracker.pop(0)
else:
print("total loss didn't decreased during {} epochs. train stop".format(patience))
return
epoch_loss = 0.
else:
epoch_loss += train_loss
steps_in_epoch = 0
"""
if (epoch + 1) % 100 == 0:
model.save("D:\yolov4-tflite-train_tf-epoch{}".format(epoch + 1))
print("{} epoch model saved".format(epoch + 1))
print("consumed time: {}".format(time.time() - start))
model.save(r"D:\yolov4-tflite-train_tf-last")
