Hi,
I have shared the configuration files for every step from the retraining to the engine file creation.
Further i will share the retrained model on custom data set along with the engine files that we created and a testing video in direct message.
Step 1: Retraining of the coco based pre-trained maskrcnn model on custom data set.
Following are the configurations for retraining of the model:
Configurations:
BACKBONE resnet101
BACKBONE_STRIDES [4, 8, 16, 32, 64]
BATCH_SIZE 1
BBOX_STD_DEV [0.1 0.1 0.2 0.2]
COMPUTE_BACKBONE_SHAPE None
DETECTION_MAX_INSTANCES 100
DETECTION_MIN_CONFIDENCE 0.7
DETECTION_NMS_THRESHOLD 0.3
FPN_CLASSIF_FC_LAYERS_SIZE 1024
GPU_COUNT 1
GRADIENT_CLIP_NORM 5.0
IMAGES_PER_GPU 1
IMAGE_CHANNEL_COUNT 3
IMAGE_MAX_DIM 1024
IMAGE_META_SIZE 14
IMAGE_MIN_DIM 800
IMAGE_MIN_SCALE 0
IMAGE_RESIZE_MODE square
IMAGE_SHAPE [1024 1024 3]
LEARNING_MOMENTUM 0.9
LEARNING_RATE 0.001
LOSS_WEIGHTS {'rpn_class_loss': 1.0, 'rpn_bbox_loss': 1.0, 'mrcnn_class_loss': 1.0, 'mrcnn_bbox_loss': 1.0, 'mrcnn_mask_loss': 1.0}
MASK_POOL_SIZE 14
MASK_SHAPE [28, 28]
MAX_GT_INSTANCES 100
MEAN_PIXEL [123.7 116.8 103.9]
MINI_MASK_SHAPE (56, 56)
NAME cafe
NUM_CLASSES 2
POOL_SIZE 7
POST_NMS_ROIS_INFERENCE 1000
POST_NMS_ROIS_TRAINING 2000
PRE_NMS_LIMIT 1024
ROI_POSITIVE_RATIO 0.33
RPN_ANCHOR_RATIOS [0.5, 1, 2]
RPN_ANCHOR_SCALES (32, 64, 128, 256, 512)
RPN_ANCHOR_STRIDE 1
RPN_BBOX_STD_DEV [0.1 0.1 0.2 0.2]
RPN_NMS_THRESHOLD 0.7
RPN_TRAIN_ANCHORS_PER_IMAGE 256
STEPS_PER_EPOCH 1000
TOP_DOWN_PYRAMID_SIZE 256
TRAIN_BN False
TRAIN_ROIS_PER_IMAGE 200
USE_MINI_MASK True
USE_RPN_ROIS True
VALIDATION_STEPS 100
WEIGHT_DECAY 0.0001
The retained model was then tested then and the results are shared with you and they were really good.
Step 2: H5 to UFF conversion
As mentioned earlier the following steps were followed for uff conversion https://github.com/NVIDIA/TensorRT/tree/master/samples/opensource/sampleUffMaskRCNN
Tensorrt versions in this conversion container:
ii graphsurgeon-tf 6.0.1-1+cuda10.1 amd64 GraphSurgeon for TensorRT package
ii libnvinfer-bin 6.0.1-1+cuda10.1 amd64 TensorRT binaries
ii libnvinfer-dev 6.0.1-1+cuda10.1 amd64 TensorRT development libraries and headers
ii libnvinfer-plugin-dev 6.0.1-1+cuda10.1 amd64 TensorRT plugin libraries
ii libnvinfer-plugin6 6.0.1-1+cuda10.1 amd64 TensorRT plugin libraries
ii libnvinfer6 6.0.1-1+cuda10.1 amd64 TensorRT runtime libraries
ii libnvonnxparsers-dev 6.0.1-1+cuda10.1 amd64 TensorRT ONNX libraries
ii libnvonnxparsers6 6.0.1-1+cuda10.1 amd64 TensorRT ONNX libraries
ii libnvparsers-dev 6.0.1-1+cuda10.1 amd64 TensorRT parsers libraries
ii libnvparsers6 6.0.1-1+cuda10.1 amd64 TensorRT parsers libraries
ii python3-libnvinfer 6.0.1-1+cuda10.1 amd64 Python 3 bindings for TensorRT
ii python3-libnvinfer-dev 6.0.1-1+cuda10.1 amd64 Python 3 development package for TensorRT
ii uff-converter-tf 6.0.1-1+cuda10.1 amd64 UFF converter for TensorRT package
mrcnn_to_trt_single.py
from keras.models import model_from_json, Model
from keras import backend as K
from keras.layers import Input, Lambda
from tensorflow.python.framework import graph_util
from tensorflow.python.framework import graph_io
from mrcnn.model import *
import mrcnn.model as modellib
from mrcnn.config import Config
import sys
import os
ROOT_DIR = os.path.abspath("./")
LOG_DIR = os.path.join(ROOT_DIR, "logs")
import argparse
import os
import uff
def parse_command_line_arguments(args=None):
parser = argparse.ArgumentParser(prog='keras_to_trt', description='Convert trained keras .hdf5 model to trt .uff')
parser.add_argument(
'-w',
'--weights',
type=str,
default=None,
required=True,
help="The checkpoint weights file of keras model."
)
parser.add_argument(
'-o',
'--output_file',
type=str,
default=None,
required=True,
help="The path to output .uff file."
)
parser.add_argument(
'-l',
'--list-nodes',
action='store_true',
help="show list of nodes contained in converted pb"
)
parser.add_argument(
'-p',
'--preprocessor',
type=str,
default=False,
help="The preprocess function for converting tf node to trt plugin"
)
return parser.parse_args(args)
class CocoConfig(Config):
"""Configuration for training on MS COCO.
Derives from the base Config class and overrides values specific
to the COCO dataset.
"""
# Give the configuration a recognizable name
NAME = "cafe"
# We use a GPU with 12GB memory, which can fit two images.
# Adjust down if you use a smaller GPU.
IMAGES_PER_GPU = 2
# Uncomment to train on 8 GPUs (default is 1)
# GPU_COUNT = 8
# Number of classes (including background)
NUM_CLASSES = 1 + 1 # COCO has 80 classes
class InferenceConfig(CocoConfig):
# Set batch size to 1 since we'll be running inference on
# one image at a time. Batch size = GPU_COUNT * IMAGES_PER_GPU
GPU_COUNT = 1
IMAGES_PER_GPU = 1
NUM_CLASSES = 1 + 1
def main(args=None):
K.set_image_data_format('channels_first')
K.set_learning_phase(0)
args = parse_command_line_arguments(args)
model_weights_path = args.weights
output_file_path = args.output_file
list_nodes = args.list_nodes
config = InferenceConfig()
config.display()
model = modellib.MaskRCNN(mode="inference", model_dir=LOG_DIR, config=config).keras_model
model.load_weights(model_weights_path, by_name=True)
model_A = Model(inputs=model.input, outputs=model.get_layer('mrcnn_mask').output)
model_A.summary()
output_nodes = ['mrcnn_detection', "mrcnn_mask/Sigmoid"]
convert_model(model_A, output_file_path, output_nodes, preprocessor=args.preprocessor,
text=True, list_nodes=list_nodes)
def convert_model(inference_model, output_path, output_nodes=[], preprocessor=None, text=False,
list_nodes=False):
# convert the keras model to pb
orig_output_node_names = [node.op.name for node in inference_model.outputs]
print("The output names of tensorflow graph nodes: {}".format(str(orig_output_node_names)))
sess = K.get_session()
constant_graph = graph_util.convert_variables_to_constants(
sess,
sess.graph.as_graph_def(),
orig_output_node_names)
temp_pb_path = "../temp.pb"
graph_io.write_graph(constant_graph, os.path.dirname(temp_pb_path), os.path.basename(temp_pb_path),
as_text=False)
predefined_output_nodes = output_nodes
if predefined_output_nodes != []:
trt_output_nodes = predefined_output_nodes
else:
trt_output_nodes = orig_output_node_names
# convert .pb to .uff
uff.from_tensorflow_frozen_model(
temp_pb_path,
output_nodes=trt_output_nodes,
preprocessor=preprocessor,
text=text,
list_nodes=list_nodes,
output_filename=output_path,
debug_mode = False
)
os.remove(temp_pb_path)
if __name__ == "__main__":
main()
Further in the config.h the number of classes was updated.
Step 3: Uff to engine conversion
In this step a deepstream 4.0.2 container was established with the same tensorrt 6 as in the previous container.
uff_to_engine.py
import os
os.environ['CUDA_VISIBLE_DEVICES'] = '0' # gpu ID
import tensorrt as trt
tile_nums = (3,5)
batch_size = tile_nums[0]*tile_nums[1]+1
batch_size = 1
MODEL_INPUT = 'input_image'
MODEL_OUTPUTS = ["mrcnn_detection", "mrcnn_mask/Sigmoid"]
INPUT_SHAPE = (3, 1024, 1024)
FP16_FLAG = False
model_file = './uff/mask_rcnn_cafe_0018_81_370train.uff'
batch_size = batch_size
TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
# must initialize plugins (custom layers for MRCNN)
trt.init_libnvinfer_plugins(TRT_LOGGER, '')
print("Building Engine file")
# Create context, engine
with trt.Builder(TRT_LOGGER) as builder, builder.create_network() as network, trt.UffParser() as parser:
parser.register_input(name=MODEL_INPUT, shape=INPUT_SHAPE, order=trt.UffInputOrder.NCHW)
for output in MODEL_OUTPUTS:
parser.register_output(name=output)
parser.parse(file=model_file, network=network, weights_type=trt.tensorrt.DataType.FLOAT)
builder.max_batch_size = batch_size
builder.max_workspace_size = int(8e8)
builder.fp16_mode = FP16_FLAG
engine = builder.build_cuda_engine(network=network)
context = engine.create_execution_context()
with open("mask_rcnn_cafe_0018_81_370_batch1.engine", "wb") as f:
f.write(engine.serialize())
print("Finished")
The final generated engine file is also shared with you with i have integrated with deepstream and the results were not as good as outside deepstream.
Further i have tried all the above steps with 81 classes that were default and no difference .
Thanks