def main():
yolov3 = YOLOv3((416, 416, 3), 80)
optimizer = tf.keras.optimizers.Adam(
tf.keras.optimizers.schedules.ExponentialDecay(1e-4,
decay_steps=110000,
decay_rate=0.99))
checkpoint = tf.train.Checkpoint(model=yolov3, optimizer=optimizer)
checkpoint.restore(tf.train.latest_checkpoint('checkpoints'))
yolov3.save('yolov3.h5')
yolov3.save_weights('yolov3_weights.h5')
def test_yolov3_num_params_NCHW(self):
"""Test to see if the number of parameters in YOLOv3 is 62001757 with tensors in NCHW format."""
c = copy.deepcopy(config)
c['DATA_FORMAT'] = 'NCHW'
with tf.variable_scope('model'):
_ = YOLOv3(c)
global_vars = tf.global_variables(scope='model')
num_params = 0
for var in global_vars:
shape = var.get_shape().as_list()
num_params += np.prod(shape)
self.assertAllEqual(num_params, 62001757)
def main():
strategy = tf.distribute.MirroredStrategy();
# yolov3 model
with strategy.scope():
yolov3 = YOLOv3((416,416,3,), 80);
@tf.function
def loss(labels, outputs):
return Loss((416,416,3,),80)([outputs[0], outputs[1], outputs[2], labels[0], labels[1], labels[2]]);
yolov3.compile(optimizer = tf.keras.optimizers.Adam(1e-4), loss = loss);
# 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);
yolov3.fit(trainset, epochs = 100, validation_data = testset);
yolov3.save('yolov3.h5');
def __init__(self, input_shape=(416, 416, 3), class_num=80, yolov3=None):
if yolov3 is None:
self.input_shape = input_shape
if exists("yolov3.h5"):
# restore from serialized file
self.yolov3 = tf.keras.models.load_model('yolov3.h5',
compile=False)
else:
# restore from checkpoint
self.yolov3 = YOLOv3(input_shape, class_num)
optimizer = tf.keras.optimizers.Adam(1e-4)
checkpoint = tf.train.Checkpoint(
model=self.yolov3,
optimizer=optimizer,
optimizer_step=optimizer.iterations)
checkpoint.restore(tf.train.latest_checkpoint('checkpoints'))
else:
self.input_shape = tuple(yolov3.input.shape[1:])
self.yolov3 = yolov3
self.parsers = [
OutputParser(tuple(self.yolov3.outputs[l].shape[1:]),
self.input_shape, self.anchors[l]) for l in range(3)
]
def main():
yolov3 = YOLOv3((416, 416, 3), 80);
yolov3.load_weights('./checkpoints/ckpt/variables/variables');
yolov3.save('yolov3.h5');
yolov3.save_weights('yolov3_weights.h5');
def main():
gpus = tf.config.experimental.list_physical_devices('GPU')
[tf.config.experimental.set_memory_growth(gpu, True) for gpu in gpus]
# yolov3 model
yolov3 = YOLOv3((416, 416, 3), 80)
loss1 = Loss((416, 416, 3), 0, 80)
loss2 = Loss((416, 416, 3), 1, 80)
loss3 = Loss((416, 416, 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=yolov3, 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 = yolov3(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, yolov3.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, yolov3.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'))
yolov3.save('yolov3.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 = yolov3(images)
loss = loss1([outputs[0], labels[0]]) + loss2(
[outputs[1], labels[1]]) + loss3([outputs[2], labels[2]])
test_loss.update_state(loss)
# visualize
image, bbox, labels = next(validationset_iter)
# image.shape = (h, w, 3)
image = image.numpy().astype('uint8')
predictor = Predictor(yolov3=yolov3)
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()
yolov3.save('yolov3.h5')
import tensorflow as tf
from tqdm import tqdm
from models import YOLOv3
from utils import postprocessing
from config import config
from utils.load_yolov3_weights import load_yolov3_weights
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument('input', help='Input image/video.')
parser.add_argument('-c', '--classes', nargs='+', help='Predict the specific class only.')
args = parser.parse_args()
with tf.variable_scope('model'):
print("Constructing computational graph...")
model = YOLOv3(config)
print("Done")
print("Loading weights...")
global_vars = tf.global_variables(scope='model')
assign_ops = load_yolov3_weights(global_vars, config['WEIGHTS_PATH'])
[print(n.name) for n in tf.get_default_graph().as_graph_def().node]
print("Done")
print("=============================================")
print("Loading class names...")
classes = []
colours = {}
f = open(config['CLASS_PATH'], 'r').read().splitlines()
for i, line in enumerate(f):
from models import YOLOv3
from config import img_size, coco_cat2int, coco_int2cat, cat2cat, cat2int, int2cat, anchors, num_classes
from utils import COCODataset, bboxes_to_labels, yolo_loss, draw_bboxes
data_dir = "data/COCO/cleaned/val2017/"
model_dir = "model/yolo.pt"
optim_dir = "model/optim.pt"
summary_dir = "summary/"
gs_dir = "summary/gs.pkl"
batch_size = 2
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
gs = 1
model = YOLOv3().to(device)
optim = torch.optim.Adam(filter(lambda p: p.requires_grad, model.parameters()),
lr=.001,
betas=(.5, .999))
resume_training = False
if resume_training:
print("Loading states to resume training...")
if os.path.exists(gs_dir):
gs = pickle.load(open(gs_dir))
else:
raise FileNotFoundError("Global steps check point not found")
if os.path.exists(model_dir):
state_dict = torch.load(model_dir, map_location=device)
model.load_state_dict(state_dict)
else:
def main():
yolov3 = YOLOv3((
416,
416,
3,
), 80)
loss1 = Loss((
416,
416,
3,
), 0, 80)
loss2 = Loss((
416,
416,
3,
), 1, 80)
loss3 = Loss((
416,
416,
3,
), 2, 80)
if exists('./checkpoints/ckpt'):
yolov3.load_weights('./checkpoints/ckpt/variables/variables')
optimizer = tf.keras.optimizers.Adam(1e-5)
yolov3.compile(optimizer=optimizer,
loss={
'output1':
lambda labels, outputs: loss1([outputs, labels]),
'output2':
lambda labels, outputs: loss2([outputs, labels]),
'output3':
lambda labels, outputs: loss3([outputs, labels])
})
class SummaryCallback(tf.keras.callbacks.Callback):
def __init__(self, eval_freq=100):
self.eval_freq = eval_freq
testset = tf.data.TFRecordDataset(testset_filenames).map(
parse_function).repeat(-1)
self.iter = iter(testset)
self.train_loss = tf.keras.metrics.Mean(name='train loss',
dtype=tf.float32)
self.log = tf.summary.create_file_writer('./checkpoints')
def on_batch_begin(self, batch, logs=None):
pass
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(yolov3=yolov3)
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 on_epoch_begin(self, epoch, logs=None):
pass
def on_epoch_end(self, batch, logs=None):
pass
# 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)).shuffle(batch_size).batch(
batch_size).prefetch(tf.data.experimental.AUTOTUNE)
testset = tf.data.TFRecordDataset(testset_filenames).map(
parse_function_generator(80)).shuffle(batch_size).batch(
batch_size).prefetch(tf.data.experimental.AUTOTUNE)
callbacks = [
tf.keras.callbacks.TensorBoard(log_dir='./checkpoints'),
tf.keras.callbacks.ModelCheckpoint(filepath='./checkpoints/ckpt',
save_freq=10000),
SummaryCallback(),
]
yolov3.fit(trainset,
epochs=100,
validation_data=testset,
callbacks=callbacks)
yolov3.save('yolov3.h5')