def on_batch_end(self, batch, logs=None):
self.train_loss.update_state(logs['loss'])
if batch % self.eval_freq == 0:
image, bbox, labels = next(self.iter)
image = image.numpy().astype('uint8')
predictor = Predictor(yolov5l=yolov5l)
boundings = predictor.predict(image)
color_map = dict()
for bounding in boundings:
if bounding[5].numpy().astype('int32') not in color_map:
color_map[bounding[5].numpy().astype('int32')] = tuple(
np.random.randint(low=0, high=256,
size=(3, )).tolist())
clr = color_map[bounding[5].numpy().astype('int32')]
cv2.rectangle(image,
tuple(bounding[0:2].numpy().astype('int32')),
tuple(bounding[2:4].numpy().astype('int32')),
clr, 1)
cv2.putText(
image,
predictor.getClsName(
bounding[5].numpy().astype('int32')),
tuple(bounding[0:2].numpy().astype('int32')),
cv2.FONT_HERSHEY_PLAIN, 1, clr, 2)
image = tf.expand_dims(image, axis=0)
with self.log.as_default():
tf.summary.scalar('train loss',
self.train_loss.result(),
step=optimizer.iterations)
tf.summary.image('detect',
image[..., ::-1],
step=optimizer.iterations)
self.train_loss.reset_states()
def main():
gpus = tf.config.experimental.list_physical_devices('GPU')
[tf.config.experimental.set_memory_growth(gpu, True) for gpu in gpus]
# yolov5l model
yolov5l = YOLOv5_large((608, 608, 3), 80)
loss1 = Loss((608, 608, 3), 0, 80)
loss2 = Loss((608, 608, 3), 1, 80)
loss3 = Loss((608, 608, 3), 2, 80)
#optimizer = tf.keras.optimizers.Adam(tf.keras.optimizers.schedules.ExponentialDecay(1e-5, decay_steps = 110000, decay_rate = 0.99));
optimizer = tf.keras.optimizers.Adam(1e-5)
checkpoint = tf.train.Checkpoint(model=yolov5l, optimizer=optimizer)
train_loss = tf.keras.metrics.Mean(name='train loss', dtype=tf.float32)
test_loss = tf.keras.metrics.Mean(name='test loss', dtype=tf.float32)
# load downloaded dataset
trainset_filenames = [
join('trainset', filename) for filename in listdir('trainset')
]
testset_filenames = [
join('testset', filename) for filename in listdir('testset')
]
trainset = tf.data.TFRecordDataset(trainset_filenames).map(
parse_function_generator(80)).repeat(-1).shuffle(batch_size).batch(
batch_size).prefetch(tf.data.experimental.AUTOTUNE)
testset = tf.data.TFRecordDataset(testset_filenames).map(
parse_function_generator(80)).repeat(-1).shuffle(batch_size).batch(
batch_size).prefetch(tf.data.experimental.AUTOTUNE)
validationset = tf.data.TFRecordDataset(testset_filenames).map(
parse_function).repeat(-1)
trainset_iter = iter(trainset)
testset_iter = iter(testset)
validationset_iter = iter(validationset)
# restore from existing checkpoint
if False == exists('checkpoints'): mkdir('checkpoints')
checkpoint.restore(tf.train.latest_checkpoint('checkpoints'))
# tensorboard summary
log = tf.summary.create_file_writer('checkpoints')
# train model
while True:
images, labels = next(trainset_iter)
labels1, labels2, labels3 = labels
with tf.GradientTape() as tape:
outputs1, outputs2, outputs3 = yolov5l(images)
loss = loss1([outputs1, labels1]) + loss2(
[outputs2, labels2]) + loss3([outputs3, labels3])
# check whether the loss numberic is correct
if tf.math.reduce_any(tf.math.is_nan(loss)) == True:
print("NaN was detected in loss, skip the following steps!")
continue
grads = tape.gradient(loss, yolov5l.trainable_variables)
# check whether the grad numerics is correct
if tf.math.reduce_any([
tf.math.reduce_any(tf.math.is_nan(grad)) for grad in grads
]) == True:
print("NaN was detected in gradients, skip gradient apply!")
continue
optimizer.apply_gradients(zip(grads, yolov5l.trainable_variables))
train_loss.update_state(loss)
# save model
if tf.equal(optimizer.iterations % 10000, 0):
# save checkpoint every 1000 steps
checkpoint.save(join('checkpoints', 'ckpt'))
yolov5l.save('yolov5l.h5')
if tf.equal(optimizer.iterations % 100, 0):
# evaluate
for i in range(10):
images, labels = next(testset_iter)
# images.shape = (b, h, w, 3)
outputs = yolov5l(images)
loss = yolov3_loss([*outputs, *labels])
test_loss.update_state(loss)
# visualize
image, bbox, labels = next(validationset_iter)
# image.shape = (h, w, 3)
image = image.numpy().astype('uint8')
predictor = Predictor(yolov5l=yolov5l)
boundings = predictor.predict(image)
color_map = dict()
for bounding in boundings:
if bounding[5].numpy().astype('int32') not in color_map:
color_map[bounding[5].numpy().astype('int32')] = tuple(
np.random.randint(low=0, high=256,
size=(3, )).tolist())
clr = color_map[bounding[5].numpy().astype('int32')]
cv2.rectangle(image,
tuple(bounding[0:2].numpy().astype('int32')),
tuple(bounding[2:4].numpy().astype('int32')),
clr, 1)
cv2.putText(
image,
predictor.getClsName(bounding[5].numpy().astype('int32')),
tuple(bounding[0:2].numpy().astype('int32')),
cv2.FONT_HERSHEY_PLAIN, 1, clr, 2)
image = tf.expand_dims(image, axis=0)
# write log
with log.as_default():
tf.summary.scalar('train loss',
train_loss.result(),
step=optimizer.iterations)
tf.summary.scalar('test loss',
test_loss.result(),
step=optimizer.iterations)
tf.summary.image('detect',
image[..., ::-1],
step=optimizer.iterations)
print('Step #%d Train Loss: %.6f Test Loss: %.6f' %
(optimizer.iterations, train_loss.result(),
test_loss.result()))
# break condition
#if train_loss.result() < 0.001: break;
# reset
train_loss.reset_states()
test_loss.reset_states()
yolov5l.save('yolov5l.h5')
