def __init__(self): self.log = logging.getLogger('EMSfIIoT') # Load TFLite model and allocate tensors. try: self.interpreter = tflite.Interpreter( model_path="weights/emsfiiot_lite.h5") except: self.interpreter = tf.lite.Interpreter( model_path="weights/emsfiiot_lite.h5") self.interpreter.allocate_tensors() # Get input and output tensors. self.input_details = self.interpreter.get_input_details() self.output_details = self.interpreter.get_output_details() # check the type of the input tensor self.floating_model = self.input_details[0]['dtype'] == np.float32 height = self.input_details[0]['shape'][1] width = self.input_details[0]['shape'][2] self.model_image_size = (width, height) self.anchors = get_anchors("configs/yolo3_anchors.txt") self.class_names = get_classes("configs/emsfiiot_classes.txt") self.colors = get_colors(self.class_names)
def main(): parser = argparse.ArgumentParser(description='validate YOLO model (h5/pb/onnx/tflite/mnn) with image') parser.add_argument('--model_path', help='model file to predict', type=str, required=True) parser.add_argument('--image_file', help='image file to predict', type=str, required=True) parser.add_argument('--anchors_path',help='path to anchor definitions', type=str, required=True) parser.add_argument('--classes_path', help='path to class definitions, default ../../configs/voc_classes.txt', type=str, default='../../configs/voc_classes.txt') parser.add_argument('--model_image_size', help='model image input size as <height>x<width>, default 416x416', type=str, default='416x416') parser.add_argument('--loop_count', help='loop inference for certain times', type=int, default=1) args = parser.parse_args() # param parse anchors = get_anchors(args.anchors_path) class_names = get_classes(args.classes_path) height, width = args.model_image_size.split('x') model_image_size = (int(height), int(width)) assert (model_image_size[0]%32 == 0 and model_image_size[1]%32 == 0), 'model_image_size should be multiples of 32' # support of tflite model if args.model_path.endswith('.tflite'): validate_yolo_model_tflite(args.model_path, args.image_file, anchors, class_names, args.loop_count) # support of MNN model elif args.model_path.endswith('.mnn'): validate_yolo_model_mnn(args.model_path, args.image_file, anchors, class_names, args.loop_count) # support of TF 1.x frozen pb model elif args.model_path.endswith('.pb'): validate_yolo_model_pb(args.model_path, args.image_file, anchors, class_names, model_image_size, args.loop_count) # support of ONNX model elif args.model_path.endswith('.onnx'): validate_yolo_model_onnx(args.model_path, args.image_file, anchors, class_names, args.loop_count) # normal keras h5 model elif args.model_path.endswith('.h5'): validate_yolo_model(args.model_path, args.image_file, anchors, class_names, model_image_size, args.loop_count) else: raise ValueError('invalid model file')
def main(): # class YOLO defines the default value, so suppress any default here parser = argparse.ArgumentParser(argument_default=argparse.SUPPRESS, description='evaluate YOLO model (h5/pb/tflite/mnn) with test dataset') ''' Command line options ''' parser.add_argument( '--model_path', type=str, required=True, help='path to model file') parser.add_argument( '--custom_objects', type=str, required=False, default=None, help="Custom objects in keras model (swish/tf). Separated with comma if more than one.") parser.add_argument( '--anchors_path', type=str, required=True, help='path to anchor definitions') parser.add_argument( '--classes_path', type=str, required=False, help='path to class definitions, default configs/voc_classes.txt', default='configs/voc_classes.txt') parser.add_argument( '--annotation_file', type=str, required=True, help='test annotation txt file') parser.add_argument( '--eval_type', type=str, help='evaluation type (VOC/COCO), default=VOC', default='VOC') parser.add_argument( '--iou_threshold', type=float, help='IOU threshold for PascalVOC mAP, default=0.5', default=0.5) parser.add_argument( '--conf_threshold', type=float, help='confidence threshold for filtering box in postprocess, default=0.001', default=0.001) parser.add_argument( '--model_image_size', type=str, help='model image input size as <num>x<num>, default 416x416', default='416x416') parser.add_argument( '--save_result', default=False, action="store_true", help='Save the detection result image in result/detection dir' ) args = parser.parse_args() # param parse anchors = get_anchors(args.anchors_path) class_names = get_classes(args.classes_path) height, width = args.model_image_size.split('x') model_image_size = (int(height), int(width)) annotation_lines = get_dataset(args.annotation_file) model, model_format = load_eval_model(args.model_path, args.custom_objects) eval_AP(model, model_format, annotation_lines, anchors, class_names, model_image_size, args.eval_type, args.iou_threshold, args.conf_threshold, args.save_result)
def main(): # class YOLO defines the default value, so suppress any default here parser = argparse.ArgumentParser(argument_default=argparse.SUPPRESS, description='evaluate YOLO model (h5/pb/onnx/tflite/mnn) with test dataset') ''' Command line options ''' parser.add_argument( '--model_path', type=str, required=True, help='path to model file') parser.add_argument( '--anchors_path', type=str, required=True, help='path to anchor definitions') parser.add_argument( '--classes_path', type=str, required=False, help='path to class definitions, default configs/voc_classes.txt', default=os.path.join('configs' , 'voc_classes.txt')) parser.add_argument( '--annotation_file', type=str, required=True, help='test annotation txt file') parser.add_argument( '--eval_type', type=str, help='evaluation type (VOC/COCO), default=VOC', default='VOC') parser.add_argument( '--iou_threshold', type=float, help='IOU threshold for PascalVOC mAP, default=0.5', default=0.5) parser.add_argument( '--conf_threshold', type=float, help='confidence threshold for filtering box in postprocess, default=0.001', default=0.001) parser.add_argument( '--model_image_size', type=str, help='model image input size as <height>x<width>, default 416x416', default='416x416') parser.add_argument( '--save_result', default=False, action="store_true", help='Save the detection result image in result/detection dir' ) args = parser.parse_args() # param parse anchors = get_anchors(args.anchors_path) class_names = get_classes(args.classes_path) height, width = args.model_image_size.split('x') model_image_size = (int(height), int(width)) assert (model_image_size[0]%32 == 0 and model_image_size[1]%32 == 0), 'model_image_size should be multiples of 32' annotation_lines = get_dataset(args.annotation_file, shuffle=False) model, model_format = load_eval_model(args.model_path) start = time.time() eval_AP(model, model_format, annotation_lines, anchors, class_names, model_image_size, args.eval_type, args.iou_threshold, args.conf_threshold, args.save_result) end = time.time() print("Evaluation time cost: {:.6f}s".format(end - start))
def __init__(self, **kwargs): super(YOLO_np, self).__init__() self.__dict__.update(self._defaults) # set up default values self.__dict__.update(kwargs) # and update with user overrides self.class_names = get_classes(self.classes_path) self.anchors = get_anchors(self.anchors_path) self.colors = get_colors(self.class_names) K.set_learning_phase(0) self.yolo_model = self._generate_model()
def main(): parser = argparse.ArgumentParser(description='validate YOLO model (h5/pb/onnx/tflite/mnn) with image') parser.add_argument('--model_path', help='model file to predict', type=str, required=True) parser.add_argument('--image_path', help='image file or directory to predict', type=str, required=True) parser.add_argument('--anchors_path', help='path to anchor definitions', type=str, required=True) parser.add_argument('--classes_path', help='path to class definitions, default=%(default)s', type=str, default='../../configs/voc_classes.txt') parser.add_argument('--model_input_shape', help='model image input shape as <height>x<width>, default=%(default)s', type=str, default='416x416') parser.add_argument('--elim_grid_sense', help="Eliminate grid sensitivity", default=False, action="store_true") parser.add_argument('--v5_decode', help="Use YOLOv5 prediction decode", default=False, action="store_true") parser.add_argument('--loop_count', help='loop inference for certain times', type=int, default=1) parser.add_argument('--output_path', help='output path to save predict result, default=%(default)s', type=str, required=False, default=None) args = parser.parse_args() # param parse anchors = get_anchors(args.anchors_path) class_names = get_classes(args.classes_path) height, width = args.model_input_shape.split('x') model_input_shape = (int(height), int(width)) assert (model_input_shape[0]%32 == 0 and model_input_shape[1]%32 == 0), 'model_input_shape should be multiples of 32' model = load_val_model(args.model_path) if args.model_path.endswith('.mnn'): #MNN inference engine need create session session = model.createSession() # get image file list or single image if os.path.isdir(args.image_path): image_files = glob.glob(os.path.join(args.image_path, '*')) assert args.output_path, 'need to specify output path if you use image directory as input.' else: image_files = [args.image_path] # loop the sample list to predict on each image for image_file in image_files: # support of tflite model if args.model_path.endswith('.tflite'): validate_yolo_model_tflite(model, image_file, anchors, class_names, args.elim_grid_sense, args.v5_decode, args.loop_count, args.output_path) # support of MNN model elif args.model_path.endswith('.mnn'): validate_yolo_model_mnn(model, session, image_file, anchors, class_names, args.elim_grid_sense, args.v5_decode, args.loop_count, args.output_path) # support of TF 1.x frozen pb model elif args.model_path.endswith('.pb'): validate_yolo_model_pb(model, image_file, anchors, class_names, model_input_shape, args.elim_grid_sense, args.v5_decode, args.loop_count, args.output_path) # support of ONNX model elif args.model_path.endswith('.onnx'): validate_yolo_model_onnx(model, image_file, anchors, class_names, args.elim_grid_sense, args.v5_decode, args.loop_count, args.output_path) # normal keras h5 model elif args.model_path.endswith('.h5'): validate_yolo_model(model, image_file, anchors, class_names, model_input_shape, args.elim_grid_sense, args.v5_decode, args.loop_count, args.output_path) else: raise ValueError('invalid model file')
def main(args): annotation_file = args.annotation_file log_dir = os.path.join('logs', '000') classes_path = args.classes_path class_names = get_classes(classes_path) num_classes = len(class_names) print('classes_path =', classes_path) print('class_names = ', class_names) print('num_classes = ', num_classes) anchors = get_anchors(args.anchors_path) num_anchors = len(anchors) # get freeze level according to CLI option if args.weights_path: freeze_level = 0 else: freeze_level = 1 if args.freeze_level is not None: freeze_level = args.freeze_level # callbacks for training process logging = TensorBoard(log_dir=log_dir, histogram_freq=0, write_graph=False, write_grads=False, write_images=False, update_freq='batch') checkpoint = ModelCheckpoint(os.path.join( log_dir, 'ep{epoch:03d}-loss{loss:.3f}-val_loss{val_loss:.3f}.h5'), monitor='val_loss', verbose=1, save_weights_only=False, save_best_only=True, period=1) reduce_lr = ReduceLROnPlateau(monitor='val_loss', factor=0.5, patience=10, verbose=1, cooldown=0, min_lr=1e-10) early_stopping = EarlyStopping(monitor='val_loss', min_delta=0, patience=50, verbose=1) terminate_on_nan = TerminateOnNaN() callbacks = [ logging, checkpoint, reduce_lr, early_stopping, terminate_on_nan ] # get train&val dataset dataset = get_dataset(annotation_file) if args.val_annotation_file: val_dataset = get_dataset(args.val_annotation_file) num_train = len(dataset) print('num_train = ', num_train) num_val = len(val_dataset) dataset.extend(val_dataset) else: val_split = args.val_split num_val = int(len(dataset) * val_split) num_train = len(dataset) - num_val # assign multiscale interval if args.multiscale: rescale_interval = args.rescale_interval else: rescale_interval = -1 #Doesn't rescale # model input shape check input_shape = args.model_image_size assert (input_shape[0] % 32 == 0 and input_shape[1] % 32 == 0), 'Multiples of 32 required' # get different model type & train&val data generator if num_anchors == 9: # YOLOv3 use 9 anchors get_train_model = get_yolo3_train_model data_generator = yolo3_data_generator_wrapper # tf.keras.Sequence style data generator #train_data_generator = Yolo3DataGenerator(dataset[:num_train], args.batch_size, input_shape, anchors, num_classes, args.enhance_augment, rescale_interval) #val_data_generator = Yolo3DataGenerator(dataset[num_train:], args.batch_size, input_shape, anchors, num_classes) tiny_version = False elif num_anchors == 6: # Tiny YOLOv3 use 6 anchors get_train_model = get_yolo3_train_model data_generator = yolo3_data_generator_wrapper # tf.keras.Sequence style data generator #train_data_generator = Yolo3DataGenerator(dataset[:num_train], args.batch_size, input_shape, anchors, num_classes, args.enhance_augment, rescale_interval) #val_data_generator = Yolo3DataGenerator(dataset[num_train:], args.batch_size, input_shape, anchors, num_classes) tiny_version = True elif num_anchors == 5: # YOLOv2 use 5 anchors get_train_model = get_yolo2_train_model data_generator = yolo2_data_generator_wrapper # tf.keras.Sequence style data generator #train_data_generator = Yolo2DataGenerator(dataset[:num_train], args.batch_size, input_shape, anchors, num_classes, args.enhance_augment, rescale_interval) #val_data_generator = Yolo2DataGenerator(dataset[num_train:], args.batch_size, input_shape, anchors, num_classes) tiny_version = False else: raise ValueError('Unsupported anchors number') # prepare online evaluation callback if args.eval_online: eval_callback = EvalCallBack( args.model_type, dataset[num_train:], anchors, class_names, args.model_image_size, args.model_pruning, log_dir, eval_epoch_interval=args.eval_epoch_interval, save_eval_checkpoint=args.save_eval_checkpoint) callbacks.append(eval_callback) # prepare train/val data shuffle callback if args.data_shuffle: shuffle_callback = DatasetShuffleCallBack(dataset) callbacks.append(shuffle_callback) # prepare model pruning config pruning_end_step = np.ceil(1.0 * num_train / args.batch_size).astype( np.int32) * args.total_epoch if args.model_pruning: pruning_callbacks = [ sparsity.UpdatePruningStep(), sparsity.PruningSummaries(log_dir=log_dir, profile_batch=0) ] callbacks = callbacks + pruning_callbacks # prepare optimizer optimizer = get_optimizer(args.optimizer, args.learning_rate, decay_type=None) # get train model model = get_train_model(args.model_type, anchors, num_classes, weights_path=args.weights_path, freeze_level=freeze_level, optimizer=optimizer, label_smoothing=args.label_smoothing, model_pruning=args.model_pruning, pruning_end_step=pruning_end_step) # support multi-gpu training template_model = None if args.gpu_num >= 2: # keep the template model for saving result template_model = model model = multi_gpu_model(model, gpus=args.gpu_num) # recompile multi gpu model model.compile(optimizer=optimizer, loss={ 'yolo_loss': lambda y_true, y_pred: y_pred }) model.summary() # Transfer training some epochs with frozen layers first if needed, to get a stable loss. initial_epoch = args.init_epoch ##################################################################################################### epochs = initial_epoch + args.transfer_epoch print("Transfer training stage") print( 'Train on {} samples, val on {} samples, with batch size {}, input_shape {}.' .format(num_train, num_val, args.batch_size, input_shape)) #model.fit_generator(train_data_generator, model.fit_generator( data_generator(dataset[:num_train], args.batch_size, input_shape, anchors, num_classes, args.enhance_augment), steps_per_epoch=max(1, num_train // args.batch_size), #validation_data=val_data_generator, validation_data=data_generator(dataset[num_train:], args.batch_size, input_shape, anchors, num_classes), validation_steps=max(1, num_val // args.batch_size), epochs=epochs, initial_epoch=initial_epoch, #verbose=1, workers=1, use_multiprocessing=False, max_queue_size=10, callbacks=callbacks) # Wait 2 seconds for next stage time.sleep(2) if args.decay_type: # rebuild optimizer to apply learning rate decay, only after # unfreeze all layers callbacks.remove(reduce_lr) steps_per_epoch = max(1, num_train // args.batch_size) decay_steps = steps_per_epoch * (args.total_epoch - args.init_epoch - args.transfer_epoch) optimizer = get_optimizer(args.optimizer, args.learning_rate, decay_type=args.decay_type, decay_steps=decay_steps) # Unfreeze the whole network for further tuning # NOTE: more GPU memory is required after unfreezing the body print("Unfreeze and continue training, to fine-tune.") for i in range(len(model.layers)): model.layers[i].trainable = True model.compile(optimizer=optimizer, loss={ 'yolo_loss': lambda y_true, y_pred: y_pred }) # recompile to apply the change print( 'Train on {} samples, val on {} samples, with batch size {}, input_shape {}.' .format(num_train, num_val, args.batch_size, input_shape)) #model.fit_generator(train_data_generator, model.fit_generator( data_generator(dataset[:num_train], args.batch_size, input_shape, anchors, num_classes, args.enhance_augment, rescale_interval), steps_per_epoch=max(1, num_train // args.batch_size), #validation_data=val_data_generator, validation_data=data_generator(dataset[num_train:], args.batch_size, input_shape, anchors, num_classes), validation_steps=max(1, num_val // args.batch_size), epochs=args.total_epoch, initial_epoch=epochs, #verbose=1, workers=1, use_multiprocessing=False, max_queue_size=10, callbacks=callbacks) # Finally store model if args.model_pruning: if template_model is not None: template_model = sparsity.strip_pruning(template_model) else: model = sparsity.strip_pruning(model) if template_model is not None: template_model.save(os.path.join(log_dir, 'trained_final.h5')) else: model.save(os.path.join(log_dir, 'trained_final.h5'))
if val == 'y': print(" Renaming has been enabled") elif val == 'a': print("Aborted") exit(0) else: print("Renaming has been disabled") args.rename = False # Load model and process input image model = tf.saved_model.load(model_file) # Load the labels if label_file: classes = load_coco_names(label_file) anchors = get_anchors(anchor_file) if not args.disable_ovtf: # Print list of available backends print('Available Backends:') backends_list = ovtf.list_backends() for backend in backends_list: print(backend) ovtf.set_backend(backend_name) else: ovtf.disable() cap = None images = [] if label_file: labels = load_labels(label_file)
def main(args): #데이터 annotation 파일 경로 annotation_file = args.annotation_file # 결과 log 및 weight가 저장될 경로 log_dir = os.path.join('logs', '000') #클래스 파일 경로 classes_path = args.classes_path class_names = get_classes(classes_path) num_classes = len(class_names) # anchors 받아오는 라인 anchors = get_anchors(args.anchors_path) num_anchors = len(anchors) # get freeze level according to CLI option if args.weights_path: freeze_level = 0 else: freeze_level = 1 if args.freeze_level is not None: freeze_level = args.freeze_level # callbacks for training process logging = TensorBoard(log_dir=log_dir, histogram_freq=0, write_graph=False, write_grads=False, write_images=False, update_freq='batch') checkpoint = ModelCheckpoint(os.path.join( log_dir, 'ep{epoch:03d}-loss{loss:.3f}-val_loss{val_loss:.3f}.h5'), monitor='val_loss', mode='min', verbose=1, save_weights_only=False, save_best_only=True, period=1) reduce_lr = ReduceLROnPlateau(monitor='val_loss', factor=0.5, mode='min', patience=10, verbose=1, cooldown=0, min_lr=1e-10) early_stopping = EarlyStopping(monitor='val_loss', min_delta=0, patience=50, verbose=1, mode='min') terminate_on_nan = TerminateOnNaN() callbacks = [ logging, checkpoint, reduce_lr, early_stopping, terminate_on_nan ] # 데이터셋 로딩 dataset = get_dataset(annotation_file) if args.val_annotation_file: val_dataset = get_dataset(args.val_annotation_file) num_train = len(dataset) num_val = len(val_dataset) dataset.extend(val_dataset) else: val_split = args.val_split num_val = int(len(dataset) * val_split) num_train = len(dataset) - num_val # assign multiscale interval if args.multiscale: rescale_interval = args.rescale_interval else: rescale_interval = -1 #Doesn't rescale # model input shape check input_shape = args.model_image_size assert (input_shape[0] % 32 == 0 and input_shape[1] % 32 == 0), 'model_image_size should be multiples of 32' # 모델종류에 따른 data generator 및 모델 생성 if num_anchors == 9: # YOLOv3 use 9 anchors get_train_model = get_yolo3_train_model data_generator = yolo3_data_generator_wrapper tiny_version = False elif num_anchors == 6: # Tiny YOLOv3 use 6 anchors get_train_model = get_yolo3_train_model data_generator = yolo3_data_generator_wrapper tiny_version = True elif num_anchors == 5: # YOLOv2 use 5 anchors get_train_model = get_yolo2_train_model data_generator = yolo2_data_generator_wrapper tiny_version = False else: raise ValueError('Unsupported anchors number') # prepare online evaluation callback if args.eval_online: eval_callback = EvalCallBack( args.model_type, dataset[num_train:], anchors, class_names, args.model_image_size, args.model_pruning, log_dir, eval_epoch_interval=args.eval_epoch_interval, save_eval_checkpoint=args.save_eval_checkpoint, elim_grid_sense=args.elim_grid_sense) callbacks.append(eval_callback) # prepare train/val data shuffle callback if args.data_shuffle: shuffle_callback = DatasetShuffleCallBack(dataset) callbacks.append(shuffle_callback) # prepare model pruning config pruning_end_step = np.ceil(1.0 * num_train / args.batch_size).astype( np.int32) * args.total_epoch if args.model_pruning: pruning_callbacks = [ sparsity.UpdatePruningStep(), sparsity.PruningSummaries(log_dir=log_dir, profile_batch=0) ] callbacks = callbacks + pruning_callbacks # prepare optimizer optimizer = get_optimizer(args.optimizer, args.learning_rate, decay_type=None) # support multi-gpu training if args.gpu_num >= 2: # devices_list=["/gpu:0", "/gpu:1"] devices_list = ["/gpu:{}".format(n) for n in range(args.gpu_num)] strategy = tf.distribute.MirroredStrategy(devices=devices_list) print('Number of devices: {}'.format(strategy.num_replicas_in_sync)) with strategy.scope(): # get multi-gpu train model model = get_train_model(args.model_type, anchors, num_classes, weights_path=args.weights_path, freeze_level=freeze_level, optimizer=optimizer, label_smoothing=args.label_smoothing, elim_grid_sense=args.elim_grid_sense, model_pruning=args.model_pruning, pruning_end_step=pruning_end_step) else: # get normal train model model = get_train_model(args.model_type, anchors, num_classes, weights_path=args.weights_path, freeze_level=freeze_level, optimizer=optimizer, label_smoothing=args.label_smoothing, elim_grid_sense=args.elim_grid_sense, model_pruning=args.model_pruning, pruning_end_step=pruning_end_step) model.summary() # Transfer training some epochs with frozen layers first if needed, to get a stable loss. initial_epoch = args.init_epoch epochs = initial_epoch + args.transfer_epoch print("Transfer training stage") print( 'Train on {} samples, val on {} samples, with batch size {}, input_shape {}.' .format(num_train, num_val, args.batch_size, input_shape)) # 성능향상을 위해 초반 일부 epoch은 Transfer Learning 진행 (Initial Epoch ~ Transfer Epoch) model.fit_generator( data_generator(dataset[:num_train], args.batch_size, input_shape, anchors, num_classes, args.enhance_augment, rescale_interval, multi_anchor_assign=args.multi_anchor_assign), steps_per_epoch=max(1, num_train // args.batch_size), #validation_data=val_data_generator, validation_data=data_generator( dataset[num_train:], args.batch_size, input_shape, anchors, num_classes, multi_anchor_assign=args.multi_anchor_assign), validation_steps=max(1, num_val // args.batch_size), epochs=epochs, initial_epoch=initial_epoch, #verbose=1, workers=1, use_multiprocessing=False, max_queue_size=10, callbacks=callbacks) # Wait 2 seconds for next stage time.sleep(2) if args.decay_type: # rebuild optimizer to apply learning rate decay, only after # unfreeze all layers callbacks.remove(reduce_lr) steps_per_epoch = max(1, num_train // args.batch_size) decay_steps = steps_per_epoch * (args.total_epoch - args.init_epoch - args.transfer_epoch) optimizer = get_optimizer(args.optimizer, args.learning_rate, decay_type=args.decay_type, decay_steps=decay_steps) # Unfreeze the whole network for further tuning # NOTE: more GPU memory is required after unfreezing the body print("Unfreeze and continue training, to fine-tune.") if args.gpu_num >= 2: with strategy.scope(): for i in range(len(model.layers)): model.layers[i].trainable = True model.compile(optimizer=optimizer, loss={ 'yolo_loss': lambda y_true, y_pred: y_pred }) # recompile to apply the change else: for i in range(len(model.layers)): model.layers[i].trainable = True model.compile(optimizer=optimizer, loss={ 'yolo_loss': lambda y_true, y_pred: y_pred }) # recompile to apply the change print( 'Train on {} samples, val on {} samples, with batch size {}, input_shape {}.' .format(num_train, num_val, args.batch_size, input_shape)) # Transfer Learning 이후 나머지 Epoch에 대하여 학습 진행 (Transfer Epoch ~ Total Epoch) # 이 부분이 필요없거나 학습 시간이 너무 오래 걸릴 경우 Total Epoch을 Transfer와 동일하게 두고, 아래 학습을 진행하지 않고 넘어갈 수 있음 # 본인 컴퓨터 사양에 맞춰서 진행 model.fit_generator( data_generator(dataset[:num_train], args.batch_size, input_shape, anchors, num_classes, args.enhance_augment, rescale_interval, multi_anchor_assign=args.multi_anchor_assign), steps_per_epoch=max(1, num_train // args.batch_size), #validation_data=val_data_generator, validation_data=data_generator( dataset[num_train:], args.batch_size, input_shape, anchors, num_classes, multi_anchor_assign=args.multi_anchor_assign), validation_steps=max(1, num_val // args.batch_size), epochs=args.total_epoch, initial_epoch=epochs, #verbose=1, workers=1, use_multiprocessing=False, max_queue_size=10, callbacks=callbacks) # Finally store model if args.model_pruning: model = sparsity.strip_pruning(model) model.save(os.path.join(log_dir, 'trained_final.h5'))
def main(args): annotation_file = args.annotation_file classes_path = args.classes_path class_names = get_classes(classes_path) num_classes = len(class_names) anchors = get_anchors(args.anchors_path) num_anchors = len(anchors) log_dir_path = args.log_directory try: log_dir = os.path.join('logs', log_dir_path) except TypeError: date_now = datetime.now() log_dir_folder_name = f'{date_now.strftime("%Y_%m_%d_%H%M%S")}_{args.model_type}_TransferEp_{args.transfer_epoch}_TotalEP_{args.total_epoch}' log_dir = os.path.realpath(os.path.join( 'logs', log_dir_folder_name )) # get freeze level according to CLI option if args.weights_path: freeze_level = 0 else: freeze_level = 1 if args.freeze_level is not None: freeze_level = args.freeze_level # How many percentage of layers to unfreeze in fine tuning unfreeze_level = args.unfreeze_level # callbacks for training process logging = TensorBoard(log_dir=log_dir, histogram_freq=0, write_graph=False, write_grads=False, write_images=False, update_freq='batch') checkpoint = ModelCheckpoint( filepath=log_dir + os.sep + 'ep{epoch:03d}-loss{loss:.3f}-val_loss{val_loss:.3f}.h5', monitor='val_loss', mode='min', verbose=1, save_weights_only=False, save_best_only=True, period=1 ) reduce_lr = ReduceLROnPlateau( monitor='val_loss', factor=0.5, mode='min', patience=10, verbose=1, cooldown=0, min_lr=1e-10 ) early_stopping = EarlyStopping(monitor='val_loss', min_delta=0, patience=50, verbose=1, mode='min') terminate_on_nan = TerminateOnNaN() callbacks = [logging, checkpoint, reduce_lr, early_stopping, terminate_on_nan] # get train&val dataset dataset = get_dataset(annotation_file) if args.val_annotation_file: val_dataset = get_dataset(args.val_annotation_file) num_train = len(dataset) num_val = len(val_dataset) dataset.extend(val_dataset) else: val_split = args.val_split num_val = int(len(dataset) * val_split) num_train = len(dataset) - num_val # assign multiscale interval if args.multiscale: rescale_interval = args.rescale_interval else: rescale_interval = -1 # Doesn't rescale # model input shape check input_shape = args.model_image_size assert (input_shape[0] % 32 == 0 and input_shape[1] % 32 == 0), 'model_image_size should be multiples of 32' # get different model type & train&val data generator if num_anchors == 9: # YOLOv3 use 9 anchors get_train_model = get_yolo3_train_model data_generator = yolo3_data_generator_wrapper # tf.keras.Sequence style data generator # train_data_generator = Yolo3DataGenerator(dataset[:num_train], args.batch_size, input_shape, anchors, num_classes, args.enhance_augment, rescale_interval, args.multi_anchor_assign) # val_data_generator = Yolo3DataGenerator(dataset[num_train:], args.batch_size, input_shape, anchors, num_classes, multi_anchor_assign=args.multi_anchor_assign) tiny_version = False elif num_anchors == 6: # Tiny YOLOv3 use 6 anchors get_train_model = get_yolo3_train_model data_generator = yolo3_data_generator_wrapper # tf.keras.Sequence style data generator # train_data_generator = Yolo3DataGenerator(dataset[:num_train], args.batch_size, input_shape, anchors, num_classes, args.enhance_augment, rescale_interval, args.multi_anchor_assign) # val_data_generator = Yolo3DataGenerator(dataset[num_train:], args.batch_size, input_shape, anchors, num_classes, multi_anchor_assign=args.multi_anchor_assign) tiny_version = True elif num_anchors == 5: # YOLOv2 use 5 anchors get_train_model = get_yolo2_train_model data_generator = yolo2_data_generator_wrapper # tf.keras.Sequence style data generator # train_data_generator = Yolo2DataGenerator(dataset[:num_train], args.batch_size, input_shape, anchors, num_classes, args.enhance_augment, rescale_interval) # val_data_generator = Yolo2DataGenerator(dataset[num_train:], args.batch_size, input_shape, anchors, num_classes) tiny_version = False else: raise ValueError('Unsupported anchors number') # prepare online evaluation callback if args.eval_online: eval_callback = EvalCallBack( model_type=args.model_type, annotation_lines=dataset[num_train:], anchors=anchors, class_names=class_names, model_image_size=args.model_image_size, model_pruning=args.model_pruning, log_dir=log_dir, eval_epoch_interval=args.eval_epoch_interval, save_eval_checkpoint=args.save_eval_checkpoint, elim_grid_sense=args.elim_grid_sense ) callbacks.append(eval_callback) # prepare train/val data shuffle callback if args.data_shuffle: shuffle_callback = DatasetShuffleCallBack(dataset) callbacks.append(shuffle_callback) # prepare model pruning config pruning_end_step = np.ceil(1.0 * num_train / args.batch_size).astype(np.int32) * args.total_epoch if args.model_pruning: pruning_callbacks = [sparsity.UpdatePruningStep(), sparsity.PruningSummaries(log_dir=log_dir, profile_batch=0)] callbacks = callbacks + pruning_callbacks # prepare optimizer optimizer = get_optimizer(args.optimizer, args.learning_rate, decay_type=None) # support multi-gpu training if args.gpu_num >= 2: # devices_list=["/gpu:0", "/gpu:1"] devices_list = ["/gpu:{}".format(n) for n in range(args.gpu_num)] strategy = tf.distribute.MirroredStrategy(devices=devices_list) print('Number of devices: {}'.format(strategy.num_replicas_in_sync)) with strategy.scope(): # get multi-gpu train model model = get_train_model( model_type=args.model_type, anchors=anchors, num_classes=num_classes, weights_path=args.weights_path, freeze_level=freeze_level, optimizer=optimizer, label_smoothing=args.label_smoothing, elim_grid_sense=args.elim_grid_sense, model_pruning=args.model_pruning, pruning_end_step=pruning_end_step ) else: # get normal train model model = get_train_model( model_type=args.model_type, anchors=anchors, num_classes=num_classes, weights_path=args.weights_path, freeze_level=freeze_level, optimizer=optimizer, label_smoothing=args.label_smoothing, elim_grid_sense=args.elim_grid_sense, model_pruning=args.model_pruning, pruning_end_step=pruning_end_step ) if args.show_history: model.summary() layers_count = len(model.layers) print(f'Total layers: {layers_count}') # Transfer training some epochs with frozen layers first if needed, to get a stable loss. initial_epoch = args.init_epoch epochs = initial_epoch + args.transfer_epoch print("Transfer training stage") print('Train on {} samples, val on {} samples, with batch size {}, input_shape {}.'.format(num_train, num_val, args.batch_size, input_shape)) # model.fit_generator(train_data_generator, """ Transfer training steps, train with freeze layers """ model.fit( data_generator( annotation_lines=dataset[:num_train], batch_size=args.batch_size, input_shape=input_shape, anchors=anchors, num_classes=num_classes, enhance_augment=args.enhance_augment, rescale_interval=rescale_interval, multi_anchor_assign=args.multi_anchor_assign ), steps_per_epoch=max(1, num_train // args.batch_size), # validation_data=val_data_generator, validation_data=data_generator( annotation_lines=dataset[num_train:], batch_size=args.batch_size, input_shape=input_shape, anchors=anchors, num_classes=num_classes, multi_anchor_assign=args.multi_anchor_assign ), validation_steps=max(1, num_val // args.batch_size), epochs=epochs, initial_epoch=initial_epoch, # verbose=1, workers=1, use_multiprocessing=False, max_queue_size=10, callbacks=callbacks ) # Wait 2 seconds for next stage time.sleep(2) if args.decay_type: # rebuild optimizer to apply learning rate decay, only after # unfreeze all layers callbacks.remove(reduce_lr) steps_per_epoch = max(1, num_train // args.batch_size) decay_steps = steps_per_epoch * (args.total_epoch - args.init_epoch - args.transfer_epoch) optimizer = get_optimizer(args.optimizer, args.learning_rate, decay_type=args.decay_type, decay_steps=decay_steps) # Unfreeze the whole network for further tuning # NOTE: more GPU memory is required after unfreezing the body fine_tune_layers = int(layers_count * unfreeze_level) print(f"Unfreeze {unfreeze_level * 100}% of layers and continue training, to fine-tune.") print(f"Unfroze {fine_tune_layers} layers of {layers_count}") if args.gpu_num >= 2: with strategy.scope(): for i in range(layers_count - fine_tune_layers, layers_count): model.layers[i].trainable = True model.compile(optimizer=optimizer, loss={'yolo_loss': lambda y_true, y_pred: y_pred}) # recompile to apply the change else: for i in range(layers_count - fine_tune_layers, layers_count): model.layers[i].trainable = True model.compile(optimizer=optimizer, loss={'yolo_loss': lambda y_true, y_pred: y_pred}) # recompile to apply the change print('Train on {} samples, val on {} samples, with batch size {}, input_shape {}.'.format(num_train, num_val, args.batch_size, input_shape)) """ Fine-tuning steps, more memory will be used. LR (Learning Rate) will be decayed """ # model.fit_generator(train_data_generator, model.fit( # The YOLO data augmentation generator tool data_generator( annotation_lines=dataset[:num_train], batch_size=args.batch_size, input_shape=input_shape, anchors=anchors, num_classes=num_classes, enhance_augment=args.enhance_augment, rescale_interval=rescale_interval, multi_anchor_assign=args.multi_anchor_assign ), steps_per_epoch=max(1, num_train // args.batch_size), # validation_data=val_data_generator, # Validation generator validation_data=data_generator( annotation_lines=dataset[num_train:], batch_size=args.batch_size, input_shape=input_shape, anchors=anchors, num_classes=num_classes, multi_anchor_assign=args.multi_anchor_assign ), validation_steps=max(1, num_val // args.batch_size), epochs=args.total_epoch, initial_epoch=epochs, # verbose=1, workers=1, use_multiprocessing=False, max_queue_size=10, callbacks=callbacks ) # Finally store model if args.model_pruning: model = sparsity.strip_pruning(model) model.save(os.path.join(log_dir, 'trained_final.h5'))
def main(args): annotation_file = args.annotation_file log_dir = os.path.join('logs', '000') classes_path = args.classes_path class_names = get_classes(classes_path) num_classes = len(class_names) anchors = get_anchors(args.anchors_path) num_anchors = len(anchors) # get freeze level according to CLI option if args.weights_path: freeze_level = 0 else: freeze_level = 1 if args.freeze_level is not None: freeze_level = args.freeze_level # callbacks for training process logging = TensorBoard(log_dir=log_dir, histogram_freq=0, write_graph=False, write_grads=False, write_images=False, update_freq='batch') checkpoint = ModelCheckpoint(os.path.join( log_dir, 'ep{epoch:03d}-loss{loss:.3f}-val_loss{val_loss:.3f}.h5'), monitor='val_loss', mode='min', verbose=1, save_weights_only=False, save_best_only=True, period=1) reduce_lr = ReduceLROnPlateau(monitor='val_loss', factor=0.5, mode='min', patience=10, verbose=1, cooldown=0, min_lr=1e-10) early_stopping = EarlyStopping(monitor='val_loss', min_delta=0, patience=50, verbose=1, mode='min') terminate_on_nan = TerminateOnNaN() callbacks = [ logging, checkpoint, reduce_lr, early_stopping, terminate_on_nan ] # get train&val dataset dataset = get_dataset(annotation_file) if args.val_annotation_file: val_dataset = get_dataset(args.val_annotation_file) num_train = len(dataset) num_val = len(val_dataset) dataset.extend(val_dataset) else: val_split = args.val_split num_val = int(len(dataset) * val_split) num_train = len(dataset) - num_val # assign multiscale interval if args.multiscale: rescale_interval = args.rescale_interval else: rescale_interval = -1 #Doesn't rescale # model input shape check input_shape = args.model_image_size assert (input_shape[0] % 32 == 0 and input_shape[1] % 32 == 0), 'model_image_size should be multiples of 32' # get different model type & train&val data generator if args.model_type.startswith( 'scaled_yolo4_') or args.model_type.startswith('yolo5_'): # Scaled-YOLOv4 & YOLOv5 entrance, use yolo5 submodule but now still yolo3 data generator # TODO: create new yolo5 data generator to apply YOLOv5 anchor assignment get_train_model = get_yolo5_train_model data_generator = yolo5_data_generator_wrapper # tf.keras.Sequence style data generator #train_data_generator = Yolo5DataGenerator(dataset[:num_train], args.batch_size, input_shape, anchors, num_classes, args.enhance_augment, rescale_interval, args.multi_anchor_assign) #val_data_generator = Yolo5DataGenerator(dataset[num_train:], args.batch_size, input_shape, anchors, num_classes, multi_anchor_assign=args.multi_anchor_assign) tiny_version = False elif args.model_type.startswith('yolo3_') or args.model_type.startswith( 'yolo4_'): #if num_anchors == 9: # YOLOv3 & v4 entrance, use 9 anchors get_train_model = get_yolo3_train_model data_generator = yolo3_data_generator_wrapper # tf.keras.Sequence style data generator #train_data_generator = Yolo3DataGenerator(dataset[:num_train], args.batch_size, input_shape, anchors, num_classes, args.enhance_augment, rescale_interval, args.multi_anchor_assign) #val_data_generator = Yolo3DataGenerator(dataset[num_train:], args.batch_size, input_shape, anchors, num_classes, multi_anchor_assign=args.multi_anchor_assign) tiny_version = False elif args.model_type.startswith( 'tiny_yolo3_') or args.model_type.startswith('tiny_yolo4_'): #elif num_anchors == 6: # Tiny YOLOv3 & v4 entrance, use 6 anchors get_train_model = get_yolo3_train_model data_generator = yolo3_data_generator_wrapper # tf.keras.Sequence style data generator #train_data_generator = Yolo3DataGenerator(dataset[:num_train], args.batch_size, input_shape, anchors, num_classes, args.enhance_augment, rescale_interval, args.multi_anchor_assign) #val_data_generator = Yolo3DataGenerator(dataset[num_train:], args.batch_size, input_shape, anchors, num_classes, multi_anchor_assign=args.multi_anchor_assign) tiny_version = True elif args.model_type.startswith('yolo2_') or args.model_type.startswith( 'tiny_yolo2_'): #elif num_anchors == 5: # YOLOv2 & Tiny YOLOv2 use 5 anchors get_train_model = get_yolo2_train_model data_generator = yolo2_data_generator_wrapper # tf.keras.Sequence style data generator #train_data_generator = Yolo2DataGenerator(dataset[:num_train], args.batch_size, input_shape, anchors, num_classes, args.enhance_augment, rescale_interval) #val_data_generator = Yolo2DataGenerator(dataset[num_train:], args.batch_size, input_shape, anchors, num_classes) tiny_version = False else: raise ValueError('Unsupported model type') # prepare online evaluation callback if args.eval_online: eval_callback = EvalCallBack( args.model_type, dataset[num_train:], anchors, class_names, args.model_image_size, args.model_pruning, log_dir, eval_epoch_interval=args.eval_epoch_interval, save_eval_checkpoint=args.save_eval_checkpoint, elim_grid_sense=args.elim_grid_sense) callbacks.append(eval_callback) # prepare train/val data shuffle callback if args.data_shuffle: shuffle_callback = DatasetShuffleCallBack(dataset) callbacks.append(shuffle_callback) # prepare model pruning config pruning_end_step = np.ceil(1.0 * num_train / args.batch_size).astype( np.int32) * args.total_epoch if args.model_pruning: pruning_callbacks = [ sparsity.UpdatePruningStep(), sparsity.PruningSummaries(log_dir=log_dir, profile_batch=0) ] callbacks = callbacks + pruning_callbacks # prepare optimizer optimizer = get_optimizer(args.optimizer, args.learning_rate, average_type=None, decay_type=None) # support multi-gpu training if args.gpu_num >= 2: # devices_list=["/gpu:0", "/gpu:1"] devices_list = ["/gpu:{}".format(n) for n in range(args.gpu_num)] strategy = tf.distribute.MirroredStrategy(devices=devices_list) print('Number of devices: {}'.format(strategy.num_replicas_in_sync)) with strategy.scope(): # get multi-gpu train model model = get_train_model(args.model_type, anchors, num_classes, weights_path=args.weights_path, freeze_level=freeze_level, optimizer=optimizer, label_smoothing=args.label_smoothing, elim_grid_sense=args.elim_grid_sense, model_pruning=args.model_pruning, pruning_end_step=pruning_end_step) else: # get normal train model model = get_train_model(args.model_type, anchors, num_classes, weights_path=args.weights_path, freeze_level=freeze_level, optimizer=optimizer, label_smoothing=args.label_smoothing, elim_grid_sense=args.elim_grid_sense, model_pruning=args.model_pruning, pruning_end_step=pruning_end_step) model.summary() # Transfer training some epochs with frozen layers first if needed, to get a stable loss. initial_epoch = args.init_epoch epochs = initial_epoch + args.transfer_epoch print("Transfer training stage") print( 'Train on {} samples, val on {} samples, with batch size {}, input_shape {}.' .format(num_train, num_val, args.batch_size, input_shape)) #model.fit_generator(train_data_generator, model.fit_generator( data_generator(dataset[:num_train], args.batch_size, input_shape, anchors, num_classes, args.enhance_augment, rescale_interval, multi_anchor_assign=args.multi_anchor_assign), steps_per_epoch=max(1, num_train // args.batch_size), #validation_data=val_data_generator, validation_data=data_generator( dataset[num_train:], args.batch_size, input_shape, anchors, num_classes, multi_anchor_assign=args.multi_anchor_assign), validation_steps=max(1, num_val // args.batch_size), epochs=epochs, initial_epoch=initial_epoch, #verbose=1, workers=1, use_multiprocessing=False, max_queue_size=10, callbacks=callbacks) # Wait 2 seconds for next stage time.sleep(2) if args.decay_type or args.average_type: # rebuild optimizer to apply learning rate decay or weights averager, # only after unfreeze all layers if args.decay_type: callbacks.remove(reduce_lr) if args.average_type == 'ema' or args.average_type == 'swa': # weights averager need tensorflow-addons, # which request TF 2.x and have version compatibility import tensorflow_addons as tfa callbacks.remove(checkpoint) avg_checkpoint = tfa.callbacks.AverageModelCheckpoint( filepath=os.path.join( log_dir, 'ep{epoch:03d}-loss{loss:.3f}-val_loss{val_loss:.3f}.h5'), update_weights=True, monitor='val_loss', mode='min', verbose=1, save_weights_only=False, save_best_only=True, period=1) callbacks.append(avg_checkpoint) steps_per_epoch = max(1, num_train // args.batch_size) decay_steps = steps_per_epoch * (args.total_epoch - args.init_epoch - args.transfer_epoch) optimizer = get_optimizer(args.optimizer, args.learning_rate, average_type=args.average_type, decay_type=args.decay_type, decay_steps=decay_steps) # Unfreeze the whole network for further tuning # NOTE: more GPU memory is required after unfreezing the body print("Unfreeze and continue training, to fine-tune.") if args.gpu_num >= 2: with strategy.scope(): for i in range(len(model.layers)): model.layers[i].trainable = True model.compile(optimizer=optimizer, loss={ 'yolo_loss': lambda y_true, y_pred: y_pred }) # recompile to apply the change else: for i in range(len(model.layers)): model.layers[i].trainable = True model.compile(optimizer=optimizer, loss={ 'yolo_loss': lambda y_true, y_pred: y_pred }) # recompile to apply the change print( 'Train on {} samples, val on {} samples, with batch size {}, input_shape {}.' .format(num_train, num_val, args.batch_size, input_shape)) #model.fit_generator(train_data_generator, model.fit_generator( data_generator(dataset[:num_train], args.batch_size, input_shape, anchors, num_classes, args.enhance_augment, rescale_interval, multi_anchor_assign=args.multi_anchor_assign), steps_per_epoch=max(1, num_train // args.batch_size), #validation_data=val_data_generator, validation_data=data_generator( dataset[num_train:], args.batch_size, input_shape, anchors, num_classes, multi_anchor_assign=args.multi_anchor_assign), validation_steps=max(1, num_val // args.batch_size), epochs=args.total_epoch, initial_epoch=epochs, #verbose=1, workers=1, use_multiprocessing=False, max_queue_size=10, callbacks=callbacks) # Finally store model if args.model_pruning: model = sparsity.strip_pruning(model) model.save(os.path.join(log_dir, 'trained_final.h5'))
def main(args): model_type = "yolo3_darknet_spp" # yolo3_darknet_spp, yolo3_darknet current_dir = os.path.dirname(__file__) + "/" print("current_dir == ", current_dir) annotation_file = current_dir + "sample/trainval/train.txt" val_annotation_file = current_dir + "sample/trainval/val.txt" classes_path = current_dir + "sample/trainval/train_classes.txt" anchors_path = current_dir + "sample/trainval/yolo_anchors.txt" weights_path = current_dir + "weights/yolov3-spp.h5" load_weights_path = None # None or "{weights path}" is_one_stage_train = True learning_rate_1 = 1e-4 learning_rate_2 = 1e-5 epoch_1 = args.max_epochs_1 epoch_2 = args.max_epochs_2 batch_size_1 = args.batch_size_1 batch_size_2 = args.batch_size_2 freeze_level = 2 model_image_size = (416, 416) val_split = 0.1 label_smoothing = 0 enhance_augment = None # enhance data augmentation type (None/mosaic) rescale_interval = 0 # Number of iteration(batches) interval to rescale input size, default=10 log_dir = os.path.join('logs', '20200602') class_names = get_classes(classes_path) num_classes = len(class_names) anchors = get_anchors(anchors_path) logging = TensorBoard(log_dir=log_dir, update_freq='batch') checkpoint = ModelCheckpoint(os.path.join( log_dir, 'ep{epoch:03d}-loss{loss:.3f}-val_loss{val_loss:.3f}.h5'), monitor='val_loss', verbose=1, save_weights_only=True, save_best_only=True, period=1) reduce_lr = ReduceLROnPlateau(monitor='val_loss', factor=0.1, patience=10, verbose=1, cooldown=0, min_lr=1e-10) early_stopping = EarlyStopping(monitor='val_loss', min_delta=0, patience=50, verbose=1) # terminate_on_nan = TerminateOnNaN() callbacks = [ logging, checkpoint, reduce_lr, early_stopping, ModertFileToObs(log_dir, args) ] # callbacks = [logging, checkpoint, reduce_lr] # get train&val dataset dataset = get_dataset(annotation_file) dataset = [current_dir + d for d in dataset] if val_annotation_file != "": val_dataset = get_dataset(val_annotation_file) num_train = len(dataset) num_val = len(val_dataset) dataset.extend(val_dataset) else: val_split = val_split num_val = int(len(dataset) * val_split) num_train = len(dataset) - num_val # num_val = 100 # num_train = 200 # model input shape check input_shape = model_image_size assert (input_shape[0] % 32 == 0 and input_shape[1] % 32 == 0), 'Multiples of 32 required' get_train_model = get_yolo3_train_model data_generator = yolo3_data_generator_wrapper # get train model model = get_train_model(model_type, anchors, num_classes, input_shape, weights_path=weights_path, freeze_level=freeze_level, label_smoothing=label_smoothing) if load_weights_path: model.load_weights(load_weights_path) print("reload weights: {}".format(load_weights_path)) if is_one_stage_train: model.compile(optimizer=get_optimizer(learning_rate_1), loss={ 'yolo_loss': lambda y_true, y_pred: y_pred }) print( 'One stage Train on {} samples, val on {} samples, with batch size {}, ' 'input_shape {}.'.format(num_train, num_val, batch_size_1, input_shape)) model.fit_generator( data_generator(dataset[:num_train], batch_size_1, input_shape, anchors, num_classes, enhance_augment), steps_per_epoch=max(1, num_train // batch_size_1), validation_data=data_generator(dataset[:num_val], batch_size_1, input_shape, anchors, num_classes), validation_steps=max(1, num_val // batch_size_1), epochs=epoch_1, initial_epoch=0, workers=1, use_multiprocessing=False, max_queue_size=10, callbacks=callbacks) model.save_weights(os.path.join(log_dir, 'trained_weights_stage_1.h5')) if True: print("Unfreeze and continue training, to fine-tune.") for i in range(len(model.layers)): model.layers[i].trainable = True model.compile(optimizer=get_optimizer(learning_rate_2), loss={ 'yolo_loss': lambda y_true, y_pred: y_pred }) print( 'Two stage Train on {} samples, val on {} samples, with batch size {}, input_shape {}.' .format(num_train, num_val, batch_size_2, input_shape)) model.fit_generator(data_generator(dataset[:num_train], batch_size_2, input_shape, anchors, num_classes, enhance_augment, rescale_interval), steps_per_epoch=max(1, num_train // batch_size_2), validation_data=data_generator( dataset[:num_val], batch_size_2, input_shape, anchors, num_classes), validation_steps=max(1, num_val // batch_size_2), epochs=epoch_2, initial_epoch=epoch_1, workers=1, use_multiprocessing=False, max_queue_size=10, callbacks=callbacks) model.save_weights(os.path.join(log_dir, 'trained_weights_final.h5')) gen_model_dir(log_dir, args, classes_path, anchors_path)
def main(): parser = argparse.ArgumentParser( description='validate YOLO model (h5/pb/tflite/mnn) with image') parser.add_argument('--model_path', help='model file to predict', type=str, required=True) parser.add_argument('--image_file', help='image file to predict', type=str, required=True) parser.add_argument('--anchors_path', help='path to anchor definitions', type=str, required=True) parser.add_argument( '--classes_path', help='path to class definitions, default ../configs/voc_classes.txt', type=str, default='../configs/voc_classes.txt') parser.add_argument( '--model_image_size', help='model image input size as <num>x<num>, default 416x416', type=str, default='416x416') parser.add_argument('--loop_count', help='loop inference for certain times', type=int, default=1) parser.add_argument( '--custom_objects', required=False, type=str, help= "Custom objects in keras model (swish/tf). Separated with comma if more than one.", default=None) args = parser.parse_args() # param parse anchors = get_anchors(args.anchors_path) class_names = get_classes(args.classes_path) height, width = args.model_image_size.split('x') model_image_size = (int(height), int(width)) # support of tflite model if args.model_path.endswith('.tflite'): validate_yolo_model_tflite(args.model_path, args.image_file, anchors, class_names, args.loop_count) # support of MNN model elif args.model_path.endswith('.mnn'): validate_yolo_model_mnn(args.model_path, args.image_file, anchors, class_names, args.loop_count) # support of TF 1.x frozen pb model elif args.model_path.endswith('.pb'): validate_yolo_model_pb(args.model_path, args.image_file, anchors, class_names, model_image_size, args.loop_count) # normal keras h5 model elif args.model_path.endswith('.h5'): validate_yolo_model(args.model_path, args.custom_objects, args.image_file, anchors, class_names, model_image_size, args.loop_count) else: raise ValueError('invalid model file')