def load_model(model_name, weight_path, input_size, framework): assert model_name in ['yolov3_tiny', 'yolov3', 'yolov4'] NUM_CLASS = len(utils.read_class_names(cfg.YOLO.CLASSES)) if framework == 'tf': input_layer = tf.keras.layers.Input([input_size, input_size, 3]) if model_name == 'yolov3_tiny': feature_maps = YOLOv3_tiny(input_layer, NUM_CLASS) bbox_tensors = [] for i, fm in enumerate(feature_maps): bbox_tensor = ops.decode(fm, NUM_CLASS) bbox_tensors.append(bbox_tensor) model = tf.keras.Model(input_layer, bbox_tensors) elif model_name == 'yolov3': feature_maps = YOLOv3(input_layer, NUM_CLASS) bbox_tensors = [] for i, fm in enumerate(feature_maps): bbox_tensor = ops.decode(fm, NUM_CLASS) bbox_tensors.append(bbox_tensor) model = tf.keras.Model(input_layer, bbox_tensors) elif model_name == 'yolov4': feature_maps = YOLOv4(input_layer, NUM_CLASS) bbox_tensors = [] for i, fm in enumerate(feature_maps): bbox_tensor = ops.decode(fm, NUM_CLASS) bbox_tensors.append(bbox_tensor) model = tf.keras.Model(input_layer, bbox_tensors) else: model = None raise ValueError if weight_path.split(".")[-1] == "weights": if model_name == 'yolov3_tiny': utils.load_weights_tiny(model, weight_path) print('load yolo tiny 3') elif model_name == 'yolov3': utils.load_weights_v3(model, weight_path) print('load yolo 3') elif model_name == 'yolov4': utils.load_weights(model, weight_path) print('load yolo 4') else: raise ValueError else: model.load_weights(weight_path).expect_partial() print('Restoring weights from: %s ' % weight_path) return model
def transfer_tflite(model_name, weight_path, output, input_size): assert model_name in ['yolov3_tiny', 'yolov3', 'yolov4'] NUM_CLASS = len(utils.read_class_names(cfg.YOLO.CLASSES)) input_layer = tf.keras.layers.Input([input_size, input_size, 3]) if model_name == 'yolov3_tiny': feature_maps = YOLOv3_tiny(input_layer, NUM_CLASS) bbox_tensors = [] for i, fm in enumerate(feature_maps): bbox_tensor = ops.decode(fm, NUM_CLASS) bbox_tensors.append(bbox_tensor) model = tf.keras.Model(input_layer, bbox_tensors) elif model_name == 'yolov3': feature_maps = YOLOv3(input_layer, NUM_CLASS) bbox_tensors = [] for i, fm in enumerate(feature_maps): bbox_tensor = ops.decode(fm, NUM_CLASS) bbox_tensors.append(bbox_tensor) model = tf.keras.Model(input_layer, bbox_tensors) elif model_name == 'yolov4': feature_maps = YOLOv4(input_layer, NUM_CLASS) bbox_tensors = [] for i, fm in enumerate(feature_maps): bbox_tensor = ops.decode(fm, NUM_CLASS) bbox_tensors.append(bbox_tensor) model = tf.keras.Model(input_layer, bbox_tensors) else: model = None raise ValueError if weight_path.split(".")[-1] == "weights": if model_name == 'yolov3_tiny': utils.load_weights_tiny(model, weight_path) elif model_name == ' yolov3': utils.load_weights_v3(model, weight_path) elif model_name == 'yolov4': utils.load_weights(model, weight_path) else: raise ValueError else: model.load_weights(weight_path).expect_partial() print('Restoring weights from: %s ... ' % weight_path) converter = tf.lite.TFLiteConverter.from_keras_model(model) # converter.optimizations = [tf.lite.Optimize.OPTIMIZE_FOR_SIZE] tflite_model = converter.convert() open(output, 'wb').write(tflite_model)
if __name__ == '__main__': classes_path = 'data/coco_classes.txt' # model_path可以是'yolov4.h5'、'./weights/step00001000.h5'这些。 # model_path = 'yolov4.h5' model_path = './weights/step00001000.h5' # input_shape越大,精度会上升,但速度会下降。 # input_shape = (320, 320) input_shape = (416, 416) # input_shape = (608, 608) num_anchors = 3 all_classes = get_classes(classes_path) num_classes = len(all_classes) inputs = layers.Input(shape=(None, None, 3)) yolo = YOLOv4(inputs, num_classes, num_anchors) yolo.load_weights(model_path, by_name=True) _decode = Decode(0.05, 0.45, input_shape, yolo, all_classes) # detect images in test floder. for (root, dirs, files) in os.walk('images/test'): if files: start = time.time() for f in files: path = os.path.join(root, f) image = cv2.imread(path) image, boxes, scores, classes = _decode.detect_image( image, draw_image=True) cv2.imwrite('images/res/' + f, image) print('total time: {0:.6f}s'.format(time.time() - start))
def detect(model_name, weight_path, input_size, image_path, framework): assert model_name in ['yolov3_tiny', 'yolov3', 'yolov4'] if model_name == 'yolov3_tiny': STRIDES = np.array(cfg.YOLO.STRIDES_TINY) ANCHORS = utils.get_anchors(cfg.YOLO.ANCHORS_TINY, True) elif model_name == 'yolov3': STRIDES = np.array(cfg.YOLO.STRIDES) ANCHORS = utils.get_anchors(cfg.YOLO.ANCHORS_V3, False) elif model_name == 'yolov4': STRIDES = np.array(cfg.YOLO.STRIDES) ANCHORS = utils.get_anchors(cfg.YOLO.ANCHORS, False) else: raise ValueError NUM_CLASS = len(utils.read_class_names(cfg.YOLO.CLASSES)) XYSCALE = cfg.YOLO.XYSCALE original_image = cv2.imread(image_path) original_image = cv2.cvtColor(original_image, cv2.COLOR_BGR2RGB) original_image_size = original_image.shape[:2] image_data = utils.image_preprocess(np.copy(original_image), [input_size, input_size]) image_data = image_data[np.newaxis, ...].astype(np.float32) if framework == 'tf': input_layer = tf.keras.layers.Input([input_size, input_size, 3]) if model_name == 'yolov3_tiny': feature_maps = YOLOv3_tiny(input_layer, NUM_CLASS) bbox_tensors = [] for i, fm in enumerate(feature_maps): bbox_tensor = ops.decode(fm, NUM_CLASS) bbox_tensors.append(bbox_tensor) model = tf.keras.Model(input_layer, bbox_tensors) elif model_name == 'yolov3': feature_maps = YOLOv3(input_layer, NUM_CLASS) bbox_tensors = [] for i, fm in enumerate(feature_maps): bbox_tensor = ops.decode(fm, NUM_CLASS) bbox_tensors.append(bbox_tensor) model = tf.keras.Model(input_layer, bbox_tensors) elif model_name == 'yolov4': feature_maps = YOLOv4(input_layer, NUM_CLASS) bbox_tensors = [] for i, fm in enumerate(feature_maps): bbox_tensor = ops.decode(fm, NUM_CLASS) bbox_tensors.append(bbox_tensor) model = tf.keras.Model(input_layer, bbox_tensors) else: model = None raise ValueError if weight_path.split(".")[-1] == "weights": if model_name == 'yolov3_tiny': utils.load_weights_tiny(model, weight_path) # utils.extract_weights_tiny(model, weight_path) print('load yolo tiny 3') elif model_name == 'yolov3': utils.load_weights_v3(model, weight_path) print('load yolo 3') elif model_name == 'yolov4': utils.load_weights(model, weight_path) print('load yolo 4') else: raise ValueError elif weight_path.split(".")[-1] == "npy": if model_name == 'yolov3_tiny': # utils.load_weights_tiny_npy(model, weight_path) print('load yolo tiny 3 npy') else: model.load_weights(weight_path) print('Restoring weights from: %s ' % weight_path) # weight = np.load('D:\\coursera\\YoLoSerirs\\checkpoint\\yolo3_tiny.npy', allow_pickle=True) # model.set_weights(weight) # model.summary() start_time = time.time() pred_bbox = model.predict(image_data) print(time.time() - start_time) else: # Load TFLite model and allocate tensors. interpreter = tf.lite.Interpreter(model_path=weight_path) interpreter.allocate_tensors() # Get input and output tensors. input_details = interpreter.get_input_details() output_details = interpreter.get_output_details() print(input_details) print(output_details) interpreter.set_tensor(input_details[0]['index'], image_data) start_time = time.time() interpreter.invoke() pred_bbox = [ interpreter.get_tensor(output_details[i]['index']) for i in range(len(output_details)) ] print(time.time() - start_time) if model_name == 'yolov4': pred_bbox = utils.postprocess_bbbox(pred_bbox, ANCHORS, STRIDES, XYSCALE) else: pred_bbox = utils.postprocess_bbbox(pred_bbox, ANCHORS, STRIDES) bboxes = utils.postprocess_boxes(pred_bbox, original_image_size, input_size, 0.5) bboxes = utils.nms(bboxes, 0.3, method='nms') image = visualize.draw_bbox(original_image, bboxes) image = Image.fromarray(image) image.show()
input_shape = (416, 416) # input_shape = (608, 608) # 验证时的分数阈值和nms_iou阈值 conf_thresh = 0.05 nms_thresh = 0.45 # 是否给图片画框。不画可以提速。读图片、后处理还可以继续优化。 draw_image = True # draw_image = False num_anchors = 3 all_classes = get_classes(classes_path) num_classes = len(all_classes) yolo = YOLOv4(num_classes, num_anchors) if torch.cuda.is_available(): # 如果有gpu可用,模型(包括了权重weight)存放在gpu显存里 yolo = yolo.cuda() yolo.load_state_dict(torch.load(model_path)) yolo.eval( ) # 必须调用model.eval()来设置dropout和batch normalization layers在运行推理前,切换到评估模式. 不这样做的化会产生不一致的推理结果. _decode = Decode(conf_thresh, nms_thresh, input_shape, yolo, all_classes) if not os.path.exists('images/res/'): os.mkdir('images/res/') path_dir = os.listdir('images/test') # warm up if use_gpu: for k, filename in enumerate(path_dir): image = cv2.imread('images/test/' + filename)
def train(model_name, weight_path, save_path, logdir=None): assert model_name in ['yolov3_tiny', 'yolov3', 'yolov4'] physical_devices = tf.config.experimental.list_physical_devices('GPU') if len(physical_devices) > 0: tf.config.experimental.set_memory_growth(physical_devices[0], True) 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) trainset = Dataset('train') testset = Dataset('test') 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 input_layer = tf.keras.layers.Input([cfg.TRAIN.INPUT_SIZE, cfg.TRAIN.INPUT_SIZE, 3]) if model_name=='yolov3_tiny': feature_maps = YOLOv3_tiny(input_layer, NUM_CLASS) bbox_tensors = [] for i, fm in enumerate(feature_maps): bbox_tensor = ops.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) elif model_name=='yolov3': feature_maps = YOLOv3(input_layer, NUM_CLASS) bbox_tensors = [] for i, fm in enumerate(feature_maps): bbox_tensor = ops.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) elif model_name=='yolov4': feature_maps = YOLOv4(input_layer, NUM_CLASS) bbox_tensors = [] for i, fm in enumerate(feature_maps): bbox_tensor = ops.decode_train(fm, NUM_CLASS, STRIDES, ANCHORS, i, XYSCALE) bbox_tensors.append(fm) bbox_tensors.append(bbox_tensor) model = tf.keras.Model(input_layer, bbox_tensors) else: raise ValueError # for name in ['conv2d_93', 'conv2d_101', 'conv2d_109']: # layer = model.get_layer(name) # print(layer.name, layer.output_shape) if weight_path: if weight_path.split(".")[-1] == "weights": if model_name == 'yolov3_tiny': utils.load_weights_tiny(model, weight_path) elif model_name=='yolov3': utils.load_weights_v3(model, weight_path) elif model_name=='yolov4': utils.load_weights(model, weight_path) else: raise ValueError else: model.load_weights(weight_path) print('Restoring weights from: %s ... ' % weight_path) optimizer = tf.keras.optimizers.Adam() if logdir: if os.path.exists(logdir): shutil.rmtree(logdir) writer = tf.summary.create_file_writer(logdir) else: writer = None 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 = ops.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" % (global_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()) # if writer: # # 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): 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 = ops.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 ['conv2d_93', 'conv2d_101', 'conv2d_109']: freeze = model.get_layer(name) ops.freeze_all(freeze) elif epoch >= first_stage_epochs: if isfreeze: isfreeze = False for name in ['conv2d_93', 'conv2d_101', 'conv2d_109']: freeze = model.get_layer(name) ops.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) if save_path: model.save_weights(save_path)
def __init__(self, batch_size: int, image_size: int): # setup anchors cfg.anchors.set_image_size(image_size) # dataset dataset_train = self.create_dataset_generator(dataset=cfg.dataset, mode=tfds.Split.TRAIN, image_size=image_size, batch_size=batch_size) dataset_val = self.create_dataset_generator(dataset=cfg.dataset, mode=tfds.Split.VALIDATION, image_size=image_size, batch_size=batch_size) self.dataset_train = dataset_train.get_dataset() self.dataset_val = dataset_val.get_dataset() self.class_names = self.create_class_names(dataset=cfg.dataset) # parameters self.batch_size = batch_size self.image_size = image_size self.buffer_size = cfg.buffer_size self.prefetch_size = cfg.prefetch_size self.num_class = dataset_train.num_class self.yolo_iou_threshold = cfg.yolo_iou_threshold self.yolo_score_threshold = cfg.yolo_score_threshold self.label_smoothing_factor = cfg.label_smoothing_factor self.lr_init = cfg.lr_init / self.batch_size self.lr_end = cfg.lr_end / self.batch_size self.warmup_epochs = cfg.warmup_epochs self.train_epochs = cfg.train_epochs self.warmup_steps = self.warmup_epochs * dataset_train.num_of_img / self.batch_size self.total_steps = self.train_epochs * dataset_train.num_of_img / self.batch_size self.step_to_log = cfg.step_to_log # define model and loss self.model = YOLOv4(num_class=self.num_class) self.lr_scheduler = WarmUpLinearCosineDecay( warmup_steps=self.warmup_steps, decay_steps=self.total_steps, initial_learning_rate=self.lr_init) self.optimizer = tf.keras.optimizers.Adam( learning_rate=self.lr_scheduler, clipnorm=1.0) self.checkpoint_dir = './checkpoints/yolov4_train.tf' self.ckpt = tf.train.Checkpoint(step=tf.Variable(1), optimizer=self.optimizer, net=self.model) self.manager = tf.train.CheckpointManager(self.ckpt, self.checkpoint_dir, max_to_keep=5) self.loss_fn = YOLOv4Loss( num_class=self.num_class, yolo_iou_threshold=self.yolo_iou_threshold, label_smoothing_factor=self.label_smoothing_factor, use_ciou_loss=True, use_focal_obj_loss=True) # metrics self.mAP = DetectionMAP(self.num_class) # summary writer self.current_time = datetime.datetime.now().strftime("%Y%m%d-%H%M%S") self.train_log_dir = 'logs/yolov4/train' self.val_log_dir = 'logs/yolov4/val' self.train_summary_writer = tf.summary.create_file_writer( self.train_log_dir) self.val_summary_writer = tf.summary.create_file_writer( self.val_log_dir)
param['keep_top_k'] = keep_top_k param['nms_top_k'] = nms_top_k param['num_classes'] = num_classes param['num_anchors'] = num_anchors # 输入字典 feed_vars = [('image', inputs), ('resize_shape', resize_shape), ('origin_shape', origin_shape)] feed_vars = OrderedDict(feed_vars) if postprocess == 'numpy_nms': param = None # 输入字典 feed_vars = [('image', inputs), ] feed_vars = OrderedDict(feed_vars) if algorithm == 'YOLOv4': if postprocess == 'fastnms': boxes, scores, classes = YOLOv4(inputs, num_classes, num_anchors, is_test=False, trainable=True, export=True, postprocess=postprocess, param=param) test_fetches = {'boxes': boxes, 'scores': scores, 'classes': classes, } if postprocess == 'multiclass_nms': pred = YOLOv4(inputs, num_classes, num_anchors, is_test=False, trainable=True, export=True, postprocess=postprocess, param=param) test_fetches = {'pred': pred, } elif algorithm == 'YOLOv3': backbone = Resnet50Vd() head = YOLOv3Head(keep_prob=1.0) # 一定要设置keep_prob=1.0, 为了得到一致的推理结果 yolov3 = YOLOv3(backbone, head) if postprocess == 'fastnms': boxes, scores, classes = yolov3(inputs, export=True, postprocess=postprocess, param=param) test_fetches = {'boxes': boxes, 'scores': scores, 'classes': classes, } if postprocess == 'multiclass_nms': pred = yolov3(inputs, export=True, postprocess=postprocess, param=param) test_fetches = {'pred': pred, } infer_prog = infer_prog.clone(for_test=True)
num_anchors = 3 all_classes = get_classes(classes_path) num_classes = len(all_classes) startup_prog = fluid.Program() eval_prog = fluid.Program() with fluid.program_guard(eval_prog, startup_prog): with fluid.unique_name.guard(): # 多尺度训练 inputs = P.data(name='input_1', shape=[-1, 3, -1, -1], append_batch_size=False, dtype='float32') output_l, output_m, output_s = YOLOv4(inputs, num_classes, num_anchors, is_test=False, trainable=True) eval_fetch_list = [output_l, output_m, output_s] eval_prog = eval_prog.clone(for_test=True) gpu_id = int(os.environ.get('FLAGS_selected_gpus', 0)) place = fluid.CUDAPlace(gpu_id) if use_gpu else fluid.CPUPlace() exe = fluid.Executor(place) exe.run(startup_prog) fluid.load(eval_prog, model_path, executor=exe) _decode = Decode(conf_thresh, nms_thresh, input_shape, exe, eval_prog, all_classes) test_dev(_decode, eval_fetch_list, images, test_pre_path, test_batch_size, draw_image)
cfg = TrainConfig() class_names = get_classes(cfg.classes_path) num_classes = len(class_names) _anchors = copy.deepcopy(cfg.anchors) num_anchors = len(cfg.anchor_masks[0]) # 每个输出层有几个先验框 _anchors = np.array(_anchors) _anchors = np.reshape(_anchors, (-1, num_anchors, 2)) _anchors = _anchors.astype(np.float32) # 步id,无需设置,会自动读。 iter_id = 0 # 多尺度训练 inputs = layers.Input(shape=(None, None, 3)) model_body = YOLOv4(inputs, num_classes, num_anchors) _decode = Decode(cfg.conf_thresh, cfg.nms_thresh, cfg.input_shape, model_body, class_names) # 模式。 0-从头训练,1-读取之前的模型继续训练(model_path可以是'yolov4.h5'、'./weights/step00001000.h5'这些。) pattern = cfg.pattern if pattern == 1: model_body.load_weights(cfg.model_path, by_name=True) # , skip_mismatch=True) strs = cfg.model_path.split('step') if len(strs) == 2: iter_id = int(strs[1][:8]) # 冻结,使得需要的显存减少。6G的卡建议这样配置。11G的卡建议不冻结。 # freeze_before = 'conv2d_60' # freeze_before = 'conv2d_72'
idx).get_tensor() tensor2 = fluid.global_scope().find_var('conv%.3d.conv.bias' % idx).get_tensor() tensor.set(w, place) tensor2.set(b, place) num_classes = 80 num_anchors = 3 use_gpu = False inputs = fluid.layers.data(name='input_1', shape=[-1, 3, -1, -1], append_batch_size=False, dtype='float32') pred_outs = YOLOv4(inputs, num_classes, num_anchors) # Create an executor using CPU as an example gpu_id = int(os.environ.get('FLAGS_selected_gpus', 0)) place = fluid.CUDAPlace(gpu_id) if use_gpu else fluid.CPUPlace() exe = fluid.Executor(place) exe.run(fluid.default_startup_program()) print('\nCopying...') for i in range(1, 94, 1): copy1(i, place) for i in range(95, 102, 1): copy1(i, place) for i in range(103, 110, 1): copy1(i, place)
draw_image = True # draw_image = False # 初始卷积核个数 initial_filters = 32 anchors = np.array([[[12, 16], [19, 36], [40, 28]], [[36, 75], [76, 55], [72, 146]], [[142, 110], [192, 243], [459, 401]]]) # 一些预处理 anchors = anchors.astype(np.float32) num_anchors = len(anchors[0]) # 每个输出层有几个先验框 all_classes = get_classes(classes_path) num_classes = len(all_classes) inputs = layers.Input(shape=(None, None, 3)) yolo = YOLOv4(inputs, num_classes, num_anchors, initial_filters, True, anchors, conf_thresh, nms_thresh, keep_top_k, nms_top_k) yolo.load_weights(model_path, by_name=True) if not os.path.exists('images/res/'): os.mkdir('images/res/') # 定义颜色 hsv_tuples = [(1.0 * x / num_classes, 1., 1.) for x in range(num_classes)] colors = list(map(lambda x: colorsys.hsv_to_rgb(*x), hsv_tuples)) colors = list( map(lambda x: (int(x[0] * 255), int(x[1] * 255), int(x[2] * 255)), colors)) random.seed(0) random.shuffle(colors) random.seed(None) path_dir = os.listdir('images/test')
def evaluate(model_name, weight_path): assert model_name in ['yolov3_tiny', 'yolov3', 'yolov4'] physical_devices = tf.config.experimental.list_physical_devices('GPU') if len(physical_devices) > 0: tf.config.experimental.set_memory_growth(physical_devices[0], True) 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) trainset = Dataset('train') 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 input_layer = tf.keras.layers.Input([cfg.TRAIN.INPUT_SIZE, cfg.TRAIN.INPUT_SIZE, 3]) if model_name=='yolov3_tiny': feature_maps = YOLOv3_tiny(input_layer, NUM_CLASS) bbox_tensors = [] for i, fm in enumerate(feature_maps): bbox_tensor = ops.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) elif model_name=='yolov3': feature_maps = YOLOv3(input_layer, NUM_CLASS) bbox_tensors = [] for i, fm in enumerate(feature_maps): bbox_tensor = ops.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) elif model_name=='yolov4': feature_maps = YOLOv4(input_layer, NUM_CLASS) bbox_tensors = [] for i, fm in enumerate(feature_maps): bbox_tensor = ops.decode_train(fm, NUM_CLASS, STRIDES, ANCHORS, i, XYSCALE) bbox_tensors.append(fm) bbox_tensors.append(bbox_tensor) model = tf.keras.Model(input_layer, bbox_tensors) else: raise ValueError if weight_path: if weight_path.split(".")[-1] == "weights": if model_name == 'yolov3_tiny': utils.load_weights_tiny(model, weight_path) elif model_name=='yolov3': utils.load_weights_v3(model, weight_path) elif model_name=='yolov4': utils.load_weights(model, weight_path) else: raise ValueError else: model.load_weights(weight_path) print('Restoring weights from: %s ... ' % weight_path) trainset = Dataset('train') for image_data, target in trainset: pred_result = model(image_data, training=True) giou_loss = conf_loss = prob_loss = 0 for i in range(3): conv, pred = pred_result[i * 2], pred_result[i * 2 + 1] loss_items = ops.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("=> 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))
def prune_train(model_name, weight_path, logdir, save_path, epoches): assert model_name in ['yolov3_tiny', 'yolov3', 'yolov4'] physical_devices = tf.config.experimental.list_physical_devices('GPU') if len(physical_devices) > 0: tf.config.experimental.set_memory_growth(physical_devices[0], True) 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) trainset = Dataset('train') 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 input_layer = tf.keras.layers.Input( [cfg.TRAIN.INPUT_SIZE, cfg.TRAIN.INPUT_SIZE, 3]) if model_name == 'yolov3_tiny': feature_maps = YOLOv3_tiny(input_layer, NUM_CLASS) bbox_tensors = [] for i, fm in enumerate(feature_maps): bbox_tensor = ops.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) elif model_name == 'yolov3': feature_maps = YOLOv3(input_layer, NUM_CLASS) bbox_tensors = [] for i, fm in enumerate(feature_maps): bbox_tensor = ops.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) elif model_name == 'yolov4': feature_maps = YOLOv4(input_layer, NUM_CLASS) bbox_tensors = [] for i, fm in enumerate(feature_maps): bbox_tensor = ops.decode_train(fm, NUM_CLASS, STRIDES, ANCHORS, i, XYSCALE) bbox_tensors.append(fm) bbox_tensors.append(bbox_tensor) model = tf.keras.Model(input_layer, bbox_tensors) else: raise ValueError if weight_path: if weight_path.split(".")[-1] == "weights": if model_name == 'yolov3_tiny': utils.load_weights_tiny(model, weight_path) elif model_name == 'yolov3': utils.load_weights_v3(model, weight_path) elif model_name == 'yolov4': utils.load_weights(model, weight_path) else: raise ValueError else: model.load_weights(weight_path) print('Restoring weights from: %s ... ' % weight_path) optimizer = tf.keras.optimizers.Adam(learning_rate=0.0001) if os.path.exists(logdir): shutil.rmtree(logdir) # for layer in model.layers: # print(layer.name, isinstance(layer, tf.keras.layers.Conv2D)) def apply_pruning_to_dense(layer): if isinstance(layer, tf.keras.layers.Conv2D): return tfmot.sparsity.keras.prune_low_magnitude(layer) return layer # Use `tf.keras.models.clone_model` to apply `apply_pruning_to_dense` # to the layers of the model. model_for_pruning = tf.keras.models.clone_model( model, clone_function=apply_pruning_to_dense, ) # model_for_pruning.summary() unused_arg = -1 model_for_pruning.optimizer = optimizer step_callback = tfmot.sparsity.keras.UpdatePruningStep() step_callback.set_model(model_for_pruning) log_callback = tfmot.sparsity.keras.PruningSummaries( log_dir=logdir) # Log sparsity and other metrics in Tensorboard. log_callback.set_model(model_for_pruning) step_callback.on_train_begin() # run pruning callback for epoch in range(epoches): log_callback.on_epoch_begin(epoch=unused_arg) # run pruning callback for image_data, target in trainset: step_callback.on_train_batch_begin( batch=unused_arg) # run pruning callback with tf.GradientTape() as tape: pred_result = model_for_pruning(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 = ops.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_for_pruning.trainable_variables) optimizer.apply_gradients( zip(gradients, model_for_pruning.trainable_variables)) tf.print( "=> STEP %4d lr: %.6f giou_loss: %4.2f conf_loss: %4.2f " "prob_loss: %4.2f total_loss: %4.2f" % (global_steps, optimizer.lr.numpy(), giou_loss, conf_loss, prob_loss, total_loss)) step_callback.on_epoch_end(batch=unused_arg) # run pruning callback model_for_export = tfmot.sparsity.keras.strip_pruning(model_for_pruning) return model_for_export
append_batch_size=False, dtype='int32') param = {} param['resize_shape'] = resize_shape param['origin_shape'] = origin_shape param['anchors'] = anchors param['conf_thresh'] = conf_thresh param['nms_thresh'] = nms_thresh param['keep_top_k'] = keep_top_k param['nms_top_k'] = nms_top_k param['use_yolo_box'] = use_yolo_box boxes, scores, classes = YOLOv4(inputs, num_classes, num_anchors, is_test=False, trainable=True, postprocess='fastnms', param=param) eval_fetch_list = [boxes, scores, classes] eval_prog = eval_prog.clone(for_test=True) gpu_id = int(os.environ.get('FLAGS_selected_gpus', 0)) place = fluid.CUDAPlace(gpu_id) if use_gpu else fluid.CPUPlace() exe = fluid.Executor(place) exe.run(startup_prog) fluid.load(eval_prog, model_path, executor=exe) if not os.path.exists('images/res/'): os.mkdir('images/res/')
def save_tflite(model_name, weight_path, quantize_mode, output, input_size): assert model_name in ['yolov3_tiny', 'yolov3', 'yolov4'] assert quantize_mode in ['int8', 'float16', 'full_int8'] NUM_CLASS = len(utils.read_class_names(cfg.YOLO.CLASSES)) input_layer = tf.keras.layers.Input([input_size, input_size, 3]) if model_name == 'yolov3_tiny': feature_maps = YOLOv3_tiny(input_layer, NUM_CLASS) bbox_tensors = [] for i, fm in enumerate(feature_maps): bbox_tensor = ops.decode(fm, NUM_CLASS) bbox_tensors.append(bbox_tensor) model = tf.keras.Model(input_layer, bbox_tensors) elif model_name == 'yolov3': feature_maps = YOLOv3(input_layer, NUM_CLASS) bbox_tensors = [] for i, fm in enumerate(feature_maps): bbox_tensor = ops.decode(fm, NUM_CLASS) bbox_tensors.append(bbox_tensor) model = tf.keras.Model(input_layer, bbox_tensors) elif model_name == 'yolov4': feature_maps = YOLOv4(input_layer, NUM_CLASS) bbox_tensors = [] for i, fm in enumerate(feature_maps): bbox_tensor = ops.decode(fm, NUM_CLASS) bbox_tensors.append(bbox_tensor) model = tf.keras.Model(input_layer, bbox_tensors) else: model = None raise ValueError if weight_path.split(".")[-1] == "weights": if model_name == 'yolov3_tiny': utils.load_weights_tiny(model, weight_path) elif model_name == ' yolov3': utils.load_weights_v3(model, weight_path) elif model_name == 'yolov4': utils.load_weights(model, weight_path) else: raise ValueError else: model.load_weights(weight_path).expect_partial() print('Restoring weights from: %s ... ' % weight_path) # model.summary() converter = tf.lite.TFLiteConverter.from_keras_model(model) if tf.__version__ >= '2.2.0': converter.experimental_new_converter = False if quantize_mode == 'int8': converter.optimizations = [tf.lite.Optimize.DEFAULT] elif quantize_mode == 'float16': converter.optimizations = [tf.lite.Optimize.DEFAULT] converter.target_spec.supported_types = [ tf.compat.v1.lite.constants.FLOAT16 ] elif quantize_mode == 'full_int8': converter.target_spec.supported_ops = [ tf.lite.OpsSet.TFLITE_BUILTINS_INT8 ] converter.optimizations = [tf.lite.Optimize.DEFAULT] converter.target_spec.supported_ops = [ tf.lite.OpsSet.TFLITE_BUILTINS, tf.lite.OpsSet.SELECT_TF_OPS ] converter.allow_custom_ops = True converter.representative_dataset = representative_data_gen else: raise ValueError tflite_model = converter.convert() open(output, 'wb').write(tflite_model) logging.info("model saved to: {}".format(output))