def train(self): for epoch in range(self.first_stage_epochs + self.second_stage_epochs): if epoch < self.first_stage_epochs: if not self.isfreeze: self.isfreeze = True for name in self.freeze_layers: freeze = self.model.get_layer(name) freeze_all(freeze) elif epoch >= self.first_stage_epochs: if self.isfreeze: self.isfreeze = False for name in self.freeze_layers: freeze = self.model.get_layer(name) unfreeze_all(freeze) for image_data, target in self.trainset: self.train_step(image_data, target) for image_data, target in self.testset: self.test_step(image_data, target) self.model.save_weights("{}/checkpoints/yolov4".format( self.base_path))
def main(_argv): physical_devices = tf.config.experimental.list_physical_devices('GPU') # if len(physical_devices) > 0: if physical_devices: # tf.config.experimental.set_memory_growth(physical_devices[0], True) tf.config.experimental.set_visible_devices(physical_devices[0], "GPU") trainset = Dataset(FLAGS, is_training=True) testset = Dataset(FLAGS, is_training=False) logdir = "./data/log" 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 # train_steps = (first_stage_epochs + second_stage_epochs) * steps_per_period input_layer = tf.keras.layers.Input([cfg.TRAIN.INPUT_SIZE, cfg.TRAIN.INPUT_SIZE, 3]) STRIDES, ANCHORS, NUM_CLASS, XYSCALE = utils.load_config(FLAGS) IOU_LOSS_THRESH = cfg.YOLO.IOU_LOSS_THRESH freeze_layers = utils.load_freeze_layer(FLAGS.model, FLAGS.tiny) feature_maps = YOLO(input_layer, NUM_CLASS, FLAGS.model, FLAGS.tiny) if FLAGS.tiny: bbox_tensors = [] for i, fm in enumerate(feature_maps): if i == 0: bbox_tensor = decode_train(fm, cfg.TRAIN.INPUT_SIZE // 16, NUM_CLASS, STRIDES, ANCHORS, i, XYSCALE) else: bbox_tensor = decode_train(fm, cfg.TRAIN.INPUT_SIZE // 32, NUM_CLASS, STRIDES, ANCHORS, i, XYSCALE) bbox_tensors.append(fm) bbox_tensors.append(bbox_tensor) else: bbox_tensors = [] for i, fm in enumerate(feature_maps): if i == 0: bbox_tensor = decode_train(fm, cfg.TRAIN.INPUT_SIZE // 8, NUM_CLASS, STRIDES, ANCHORS, i, XYSCALE) elif i == 1: bbox_tensor = decode_train(fm, cfg.TRAIN.INPUT_SIZE // 16, NUM_CLASS, STRIDES, ANCHORS, i, XYSCALE) else: bbox_tensor = decode_train(fm, cfg.TRAIN.INPUT_SIZE // 32, NUM_CLASS, STRIDES, ANCHORS, i, XYSCALE) bbox_tensors.append(fm) bbox_tensors.append(bbox_tensor) model = tf.keras.Model(input_layer, bbox_tensors) model.summary() if FLAGS.weights == None: print("Training from scratch") else: if FLAGS.weights.split(".")[len(FLAGS.weights.split(".")) - 1] == "weights": utils.load_weights(model, FLAGS.weights, FLAGS.model, FLAGS.tiny) else: model.load_weights(FLAGS.weights) print('Restoring weights from: %s ... ' % FLAGS.weights) optimizer = tf.keras.optimizers.Adam() if os.path.exists(logdir): shutil.rmtree(logdir) writer = tf.summary.create_file_writer(logdir) # define training step function # @tf.function 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(len(freeze_layers)): conv, pred = pred_result[i * 2], pred_result[i * 2 + 1] loss_items = 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/%4d lr: %.6f giou_loss: %4.2f conf_loss: %4.2f " "prob_loss: %4.2f total_loss: %4.2f" % (global_steps, total_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()) # 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): 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(len(freeze_layers)): conv, pred = pred_result[i * 2], pred_result[i * 2 + 1] loss_items = 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 freeze_layers: freeze = model.get_layer(name) freeze_all(freeze) elif epoch >= first_stage_epochs: if isfreeze: isfreeze = False for name in freeze_layers: freeze = model.get_layer(name) 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) model.save_weights("./checkpoints/yolov4")
def main(_argv): physical_devices = tf.config.experimental.list_physical_devices('GPU') if len(physical_devices) > 0: tf.config.experimental.set_memory_growth(physical_devices[0], True) model = Yolo(FLAGS.size, training=True, classes=FLAGS.num_classes, aux=True) anchors, mask_len = get_anchors(FLAGS.size) train_dataset = dataset.load_tfrecord_dataset(FLAGS.dataset, FLAGS.size) if FLAGS.check: logging.info('Current dataset check') with open(FLAGS.classes) as classes_f: class_names = [c.strip() for c in classes_f.readlines()] logging.info('Classes loaded') check_dataset(train_dataset, class_names) train_dataset = train_dataset.repeat() train_dataset = train_dataset.shuffle(buffer_size=200) train_dataset = train_dataset.batch(FLAGS.batch_size) train_dataset = train_dataset.map(lambda x, y: ( dataset.transform_images(x, FLAGS.size), dataset.transform_targets(y, anchors, mask_len, FLAGS.size))) train_dataset = train_dataset.prefetch( buffer_size=tf.data.experimental.AUTOTUNE) if FLAGS.validation_dataset: validation_dataset = dataset.load_tfrecord_dataset( FLAGS.validation_dataset, FLAGS.classes, FLAGS.size) else: validation_dataset = dataset.load_dummy_validation() validation_dataset = validation_dataset.batch(FLAGS.batch_size) validation_dataset = validation_dataset.map(lambda x, y: ( dataset.transform_images(x, FLAGS.size), dataset.transform_targets(y, anchors, mask_len, FLAGS.size))) if FLAGS.weights: if FLAGS.pretrained: model_pretrained = Yolo(FLAGS.size, training=True, classes=FLAGS.weights_num_classes) model_pretrained.load_weights(FLAGS.weights) model.get_layer('darknet_recursive').set_weights( model_pretrained.get_layer('darknet_recursive').get_weights()) else: model.load_weights(FLAGS.weights) if FLAGS.freeze: freeze_all(model.get_layer('darknet_recursive')) optimizer = tf.keras.optimizers.Adam(lr=FLAGS.learning_rate) loss = [ Loss(anchors[:mask_len], classes=FLAGS.num_classes), Loss(anchors[mask_len:], classes=FLAGS.num_classes) ] model.compile(optimizer=optimizer, loss=loss, run_eagerly=True, metrics=['accuracy']) callbacks = [ ReduceLROnPlateau(verbose=1, patience=10), EarlyStopping(patience=3, verbose=1), ModelCheckpoint('checkpoints/yolov3_face.tf', verbose=1, save_weights_only=True), TensorBoard(log_dir='logs') ] history = model.fit(train_dataset, epochs=FLAGS.epochs, callbacks=callbacks, validation_data=validation_dataset, steps_per_epoch=10)
def main(_argv): input_channel = 3 patience = 30 steps_in_epoch = 0 epoch_loss = np.inf prev_minloss = np.inf trainset = Dataset(FLAGS,input_channel, is_training=True) # testset = Dataset(FLAGS, input_channel, is_training=False) logdir = "./data/log" isfreeze = False steps_per_epoch = len(trainset) print("steps_per_epoch:{}".format(steps_per_epoch)) 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 # train_steps = (first_stage_epochs + second_stage_epochs) * steps_per_period loss_tracker = [] losses_in_epoch = 0. input_layer = tf.keras.layers.Input([cfg.TRAIN.INPUT_SIZE, cfg.TRAIN.INPUT_SIZE, 3]) STRIDES, ANCHORS, NUM_CLASS, XYSCALE = utils.load_config(FLAGS) IOU_LOSS_THRESH = cfg.YOLO.IOU_LOSS_THRESH freeze_layers = utils.load_freeze_layer(FLAGS.model, FLAGS.num_detection_layer) # feature_maps = YOLO(input_layer, NUM_CLASS, FLAGS.model, FLAGS.tiny) feature_maps = YOLOv4_more_tiny(input_layer, NUM_CLASS) bbox_tensors = [] for i, fm in enumerate(feature_maps): # fm shape: (None, featw, feath, filters) bbox_tensor = decode_train(fm, cfg.TRAIN.INPUT_SIZE // 16, NUM_CLASS, STRIDES, ANCHORS, i, XYSCALE) bbox_tensors.append(fm) bbox_tensors.append(bbox_tensor) if cfg.YOLO.NUM_YOLOLAYERS == 3: # yolov4 bbox_tensors = [] for i, fm in enumerate(feature_maps): if i == 0: bbox_tensor = decode_train(fm, cfg.TRAIN.INPUT_SIZE // 8, NUM_CLASS, STRIDES, ANCHORS, i, XYSCALE) elif i == 1: bbox_tensor = decode_train(fm, cfg.TRAIN.INPUT_SIZE // 16, NUM_CLASS, STRIDES, ANCHORS, i, XYSCALE) else: bbox_tensor = decode_train(fm, cfg.TRAIN.INPUT_SIZE // 32, NUM_CLASS, STRIDES, ANCHORS, i, XYSCALE) bbox_tensors.append(fm) bbox_tensors.append(bbox_tensor) elif cfg.YOLO.NUM_YOLOLAYERS == 2: # yolo tiny bbox_tensors = [] for i, fm in enumerate(feature_maps): if i == 0: bbox_tensor = decode_train(fm, cfg.TRAIN.INPUT_SIZE // 16, NUM_CLASS, STRIDES, ANCHORS, i, XYSCALE) else: bbox_tensor = decode_train(fm, cfg.TRAIN.INPUT_SIZE // 32, NUM_CLASS, STRIDES, ANCHORS, i, XYSCALE) bbox_tensors.append(fm) bbox_tensors.append(bbox_tensor) elif cfg.YOLO.NUM_YOLOLAYERS == 1: # custom yolo bbox_tensors = [] bbox_tensor = decode_train(feature_maps[0], cfg.TRAIN.INPUT_SIZE // cfg.YOLO.STRIDES_CUSTOM[0], NUM_CLASS, STRIDES, ANCHORS, 0, XYSCALE) bbox_tensors.append(feature_maps[0]) bbox_tensors.append(bbox_tensor) model = tf.keras.Model(input_layer, bbox_tensors) model.summary() if FLAGS.weights == None: print("Training from scratch") else: if FLAGS.weights.split(".")[len(FLAGS.weights.split(".")) - 1] == "weights": utils.load_weights(model, FLAGS.weights, FLAGS.model, FLAGS.num_detection_layer) else: model.load_weights(FLAGS.weights) print('Restoring weights from: %s ... ' % FLAGS.weights) optimizer = tf.keras.optimizers.Adam() if os.path.exists(logdir): shutil.rmtree(logdir) writer = tf.summary.create_file_writer(logdir) # define training step function # @tf.function def train_step(image_data, target): with tf.GradientTape() as tape: pred_result = model(image_data, training=True) ciou_loss = conf_loss = prob_loss = 0 # optimizing process for i in range(len(freeze_layers)): conv, pred = pred_result[i * 2], pred_result[i * 2 + 1] # target[i][1]:(32, 150, 4) # print("conv shape:{} pred shape:{}".format(tf.keras.backend.int_shape(conv), tf.keras.backend.int_shape(pred))) # print("target[i][0]:{} target[i][1]:{}".format(np.array(target[i][0]).shape, np.array(target[i][1]).shape)) loss_items = compute_loss(pred, conv, target[i][0], target[i][1], STRIDES=STRIDES, NUM_CLASS=NUM_CLASS, IOU_LOSS_THRESH=IOU_LOSS_THRESH, i=i) ciou_loss += loss_items[0] conf_loss += loss_items[1] prob_loss += loss_items[2] total_loss = ciou_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/%4d lr: %.6f ciou_loss: %4.2f conf_loss: %4.2f " "prob_loss: %4.2f total_loss: %4.2f" % (global_steps, total_steps, optimizer.lr.numpy(), ciou_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()) # 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/ciou_loss", ciou_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() # return total_loss @tf.function def test_step(image_data, target): with tf.GradientTape() as tape: pred_result = model(image_data, training=True) ciou_loss = conf_loss = prob_loss = 0 preds = None gt_infos = None gt_bboxes = None # optimizing process for i in range(len(freeze_layers)): conv, pred = pred_result[i * 2], pred_result[i * 2 + 1] loss_items = compute_loss(pred, conv, target[i][0], target[i][1], STRIDES=STRIDES, NUM_CLASS=NUM_CLASS, IOU_LOSS_THRESH=IOU_LOSS_THRESH, i=i) ciou_loss += loss_items[0] conf_loss += loss_items[1] prob_loss += loss_items[2] if FLAGS.show_map: batch_, _, _, _, ch_ = tf.keras.backend.int_shape(pred) if preds == None: preds = tf.reshape(pred, (batch_, -1, ch_)) else: preds = tf.concat([preds, tf.reshape(preds, (batch_, -1, ch_))], axis=1) if gt_infos == None: gt_infos = tf.reshape(target[i][0], (batch_, -1, ch_)) gt_bboxes = target[i][1] else: gt_infos = tf.concat([gt_infos, tf.reshape(target[i][0], (batch_, -1, ch_))], axis=1) gt_bboxes = tf.concat([gt_bboxes, tf.reshape(target[i][1], (batch_, -1, 4))], axis=1) if FLAGS.show_map: map = compute_map(preds, gt_bboxes, gt_infos, NUM_CLASS) total_map = map total_loss = ciou_loss + conf_loss + prob_loss tf.print("=> TEST STEP %4d ciou_loss: %4.2f conf_loss: %4.2f " "prob_loss: %4.2f total_loss: %4.2f" % (global_steps, ciou_loss, conf_loss, prob_loss, total_loss)) tf.print("=> TEST STEP %4d map: %4.2f" % (total_map)) # 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): start_ch = time.time() train_loss = 0. if epoch < first_stage_epochs: if not isfreeze: isfreeze = True for name in freeze_layers: try: freeze = model.get_layer(name) freeze_all(freeze) except ValueError: pass elif epoch >= first_stage_epochs: if isfreeze: isfreeze = False for name in freeze_layers: try: # try구문 추가 freeze = model.get_layer(name) freeze_all(freeze) except ValueError: pass for image_data, target in trainset: # bboxes_list = target[0][1] # label_data = target[0][0] # # for batch_idx, bboxes in enumerate(bboxes_list): # # print("bboxes:{}".format(bboxes)) # class_inds = [] # check = np.array(image_data[batch_idx]).reshape(cfg.TRAIN.INPUT_SIZE, cfg.TRAIN.INPUT_SIZE, 3) # # print(r"D:\tf_data_not\preprocessed_{}.npy".format(batch_idx)) # # np.save(r"D:\tf_data_not\image_{}.npy".format(batch_idx), check) # # np.save(r"D:\tf_data_not\preprocessed_{}.npy".format(batch_idx), label_data[batch_idx]) # # np.save(r"D:\tf_data_not\bboxes_{}.npy".format(batch_idx), bboxes) # label = label_data[batch_idx] # label_class = label[...,5:] # # print("label shape:{}".format(np.array(label).shape)) # # for line_label in label_class: # if np.sum(line_label) > 0: # print(line_label) # class_ = np.array(label[..., 5:]).flatten() # class_ = np.where(class_>0.1, class_, 0) # for class_idx, class_label in enumerate(class_): # if class_label > 0: # class_inds.append(class_idx % cfg.YOLO.NUM_CLASSES) # # for bbox_idx, bbox in enumerate(bboxes): # half_h = bbox[3] / 2 # half_w = bbox[2] / 2 # # if np.sum(bbox) > 0: # # print("class:{}".format(class_inds)) # cv2.rectangle(check, (int(bbox[0] - half_w), int(bbox[1] - half_h)), # (int(bbox[0] + half_w), int(bbox[1] + half_h)), color=(0, 255, 0), thickness=3) # cv2.putText(check, text="{}".format(class_inds[bbox_idx]), # org=(int(bbox[0] + half_w)-10, int(bbox[1] + half_h)-30), thickness=2, color=(0, 255, 0), # fontFace=cv2.FONT_HERSHEY_SIMPLEX, fontScale=0.7) # # cv2.imshow("check_aug", check) # cv2.waitKey(0) # cv2.destroyAllWindows() # cv2.imwrite(os.path.join(r"D:\Public\JHS\SIMPLE_DATA_INPUT_CHECK", "{}.jpg".format(batch_idx)), check*255) train_step(image_data, target) # for image_data, target in testset: # test_step(image_data, target) ### for loss graph """ if steps_in_epoch >= steps_per_epoch: loss_tracker.append(losses_in_epoch / steps_per_epoch) plt.plot(loss_tracker) plt.xlabel('epoch') plt.ylabel('loss') # plt.show() plt.savefig("D:/checkpoint/loss.png") """ # print("steps_in_epoch:{}, steps_per_epoch:{}".format(global_steps, steps_per_epoch)) # print("prev_minloss:{}, epoch_loss:{}".format(prev_minloss, epoch_loss)) # # early stopping # if len(losses_in_epoch) >= steps_per_epoch: # # epoch_loss = losses_in_epoch / steps_per_epoch # if prev_minloss > epoch_loss: # save best weight # prev_minloss = epoch_loss # # if epoch > 800: # 최소학습 에폭 # model.save("D:\ckpt_best") # print("{} epoch save best weights".format(epoch)) # # if len(loss_tracker) > patience: # print("check loss_tracker len:{}".format(len(loss_tracker))) # if loss_tracker[0] > np.min(loss_tracker[1:]): # loss_tracker.pop(0) # else: # print("total loss didn't decreased during {} epochs. train stop".format(patience)) # return # steps_in_epoch = 0 # epoch_loss = train_loss # else: # epoch_loss += train_loss if (epoch+1) % 500 == 0: model.save(r"D:\notf_ckpt-epoch{}".format(epoch)) print("{} epoch model saved".format(epoch)) model.save(r"D:\notf_ckpt-last")
def main(): strategy = tf.distribute.MirroredStrategy() num_gpus = strategy.num_replicas_in_sync with strategy.scope(): train_batch_size = cfg.TRAIN.BATCH_SIZE val_batch_size = cfg.VAL.BATCH_SIZE trainset = Dataset('train') valset = Dataset('val') train_generator = tf.data.Dataset.from_generator( lambda: trainset, (tf.float32, ((tf.float32, tf.float32), (tf.float32, tf.float32), (tf.float32, tf.float32)))).batch(train_batch_size) val_generator = tf.data.Dataset.from_generator( lambda: valset, (tf.float32, ((tf.float32, tf.float32), (tf.float32, tf.float32), (tf.float32, tf.float32)))).batch(val_batch_size) dist_trainset = strategy.experimental_distribute_dataset( train_generator) dist_valset = strategy.experimental_distribute_dataset(val_generator) isfreeze = False steps_per_epoch = math.ceil(len(trainset) / train_batch_size) val_steps_per_epoch = math.ceil(len(valset) / val_batch_size) 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) val_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 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) model = YOLOv4(NUM_CLASS, STRIDES, ANCHORS, XYSCALE) optimizer = tf.keras.optimizers.Adam() ckpt = tf.train.Checkpoint(optimizer=optimizer, model=model) ckpt_manager = tf.train.CheckpointManager(ckpt, args.logdir, max_to_keep=3) writer = tf.summary.create_file_writer(args.logdir) start_iteration = int(ckpt_manager.latest_checkpoint.split('-') [-1]) if args.weights else 0 if args.weights: print('Restoring weights from: %s ... ' % ckpt_manager.latest_checkpoint) ckpt.restore(ckpt_manager.latest_checkpoint) else: print('Training from scratch') print('Start iteration:', start_iteration) @tf.function 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 = utils.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 giou_loss = tf.reduce_sum(giou_loss) * (1.0 / train_batch_size) conf_loss = tf.reduce_sum(conf_loss) * (1.0 / train_batch_size) prob_loss = tf.reduce_sum(prob_loss) * (1.0 / train_batch_size) total_loss = tf.reduce_sum(total_loss) * (1.0 / train_batch_size) gradients = tape.gradient(total_loss, model.trainable_variables) optimizer.apply_gradients(zip(gradients, model.trainable_variables)) return total_loss, giou_loss, conf_loss, prob_loss def distribute_train_step(image_data, target): with strategy.scope(): tloss_temp, gloss_temp, closs_temp, ploss_temp = strategy.experimental_run_v2( train_step, args=(image_data, target)) tloss = strategy.reduce(tf.distribute.ReduceOp.MEAN, tloss_temp, axis=None) gloss = strategy.reduce(tf.distribute.ReduceOp.MEAN, gloss_temp, axis=None) closs = strategy.reduce(tf.distribute.ReduceOp.MEAN, closs_temp, axis=None) ploss = strategy.reduce(tf.distribute.ReduceOp.MEAN, ploss_temp, axis=None) return tloss, gloss, closs, ploss @tf.function def val_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 = utils.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 giou_loss = tf.reduce_sum(giou_loss) * (1.0 / val_batch_size) conf_loss = tf.reduce_sum(conf_loss) * (1.0 / val_batch_size) prob_loss = tf.reduce_sum(prob_loss) * (1.0 / val_batch_size) total_loss = tf.reduce_sum(total_loss) * (1.0 / val_batch_size) return total_loss, giou_loss, conf_loss, prob_loss def distribute_val_step(image_data, target): with strategy.scope(): tloss_temp, gloss_temp, closs_temp, ploss_temp = strategy.experimental_run_v2( val_step, args=(image_data, target)) tloss = strategy.reduce(tf.distribute.ReduceOp.MEAN, tloss_temp, axis=None) gloss = strategy.reduce(tf.distribute.ReduceOp.MEAN, gloss_temp, axis=None) closs = strategy.reduce(tf.distribute.ReduceOp.MEAN, closs_temp, axis=None) ploss = strategy.reduce(tf.distribute.ReduceOp.MEAN, ploss_temp, axis=None) return tloss, gloss, closs, ploss def train(epoch, max_epoch): for image_data, target in dist_trainset: total_loss, giou_loss, conf_loss, prob_loss = distribute_train_step( image_data, target) tf.print( "=> EPOCH [%3d/%3d] ITER [%5d/%5d] lr: %.6f giou_loss: %4.2f conf_loss: %4.2f " "prob_loss: %4.2f total_loss: %4.2f" % (epoch, max_epoch, global_steps - (epoch - 1) * steps_per_epoch, steps_per_epoch, optimizer.lr.numpy(), giou_loss, conf_loss, prob_loss, total_loss)) # 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() # 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()) def val(epoch, max_epoch): for image_data, target in dist_valset: total_loss, giou_loss, conf_loss, prob_loss = distribute_val_step( image_data, target) tf.print( "=> VALIDATION EPOCH [%3d/%3d] ITER [%5d/%5d] giou_loss: %4.2f conf_loss: %4.2f " "prob_loss: %4.2f total_loss: %4.2f" % (epoch, max_epoch, val_global_steps - (epoch - 1) * val_steps_per_epoch, val_steps_per_epoch, giou_loss, conf_loss, prob_loss, total_loss)) # writing summary data with writer.as_default(): tf.summary.scalar("val_loss/total_loss", total_loss, step=val_global_steps) tf.summary.scalar("val_loss/giou_loss", giou_loss, step=val_global_steps) tf.summary.scalar("val_loss/conf_loss", conf_loss, step=val_global_steps) tf.summary.scalar("val_loss/prob_loss", prob_loss, step=val_global_steps) writer.flush() val_global_steps.assign_add(1) for epoch in range(first_stage_epochs + second_stage_epochs): if epoch < first_stage_epochs: if not isfreeze: isfreeze = True for name in ['conv_93', 'conv_101', 'conv_109']: freeze = model.get_layer(name) utils.freeze_all(freeze) elif epoch >= first_stage_epochs: if isfreeze: isfreeze = False for name in ['conv_93', 'conv_101', 'conv_109']: freeze = model.get_layer(name) utils.unfreeze_all(freeze) train(epoch + 1, first_stage_epochs + second_stage_epochs) val(epoch + 1, first_stage_epochs + second_stage_epochs) ckpt_manager.save(checkpoint_number=epoch)
def main(_argv): physical_devices = tf.config.experimental.list_physical_devices('GPU') if len(physical_devices) > 0: tf.config.experimental.set_memory_growth(physical_devices[0], True) trainset = Dataset('train') testset = Dataset('test') logdir = "./data/log" 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 = warmup_steps = cfg.TRAIN.WARMUP_EPOCHS * steps_per_epoch total_steps = (first_stage_epochs + second_stage_epochs) * steps_per_epoch # train_steps = (first_stage_epochs + second_stage_epochs) * steps_per_period 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) input_layer = tf.keras.layers.Input([cfg.TRAIN.INPUT_SIZE, cfg.TRAIN.INPUT_SIZE, 3]) #YOLOV4 feature_maps = YOLOv4(input_layer, NUM_CLASS) bbox_tensors = [] for i, fm in enumerate(feature_maps): bbox_tensor = 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) if FLAGS.weights == "none": print("Training from scratch") else: if FLAGS.weights.split(".")[len(FLAGS.weights.split(".")) - 1] == "weights": if FLAGS.tiny: utils.load_weights_tiny(model, FLAGS.weights) else: if FLAGS.model == 'yolov3': utils.load_weights_v3(model, FLAGS.weights) else: utils.load_weights(model, FLAGS.weights) else: model.load_weights(FLAGS.weights) print('Restoring weights from: %s ... ' % FLAGS.weights) if os.path.exists(logdir): shutil.rmtree(logdir) optimizer = tf.keras.optimizers.Adam() writer = tf.summary.create_file_writer(logdir) ckpt = tf.train.Checkpoint(step=tf.Variable(1, trainable=False, dtype=tf.int64), optimizer=optimizer, net=model) manager = tf.train.CheckpointManager(ckpt, './checkpoints', max_to_keep=10) 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 = 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" % (ckpt.step, optimizer.lr.numpy(), giou_loss, conf_loss, prob_loss, total_loss)) # update learning rate ckpt.step.assign_add(1) if ckpt.step < warmup_steps: lr = ckpt.step / 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((ckpt.step - warmup_steps) / (total_steps - warmup_steps) * np.pi)) ) optimizer.lr.assign(lr.numpy()) # writing summary data with writer.as_default(): tf.summary.scalar("lr", optimizer.lr, step=ckpt.step) tf.summary.scalar("loss/total_loss", total_loss, step=ckpt.step) tf.summary.scalar("loss/giou_loss", giou_loss, step=ckpt.step) tf.summary.scalar("loss/conf_loss", conf_loss, step=ckpt.step) tf.summary.scalar("loss/prob_loss", prob_loss, step=ckpt.step,) writer.flush() def test_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 = 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" % (ckpt.step, giou_loss, conf_loss, prob_loss, total_loss)) ckpt.restore(manager.latest_checkpoint) if manager.latest_checkpoint: print("Restored from {}".format(manager.latest_checkpoint)) else: print("Initializing from scratch.") for epoch in range(first_stage_epochs + second_stage_epochs): ckpt.step.assign_add(1) print("----------- EPOCH :: ", str(epoch), " :: -----------") if epoch < first_stage_epochs: print("epoch < first_stage_epochs") if not isfreeze: isfreeze = True for name in ['conv2d_93', 'conv2d_101', 'conv2d_109']: freeze = model.get_layer(name) freeze_all(freeze) elif epoch >= first_stage_epochs: print("epoch < first_stage_epochs") if isfreeze: isfreeze = False for name in ['conv2d_93', 'conv2d_101', 'conv2d_109']: freeze = model.get_layer(name) unfreeze_all(freeze) for image_data, target in trainset: train_step(image_data, target) if int(ckpt.step) % 10 == 0: save_path = manager.save() print("Saved checkpoint for step {}: {}".format(int(ckpt.step), save_path))
def main(_argv): gpus = tf.config.experimental.list_physical_devices('GPU') if gpus: # 텐서플로가 첫 번째 GPU만 사용하도록 제한 try: tf.config.experimental.set_visible_devices(gpus[7], 'GPU') except RuntimeError as e: # 프로그램 시작시에 접근 가능한 장치가 설정되어야만 합니다 print(e) trainset = Dataset(FLAGS, is_training=True) #print(next(iter(trainset))) #_train = next(iter(trainset)) #_train[0] = (3, 416, 416, 3), len(train[1] =3, #len(train[1][0])) =2, # _train[1][0][0].shape : (3, 26, 26, 3, 25) # _train[1][0][1].shape : (3, 150, 4) # len(_train[1][1]) =2, # _train[1][1][0].shape : (3, 26, 26, 3, 25) # _train[1][1][1].shape : (3, 150, 4) testset = Dataset(FLAGS, is_training=False) logdir = "./data/log" 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 # train_steps = (first_stage_epochs + second_stage_epochs) * steps_per_period input_layer = tf.keras.layers.Input([cfg.TRAIN.INPUT_SIZE, cfg.TRAIN.INPUT_SIZE, 3]) STRIDES, ANCHORS, NUM_CLASS, XYSCALE = utils.load_config(FLAGS) IOU_LOSS_THRESH = cfg.YOLO.IOU_LOSS_THRESH freeze_layers = utils.load_freeze_layer(FLAGS.model, FLAGS.tiny) #yolov4 : ['conv2d_93', 'conv2d_101', 'conv2d_109'] #yolov4_tiny: ['conv2d_17', 'conv2d_20'] feature_maps = YOLO(input_layer, NUM_CLASS, FLAGS.model, FLAGS.tiny) # ''' # import core.backbone as backbone # route_1, route_2, conv = backbone.cspdarknet53(input_layer) # route_1 # Out[97]: <tf.Tensor 'Mul_355:0' shape=(None, 52, 52, 256) dtype=float32> # route_2 # Out[98]: <tf.Tensor 'Mul_376:0' shape=(None, 26, 26, 512) dtype=float32> # conv # Out[99]: <tf.Tensor 'LeakyRelu_164:0' shape=(None, 13, 13, 512) dtype=float32> # ''' #yolov4 : [<tf.Tensor 'conv2d_93/BiasAdd:0' shape=(None, 52, 52, 75) dtype=float32>, #<tf.Tensor 'conv2d_101/BiasAdd:0' shape=(None, 26, 26, 75) dtype=float32>, #<tf.Tensor 'conv2d_109/BiasAdd:0' shape=(None, 13, 13, 75) dtype=float32>] #yolov4_tiny : [<tf.Tensor 'conv2d_130/BiasAdd:0' shape=(None, 26, 26, 75) dtype=float32>, #<tf.Tensor 'conv2d_127/BiasAdd:0' shape=(None, 13, 13, 75) dtype=float32>] if FLAGS.tiny: bbox_tensors = [] for i, fm in enumerate(feature_maps): if i == 0: bbox_tensor = decode_train(fm, cfg.TRAIN.INPUT_SIZE // 16, NUM_CLASS, STRIDES, ANCHORS, i, XYSCALE) else: bbox_tensor = decode_train(fm, cfg.TRAIN.INPUT_SIZE // 32, NUM_CLASS, STRIDES, ANCHORS, i, XYSCALE) bbox_tensors.append(fm) bbox_tensors.append(bbox_tensor) else: bbox_tensors = [] for i, fm in enumerate(feature_maps): if i == 0: bbox_tensor = decode_train(fm, cfg.TRAIN.INPUT_SIZE // 8, NUM_CLASS, STRIDES, ANCHORS, i, XYSCALE) elif i == 1: bbox_tensor = decode_train(fm, cfg.TRAIN.INPUT_SIZE // 16, NUM_CLASS, STRIDES, ANCHORS, i, XYSCALE) else: bbox_tensor = decode_train(fm, cfg.TRAIN.INPUT_SIZE // 32, NUM_CLASS, STRIDES, ANCHORS, i, XYSCALE) bbox_tensors.append(fm) bbox_tensors.append(bbox_tensor) model = tf.keras.Model(input_layer, bbox_tensors) model.summary() # from contextlib import redirect_stdout # with open('modelsummary.txt', 'w') as f: # with redirect_stdout(f): # model.summary() if FLAGS.weights == None: print("Training from scratch") else: if FLAGS.weights.split(".")[len(FLAGS.weights.split(".")) - 1] == "weights": utils.load_weights(model, FLAGS.weights, FLAGS.model, FLAGS.tiny) else: model.load_weights(FLAGS.weights) print('Restoring weights from: %s ... ' % FLAGS.weights) # wf = open(FLAGS.weights, 'rb') # major, minor, revision, seen, _ = np.fromfile(wf, dtype=np.int32, count=5) # wf.close() # layer_size = 110 # output_pos = [93, 101, 109] # j = 0 # for i in range(layer_size): # conv_layer_name = 'conv2d_%d' %i if i > 0 else 'conv2d' # bn_layer_name = 'batch_normalization_%d' %j if j > 0 else 'batch_normalization' # conv_layer = model.get_layer(conv_layer_name) # filters = conv_layer.filters # k_size = conv_layer.kernel_size[0] # in_dim = conv_layer.input_shape[-1] optimizer = tf.keras.optimizers.Adam() if os.path.exists(logdir): shutil.rmtree(logdir) writer = tf.summary.create_file_writer(logdir) # define training step function # @tf.function 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(len(freeze_layers)): conv, pred = pred_result[i * 2], pred_result[i * 2 + 1] # print(pred) loss_items = 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/%4d lr: %.6f giou_loss: %4.2f conf_loss: %4.2f " "prob_loss: %4.2f total_loss: %4.2f" % (global_steps, total_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()) # 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): 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(len(freeze_layers)): conv, pred = pred_result[i * 2], pred_result[i * 2 + 1] loss_items = 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 freeze_layers: freeze = model.get_layer(name) freeze_all(freeze) elif epoch >= first_stage_epochs: if isfreeze: isfreeze = False for name in freeze_layers: freeze = model.get_layer(name) 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) model.save_weights("./checkpoints/yolov4")
def main(): INPUT_SIZE = FLAGS.size STRIDES, ANCHORS, NUM_CLASS, XYSCALE = utils.load_config(FLAGS) CLASSES = utils.read_class_names(cfg.YOLO.CLASSES) IOU_LOSS_THRESH = cfg.YOLO.IOU_LOSS_THRESH predicted_dir_path = './AP/predicted' loss_dir_path = './AP/loss' text_result_path = './AP/detect' trainset = Dataset(FLAGS, is_training=True) adv_lambda = 0.001 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]) feature_maps = YOLO(FLAGS.scale_v5, input_layer, NUM_CLASS) bbox_tensors = [] for i, fm in enumerate(feature_maps): if i == 0: bbox_tensor = decode_train(fm, cfg.TRAIN.INPUT_SIZE // 8, NUM_CLASS, STRIDES, ANCHORS, i, XYSCALE) bbox_tensors.append(fm) bbox_tensors.append(bbox_tensor) elif i == 1: bbox_tensor = decode_train(fm, cfg.TRAIN.INPUT_SIZE // 16, NUM_CLASS, STRIDES, ANCHORS, i, XYSCALE) bbox_tensors.append(fm) bbox_tensors.append(bbox_tensor) elif i==2: bbox_tensor = decode_train(fm, cfg.TRAIN.INPUT_SIZE // 32, NUM_CLASS, STRIDES, ANCHORS, i, XYSCALE) bbox_tensors.append(fm) bbox_tensors.append(bbox_tensor) else: result_G = fm bbox_tensors.append(result_G) D2_model = tf.keras.Model(input_layer, bbox_tensors) full_model = NLayerDiscriminator(ndf=64, n_layers=5) full_model.build([None, cfg.TRAIN.INPUT_SIZE, cfg.TRAIN.INPUT_SIZE, 3]) optimizer_D2 = tf.keras.optimizers.Adam() optimizer_Dis = tf.keras.optimizers.SGD() criterionG, criterionD = get_loss() if cfg.TRAIN.DoubleGAN: patch_model = NLayerDiscriminator(ndf=64, n_layers=3) patch_model.build([None, cfg.TRAIN.INPUT_SIZE, cfg.TRAIN.INPUT_SIZE, 3]) adv_trainer = DoubleGAN(patch_model, full_model, criterionD) else: adv_trainer = SingleGAN(full_model, criterionD) if cfg.TRAIN.GRADNORM: optimizer_P = tf.keras.optimizers.SGD() D2_model.summary() T_freeze_layers = utils.load_True_freeze_layer(FLAGS.scale_v5) # metrics draw now total_loss_metric = tf.metrics.Mean() loss1_metric = tf.metrics.Mean() loss2_metric = tf.metrics.Mean() total_loss_result= [] loss1_result= [] loss2_result= [] Weightloss1 = tf.Variable(1.0) Weightloss2 = tf.Variable(1.0) params = [Weightloss1, Weightloss2] alph = 0.16 ## @tf.function def train_step(image_data, target): start_time = time.time() # with tf.GradientTape() as tape1,tf.GradientTape() as tape2: ##Experiments have found that this performance is better than "persistent=True". with tf.GradientTape() as tape1,tf.GradientTape() as tape2,tf.GradientTape() as tape3,tf.GradientTape() as tape4,tf.GradientTape() as tape5: pred_result = D2_model(image_data[0], training=True) G_im = pred_result[-1] loss_D = loss_content = loss_adv = loss_G = giou_loss = conf_loss = prob_loss = 0 #update Discriminator loss_D = 1000 * adv_lambda * adv_trainer.loss_d(G_im, image_data[1]) gradients_Dis = tape1.gradient(loss_D, adv_trainer.get_params()) optimizer_Dis.apply_gradients(zip(gradients_Dis, adv_trainer.get_params())) #update D2Net loss_content = criterionG(G_im, image_data[1]) loss_adv = adv_trainer.loss_g(G_im, image_data[1]) loss_G = 1000*(loss_content + adv_lambda * loss_adv) for i in range(ANCHORS.shape[0]): conv, pred = pred_result[i * 2], pred_result[i * 2 + 1] loss_items = 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] yolo_loss = giou_loss + conf_loss + prob_loss l1 = params[0]*yolo_loss l2 = params[1]*loss_G total_loss = (l1 + l2)/2 gradients_D2 = tape2.gradient(total_loss, D2_model.trainable_variables) optimizer_D2.apply_gradients(zip(gradients_D2, D2_model.trainable_variables)) ###Gradnorm### L0 = 183 LP = D2_model.trainable_variables[162] #D_conv2d_54 G1R = tape3.gradient(l1, LP) G1 = tf.norm(G1R, ord=2) G2R = tape4.gradient(l2, LP) G2 = tf.norm(G2R, ord=2) G_avg = (G1+G2)/2 # Calculating relative losses lhat1 = (l1)/L0 lhat2 = (l2)/L0 lhat_avg = (lhat1 + lhat2)/2 inv_rate1 = lhat1/lhat_avg inv_rate2 = lhat2/lhat_avg C1 = G_avg*(inv_rate1)**alph C2 = G_avg*(inv_rate2)**alph C1 = tf.stop_gradient(tf.identity(C1)) C2 = tf.stop_gradient(tf.identity(C2)) # Gradnorm loss loss_gradnorm = tf.math.reduce_sum(tf.math.abs(G1-C1)) + tf.math.reduce_sum(tf.math.abs(G2-C2)) grad_grad = tape5.gradient(loss_gradnorm, params) optimizer_P.apply_gradients(grads_and_vars=zip(grad_grad, params)) total_loss_metric.update_state(values=total_loss) loss1_metric.update_state(values=yolo_loss) loss2_metric.update_state(values=loss_G) time_per_step = time.time() - start_time print("Step: {}/{},lr: {:.6f}, {:.2f}s/step, total_loss: {:.5f}, " "yolo loss: {:.5f}, G_loss: {:.5f}, loss_adv: {:.5f}, D_loss: {:.5f}".format( global_steps.numpy(), total_steps, optimizer_D2.lr.numpy(), time_per_step, total_loss, yolo_loss, loss_G, adv_lambda*loss_adv, loss_D )) # 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_D2.lr.assign(lr.numpy()) optimizer_Dis.lr.assign(10*lr.numpy()) loss_list_step = [optimizer_D2.lr.numpy(),total_loss,yolo_loss, loss_G,loss_D,giou_loss,conf_loss, prob_loss,loss_content,adv_lambda * loss_adv] return np.array(loss_list_step) ## @tf.function def test_epoch(D2_model,dectect_epoch_path): with open(cfg.TEST.ANNOT_PATH, 'r') as annotation_file: for num, line in enumerate(annotation_file): annotation = line.strip().split() image_path = annotation[0] image_name = image_path.split('/')[-1] predict_result_path = os.path.join(predicted_epoch_path, str(image_name) + '.txt') original_image = cv2.imread(image_path) image = cv2.cvtColor(original_image, cv2.COLOR_BGR2RGB) # Predict Process ## image_letter, ratio, (dw, dh) = utils.letterbox(image) image_letter = utils.test_image_preprocess(np.copy(image), [INPUT_SIZE, INPUT_SIZE]) image_data = image_letter[np.newaxis, ...].astype(np.float32) batch_data = tf.constant(image_data) bbox_tensors = [] prob_tensors = [] pred_result = D2_model(batch_data,training=False) G_im = pred_result[-1][0] for i in range(ANCHORS.shape[0]): fm = pred_result[i * 2] if i == 0: output_tensors = decode(fm, FLAGS.size // 8, NUM_CLASS, STRIDES, ANCHORS, i, XYSCALE) elif i == 1: output_tensors = decode(fm, FLAGS.size // 16, NUM_CLASS, STRIDES, ANCHORS, 1, XYSCALE) elif i==2: output_tensors = decode(fm, FLAGS.size // 32, NUM_CLASS, STRIDES, ANCHORS, 2, XYSCALE) bbox_tensors.append(output_tensors[0]) prob_tensors.append(output_tensors[1]) pred_bbox = tf.concat(bbox_tensors, axis=1) pred_prob = tf.concat(prob_tensors, axis=1) boxes, pred_conf = filter_boxes(pred_bbox, pred_prob, score_threshold=FLAGS.score_thres, input_shape=tf.constant([FLAGS.size, FLAGS.size])) pred_bbox = tf.concat([boxes, pred_conf], axis=-1) boxes = pred_bbox[:, :, 0:4] pred_conf = pred_bbox[:, :, 4:] boxes, scores, classes, valid_detections = tf.image.combined_non_max_suppression( boxes=tf.reshape(boxes, (tf.shape(boxes)[0], -1, 1, 4)), scores=tf.reshape( pred_conf, (tf.shape(pred_conf)[0], -1, tf.shape(pred_conf)[-1])), max_output_size_per_class=1, max_total_size=1, iou_threshold=FLAGS.iou, score_threshold=FLAGS.score ) boxes, scores, classes, valid_detections = [boxes.numpy(), scores.numpy(), classes.numpy(), valid_detections.numpy()] if num % 1 ==0: G_im = pred_result[-1][0] G_im = G_im * 255 G_im = np.array(G_im).astype(np.int32) image_result = utils.draw_bbox(np.copy(G_im), [boxes, scores, classes, valid_detections]) image_result = image_result[:,:,::-1] filepath = dectect_epoch_path+"/"+ str(image_name) cv2.imwrite(filepath, image_result, [int(cv2.IMWRITE_JPEG_QUALITY), 100]) ################################################################ ################################################################ ################################################################ if os.path.exists(loss_dir_path): shutil.rmtree(loss_dir_path) os.mkdir(loss_dir_path) for epoch in range(first_stage_epochs + second_stage_epochs): if epoch >= first_stage_epochs: for name in T_freeze_layers: try: freeze = D2_model.get_layer(name) freeze_all(freeze) print("Successfully freeze {}...".format(name)) except: print("{} not exist...".format(name)) loss_epoch = np.zeros((steps_per_epoch,10),dtype=np.float32) for index, (image_data, target) in enumerate(trainset): loss_step = train_step(image_data, target) loss_epoch[index] = loss_step mask = loss_epoch[:,0] >0 loss_mean = np.mean(tf.boolean_mask(loss_epoch,mask),0) loss_list_step = {"D2:lr":loss_mean[0],"total_loss":loss_mean[1],"loss/yolo_loss":loss_mean[2], "G_loss":loss_mean[3],"D_loss":loss_mean[4],"loss/giou_loss":loss_mean[5],"loss/conf_loss":loss_mean[6], "loss/prob_loss":loss_mean[7],"loss_content":loss_mean[8],"adv_lambda * loss_adv":loss_mean[9]} loss_epoch_path = os.path.join(loss_dir_path, "epoch-{}".format(epoch) + '.txt') with open(loss_epoch_path, 'w') as f: for vm in loss_list_step.values(): loss_mess = ' '.join([str(vm)]) + '\n' f.write(loss_mess) print("No {} epoch params are {} and {}:".format(epoch,params[0].numpy(),params[1].numpy())) total_loss_result.append(total_loss_metric.result()) loss1_result.append(loss1_metric.result()) loss2_result.append(loss2_metric.result()) total_loss_metric.reset_states() loss1_metric.reset_states() loss2_metric.reset_states() if epoch % FLAGS.save_frequency == 0: D2_model.save_weights(filepath=FLAGS.save_model_dir+"epoch-{}.h5".format(epoch), save_format="h5") full_model.save_weights(filepath=FLAGS.save_model_dir+"Dis"+"epoch-{}".format(epoch), save_format="h5") print("No {} epoch saved successfully...".format(epoch)) #Evaluation model dectect_epoch_path = text_result_path + "-epoch-{}".format(epoch) if os.path.exists(dectect_epoch_path): shutil.rmtree(dectect_epoch_path) os.mkdir(dectect_epoch_path) test_epoch(D2_model,dectect_epoch_path) print("Evaluation completed...") #####draw##### total_loss = np.array(total_loss_result) Yolo_loss = np.array(loss1_result) G_loss = np.array(loss2_result) epochs_range = np.arange(0,epoch+1,1) plt.figure(dpi=1000,num=1,figsize=(6, 3)) plt.plot(epochs_range, total_loss, marker='*',linestyle='-',linewidth=1, markersize=2,label='total_loss') plt.plot(epochs_range, Yolo_loss,marker='o', linestyle='-',linewidth=1, markersize=2,label='Yolo_loss') plt.plot(epochs_range, G_loss, label='Deblur_loss') plt.xlabel('Epoch') plt.ylabel('Loss') plt.legend(loc='upper right') plt.savefig('Tranin_Loss_result.png',bbox_inches="tight",dpi=1000) plt.show()
def main(_argv): input_channel = 3 patience = 30 steps_in_epoch = 0 epoch_loss = 0 prev_minloss = np.inf trainset = DatasetTF(FLAGS, input_channel, is_training=True) testset = DatasetTF(FLAGS, input_channel, is_training=False) # testset = Dataset(FLAGS, input_channel, is_training=False) logdir = "./data/log_tf" isfreeze = False steps_per_epoch = len(trainset) ## for plot graph , early stopping loss_tracker = [] # print("steps_per_epoch:{}".format(steps_per_epoch)) 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 # train_steps = (first_stage_epochs + second_stage_epochs) * steps_per_period input_layer = tf.keras.layers.Input( [cfg.TRAIN.INPUT_SIZE, cfg.TRAIN.INPUT_SIZE, 3]) STRIDES, ANCHORS, NUM_CLASS, XYSCALE = utils.load_config(FLAGS) IOU_LOSS_THRESH = cfg.YOLO.IOU_LOSS_THRESH freeze_layers = utils.load_freeze_layer(FLAGS.model, FLAGS.num_detection_layer) # feature_maps = YOLO(input_layer, NUM_CLASS, FLAGS.model, FLAGS.tiny) feature_maps = YOLOv4_more_tiny(input_layer, NUM_CLASS) bbox_tensors = [] for i, fm in enumerate(feature_maps): # fm shape: (None, 40, 40, 21) bbox_tensor = decode_train(fm, cfg.TRAIN.INPUT_SIZE // 16, NUM_CLASS, STRIDES, ANCHORS, i, XYSCALE) bbox_tensors.append(fm) bbox_tensors.append(bbox_tensor) if cfg.YOLO.NUM_YOLOLAYERS == 3: # yolov4 bbox_tensors = [] for i, fm in enumerate(feature_maps): if i == 0: bbox_tensor = decode_train(fm, cfg.TRAIN.INPUT_SIZE // 8, NUM_CLASS, STRIDES, ANCHORS, i, XYSCALE) elif i == 1: bbox_tensor = decode_train(fm, cfg.TRAIN.INPUT_SIZE // 16, NUM_CLASS, STRIDES, ANCHORS, i, XYSCALE) else: bbox_tensor = decode_train(fm, cfg.TRAIN.INPUT_SIZE // 32, NUM_CLASS, STRIDES, ANCHORS, i, XYSCALE) bbox_tensors.append(fm) bbox_tensors.append(bbox_tensor) elif cfg.YOLO.NUM_YOLOLAYERS == 2: # yolo tiny bbox_tensors = [] for i, fm in enumerate(feature_maps): if i == 0: bbox_tensor = decode_train(fm, cfg.TRAIN.INPUT_SIZE // 16, NUM_CLASS, STRIDES, ANCHORS, i, XYSCALE) else: bbox_tensor = decode_train(fm, cfg.TRAIN.INPUT_SIZE // 32, NUM_CLASS, STRIDES, ANCHORS, i, XYSCALE) bbox_tensors.append(fm) bbox_tensors.append(bbox_tensor) elif cfg.YOLO.NUM_YOLOLAYERS == 1: # custom yolo bbox_tensors = [] bbox_tensor = decode_train( feature_maps[0], cfg.TRAIN.INPUT_SIZE // cfg.YOLO.STRIDES_CUSTOM[0], NUM_CLASS, STRIDES, ANCHORS, 0, XYSCALE) bbox_tensors.append(feature_maps[0]) bbox_tensors.append(bbox_tensor) model = tf.keras.Model(input_layer, bbox_tensors) model.summary() try: if FLAGS.weights.split(".")[len(FLAGS.weights.split(".")) - 1] == "weights": utils.load_weights(model, FLAGS.weights, FLAGS.model, FLAGS.num_detection_layer) else: model.load_weights(FLAGS.weights) print('Restoring weights from: %s ... ' % FLAGS.weights) except Exception: print("Training from scratch") FLAGS.weights = None optimizer = tf.keras.optimizers.Adam() if os.path.exists(logdir): shutil.rmtree(logdir) writer = tf.summary.create_file_writer(logdir) # define training step function @tf.function def train_step(image_data, target): with tf.GradientTape() as tape: pred_result = model(image_data, training=True) ciou_loss = conf_loss = prob_loss = 0 # optimizing process for i in range(len(freeze_layers)): conv, pred = pred_result[i * 2], pred_result[i * 2 + 1] loss_items = compute_loss( pred, conv, target[i][0], target[i][1], STRIDES=STRIDES, NUM_CLASS=NUM_CLASS, # label, bbox IOU_LOSS_THRESH=IOU_LOSS_THRESH, i=i) ciou_loss += loss_items[0] conf_loss += loss_items[1] prob_loss += loss_items[2] total_loss = ciou_loss + conf_loss + prob_loss gradients = tape.gradient(total_loss, model.trainable_variables) optimizer.apply_gradients(zip(gradients, model.trainable_variables)) tf.print("=> STEP ", global_steps, "/", total_steps, " lr: ", optimizer.lr, " ciou_loss: ", ciou_loss, " conf_loss: ", conf_loss, " prob_loss: ", prob_loss, " total_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(tf.cast(lr, tf.float32)) # 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/ciou_loss", ciou_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() return total_loss @tf.function def test_step(image_data, target): with tf.GradientTape() as tape: pred_result = model(image_data, training=True) ciou_loss = conf_loss = prob_loss = 0 # optimizing process for i in range(len(freeze_layers)): conv, pred = pred_result[i * 2], pred_result[i * 2 + 1] loss_items = compute_loss(pred, conv, target[i][0], target[i][1], STRIDES=STRIDES, NUM_CLASS=NUM_CLASS, IOU_LOSS_THRESH=IOU_LOSS_THRESH, i=i) # metric_items = compute_map(pred, conv, target[i][0], target[i][1], STRIDES=STRIDES, NUM_CLASS=NUM_CLASS, i=i) # loss_items = compute_loss(pred, conv, target[1][0], target[1][1], STRIDES=STRIDES, NUM_CLASS=NUM_CLASS, IOU_LOSS_THRESH=IOU_LOSS_THRESH, i=i) ciou_loss += loss_items[0] conf_loss += loss_items[1] prob_loss += loss_items[2] total_loss = ciou_loss + conf_loss + prob_loss tf.print("==> TEST STEP: ", global_steps) tf.print(" ciou_loss: ", ciou_loss) tf.print(" conf_loss: ", conf_loss) tf.print(" prob_loss: ", prob_loss) tf.print("==> total_loss", total_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)) it = iter(trainset.get_dataset()) it_test = iter(testset.get_dataset()) start = time.time() for epoch in range(first_stage_epochs + second_stage_epochs): train_loss = 0. if epoch < first_stage_epochs: if not isfreeze: isfreeze = True for name in freeze_layers: try: # try구문 추가 freeze = model.get_layer(name) freeze_all(freeze) except ValueError: pass elif epoch >= first_stage_epochs: if isfreeze: isfreeze = False for name in freeze_layers: try: # try구문 추가 freeze = model.get_layer(name) freeze_all(freeze) except ValueError: pass # freeze = model.get_layer(name) # unfreeze_all(freeze) for i in range(trainset.steps_for_train): data = it.get_next() ### check input train images # image_data = data[0] # label_data = np.array(data[1]) # # bboxes_list = np.array(label_data[0][1]) # label_list = np.array(label_data[0][0]) # for batch_idx, bboxes in enumerate(bboxes_list): # class_inds = [] # check = np.array(image_data[batch_idx]).reshape(cfg.TRAIN.INPUT_SIZE, cfg.TRAIN.INPUT_SIZE, 3) # # print(r"D:\tf_data_\preprocessed_{}.npy".format(batch_idx)) # # np.save(r"D:\tf_data_\image_{}.npy".format(batch_idx), check) # # np.save(r"D:\tf_data_\preprocessed_{}.npy".format(batch_idx), label_list[batch_idx]) # # np.save(r"D:\tf_data_\bboxes_{}.npy".format(batch_idx), bboxes_list[batch_idx]) # label = np.array(label_list[batch_idx]) # # class_ = label[...,5:] # class_ = np.array(class_).flatten() # class_ = np.where(class_ > 0.1, class_, 0) # # for class_idx, class_label in enumerate(class_): # # print("class_label:{}".format(class_label)) # if class_label > 0: # class_inds.append(class_idx % cfg.YOLO.NUM_CLASSES) # # bboxes = np.array(bboxes) # for bbox_idx, bbox in enumerate(bboxes): # half_h = bbox[3] / 2 # half_w = bbox[2] / 2 # # if half_h + half_w > 0: # cv2.rectangle(check, (int(bbox[0] - half_w), int(bbox[1] - half_h)), # (int(bbox[0] + half_w), int(bbox[1] + half_h)), color=(0, 255, 0), thickness=3) # cv2.putText(check, text="{}".format(class_inds), # org=(int(bbox[0] + half_w)-10, int(bbox[1] + half_h)-30), thickness=2, color=(0, 255, 0), # fontFace=cv2.FONT_HERSHEY_SIMPLEX, fontScale=0.7) # # cv2.imshow("check_aug", check) # cv2.waitKey(0) # cv2.destroyAllWindows() # cv2.imwrite(os.path.join(r"D:\Public\JHS\SIMPLE_DATA_INPUT_CHECK", "{}.jpg".format(batch_idx)), check*255) train_loss = train_step(data[0], data[1]) # epoch_loss += train_loss # for i in range(testset.steps_for_train): # data = it_test.get_next() # test_step(data[0], data[1]) # CHECK TEST # image_data = data[0] # label_data = data[1] # bboxes_list = label_data[0][1] # # for batch_idx, bboxes in enumerate(bboxes_list): # check = np.array(image_data[batch_idx]).reshape(640, 640, 3) # for bbox in bboxes: # half_h = bbox[3] / 2 # half_w = bbox[2] / 2 # # bbox_upper_left = (bbox[0]) # if np.sum(bbox) > 0: # cv2.rectangle(check, (bbox[0] - half_w, bbox[1] - half_h),(bbox[0] + half_w,bbox[1]+half_h), color=(0,255,0), thickness=3) # cv2.imshow("check_aug TEST", check) # cv2.waitKey(0) # cv2.destroyAllWindows() # for image_data, target in testset: # test_step(image_data, target) # for loss graph """ if steps_in_epoch >= steps_per_epoch: loss_tracker.append(losses_in_epoch / steps_per_epoch) ## 밑으로 # plt.plot(loss_tracker) # plt.xlabel('epoch') # plt.ylabel('loss') # # plt.show() # plt.savefig("D:/checkpoint/loss.png") """ # early stopping """ epoch_loss = epoch_loss / steps_per_epoch loss_tracker.append(epoch_loss) if prev_minloss > epoch_loss: # save best weight prev_minloss = epoch_loss if epoch > 2000: # 최소학습 에폭 model.save("D:\ckpt_best") print("{} epoch save best weights".format(epoch)) if len(loss_tracker) > patience: print("check loss_tracker len:{}".format(len(loss_tracker))) if loss_tracker[0] > loss_tracker[-1]: loss_tracker.pop(0) else: print("total loss didn't decreased during {} epochs. train stop".format(patience)) return epoch_loss = 0. else: epoch_loss += train_loss steps_in_epoch = 0 """ if (epoch + 1) % 100 == 0: model.save("D:\yolov4-tflite-train_tf-epoch{}".format(epoch + 1)) print("{} epoch model saved".format(epoch + 1)) print("consumed time: {}".format(time.time() - start)) model.save(r"D:\yolov4-tflite-train_tf-last")