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')