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)
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():
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):
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")
