def main(): # dataset dataset = TrashData() dataset.load('data/dataset_surfrider_cleaned') # model model = faster_rcnn.FasterRCNN(num_classes=dataset.class_count) optimizer = keras.optimizers.SGD(1e-3, momentum=0.9, nesterov=True) for epoch in range(100): loss_history = [] for (batch, inputs) in enumerate(train_tf_dataset): batch_imgs, batch_metas, batch_bboxes, batch_labels = inputs with tf.GradientTape() as tape: rpn_class_loss, rpn_bbox_loss, rcnn_class_loss, rcnn_bbox_loss = model( (batch_imgs, batch_metas, batch_bboxes, batch_labels), training=True) loss_value = rpn_class_loss + rpn_bbox_loss + rcnn_class_loss + rcnn_bbox_loss grads = tape.gradient(loss_value, model.trainable_variables) optimizer.apply_gradients(zip(grads, model.trainable_variables)) loss_history.append(loss_value.numpy()) if batch % 10 == 0: print('epoch', epoch, batch, np.mean(loss_history))
train_datasets = ZiptrainDataset( tf_record_path, 1, 96, crop_size=crop_size, roi_path='D:/datasets/bjod/roi_test/').prepare(True) val_train = Zipvaluedata( tf_record_path, crop_size=crop_size).prepare('train_api_97.record') val_test = Zipvaluedata( tf_record_path, crop_size=crop_size).prepare('val_api_19.record') one_imgs, one_metas, one_bboxes, one_labels, _ = next(iter(val_train)) one_imgs = tf.expand_dims(tf.cast(one_imgs[0], tf.float32), axis=0) one_metas = tf.expand_dims(tf.cast(one_metas[0], tf.float32), axis=0) model = faster_rcnn.FasterRCNN(num_classes=2) _ = model((one_imgs, one_metas), training=False) model_ori = faster_rcnn.FasterRCNN(num_classes=81) _ = model_ori((one_imgs, one_metas), training=False) model_ori.load_weights( './weights/faster_rcnn_resnet101_fpn_coco2017_map35.h5', by_name=True) model.backbone.set_weights(model_ori.backbone.get_weights()) model.neck.set_weights(model_ori.neck.get_weights()) model.rpn_head.set_weights(model_ori.rpn_head.get_weights()) model.roi_align.set_weights(model_ori.roi_align.get_weights()) # print(cc) model.summary() def __init__(self,
pad_mode='fixed', mean=img_mean, std=img_std, scale=(640, 832), debug=True) train_generator = data_generator.DataGenerator(train_dataset) # load model from detection.models.detectors import faster_rcnn strategy = tf.distribute.MirroredStrategy() with strategy.scope(): model = faster_rcnn.FasterRCNN( num_classes=len(train_dataset.get_categories())) img, img_meta, _, _ = train_dataset[0] batch_imgs = tf.Variable(np.expand_dims(img, 0)) batch_metas = tf.Variable(np.expand_dims(img_meta, 0)) _ = model((batch_imgs, batch_metas), training=False) model.load_weights('weights/faster_rcnn.h5', by_name=True) # use tf.data global_batch_size = 2 with strategy.scope(): train_tf_dataset = tf.data.Dataset.from_generator( train_generator, (tf.float32, tf.float32, tf.float32, tf.int32)) train_tf_dataset = train_tf_dataset.padded_batch(
flip_ratio=0.5, pad_mode='fixed', mean=img_mean, std=img_std, scale=(800, 1216)) num_classes = len(train_dataset.get_categories()) train_generator = data_generator.DataGenerator(train_dataset) train_tf_dataset = tf.data.Dataset.from_generator( train_generator, (tf.float32, tf.float32, tf.float32, tf.int32)) train_tf_dataset = train_tf_dataset.batch(batch_size).prefetch(100).shuffle( 100) # train_tf_dataset = train_tf_dataset.padded_batch( # batch_size, padded_shapes=([None, None, None], [None], [None, None], [None])) model = faster_rcnn.FasterRCNN(num_classes=num_classes) optimizer = keras.optimizers.SGD(1e-3, momentum=0.9, nesterov=True) img, img_meta, bboxes, labels = train_dataset[6] # [N, 4], shape:[N]=data:[62] rgb_img = np.round(img + img_mean) ori_img = get_original_image(img, img_meta, img_mean) # visualize.display_instances(rgb_img, bboxes, labels, train_dataset.get_categories()) batch_imgs = tf.convert_to_tensor(np.expand_dims(img, 0)) # [1, 1216, 1216, 3] batch_metas = tf.convert_to_tensor(np.expand_dims(img_meta, 0)) # [1, 11] # dummpy forward to build network variables _ = model((batch_imgs, batch_metas), training=False) proposals = model.simple_test_rpn(img, img_meta) res = model.simple_test_bboxes(img, img_meta, proposals)