diff --git a/docker/Dockerfile b/docker/Dockerfile new file mode 100644 index 0000000..d582241 --- /dev/null +++ b/docker/Dockerfile @@ -0,0 +1,43 @@ +## Basic Docker Enviroment for correct OpenPCDet for CUDA-PointPillars +## +## Prerequisites: +## - docker +## - nvidia driver +## - nvidia container toolkit +## +## Build with: +## docker build -t open_pcdet:latest . +## +## Save with: +## docker save open_pcdet:latest | gzip > open_pcdet.tar.gz +## +## Run with: +## docker run --network host --gpus all --name PointPillars --rm -it open_pcdet:latest +## Debug with: +## docker run --network host --gpus all --name PointPillars --entrypoint /bin/bash -it open_pcdet:latest + +FROM nvcr.io/nvidia/pytorch:20.12-py3 + +## Meta information +LABEL cuda.version="11.1" maintainers="Robert Belshaw " + +WORKDIR / +RUN DEBIAN_FRONTEND=noninteractive apt-get update +RUN pip install --upgrade pip wheel setuptools +RUN pip install spconv-cu113 +RUN git clone --branch v0.5.2 https://github.com/open-mmlab/OpenPCDet.git +WORKDIR /OpenPCDet +RUN python setup.py develop +RUN DEBIAN_FRONTEND=noninteractive apt-get install ffmpeg libsm6 libxext6 -y +RUN DEBIAN_FRONTEND=noninteractive pip install onnxsim nvidia-pyindex +RUN DEBIAN_FRONTEND=noninteractive pip install onnx-graphsurgeon +RUN DEBIAN_FRONTEND=noninteractive pip install kornia==0.4.1 + +# Docker clean-up +RUN rm -rf /var/lib/apt/lists/* + +## Create entry point for image +WORKDIR / +COPY ./entrypoint.sh . +ENTRYPOINT ["/bin/bash"] +CMD ["/bin/bash", "-c"] \ No newline at end of file diff --git a/docker/README.md b/docker/README.md new file mode 100644 index 0000000..380b965 --- /dev/null +++ b/docker/README.md @@ -0,0 +1,9 @@ +# Cuda PointPillars Exporter Docker + +## usage example: + +``` +bash export.sh -m ../model/pointpillar_7728.pth -c ../tool/cfgs/kitti_models/pointpillar.yaml +``` + +tested with docker-compose 1.25.0 and docker 20.10.2120.10.21 \ No newline at end of file diff --git a/docker/docker-compose.yml b/docker/docker-compose.yml new file mode 100644 index 0000000..c07f000 --- /dev/null +++ b/docker/docker-compose.yml @@ -0,0 +1,16 @@ +# email: rbelshaw@sagarobotics.com +# Need to docker login +version: '2.4' +services: + open_pcdet_export: + image: open_pcdet:latest + build: . + user: ${_UID}:${_UID} + restart: on-failure:2 + entrypoint: /bin/bash ./entrypoint.sh + volumes: + - ./entrypoint.sh:/entrypoint.sh + - ../tool:/exporter + - ./${_CKPT}:/checkpoint.pth + - ../tool/cfgs/dataset_configs/kitti_dataset.yaml:/exporter/dataset.yaml + - ./${_CONFIG}:/pointpillar.yaml \ No newline at end of file diff --git a/docker/entrypoint.sh b/docker/entrypoint.sh new file mode 100644 index 0000000..06795cb --- /dev/null +++ b/docker/entrypoint.sh @@ -0,0 +1,2 @@ +cd /exporter +python exporter.py --cfg_file /pointpillar.yaml --ckpt /checkpoint.pth --data_path /exporter/dataset.yaml \ No newline at end of file diff --git a/docker/export.sh b/docker/export.sh new file mode 100755 index 0000000..26e4f9b --- /dev/null +++ b/docker/export.sh @@ -0,0 +1,42 @@ +#!/usr/bin/env bash +set -e +cd "$(cd -P -- "$(dirname -- "$0")" && pwd -P)" + +## Usage ## +# +# bash export.sh -m path/to/ckpt.pth -c path/to/pointpillar.yaml +# +# note: path is relative to this script +# + +while getopts m:c: flag +do + case "${flag}" in + m) _CKPT=${OPTARG};; + c) _CONFIG=${OPTARG};; + esac +done + +if [ ! -f "${_CONFIG}" ] +then + echo "Config .yaml Does not exist at: ${_CONFIG}" + exit +fi + +if [ ! -f "${_CKPT}" ] +then + echo "Ckpt .pth Does not exist at: ${_CKPT}" + exit +fi + +_CONFIG=$_CONFIG +export _CONFIG +echo "CONFIG: $_CONFIG"; +_CKPT=$_CKPT +export _CKPT +echo "CKPT: $_CKPT"; + +export _UID=$(id -u) + +# Run backend +docker-compose run --rm open_pcdet_export diff --git a/tool/exporter.py b/tool/exporter.py index 12fff84..57dd534 100644 --- a/tool/exporter.py +++ b/tool/exporter.py @@ -101,9 +101,32 @@ def main(): with torch.no_grad(): MAX_VOXELS = 10000 + NUMBER_OF_CLASSES = len(cfg.CLASS_NAMES) + + MAX_POINTS_PER_VOXEL = None + + DATA_PROCESSOR = cfg.DATA_CONFIG.DATA_PROCESSOR + POINT_CLOUD_RANGE = cfg.DATA_CONFIG.POINT_CLOUD_RANGE + for i in DATA_PROCESSOR: + if i['NAME'] == "transform_points_to_voxels": + MAX_POINTS_PER_VOXEL = i['MAX_POINTS_PER_VOXEL'] + VOXEL_SIZES = i['VOXEL_SIZE'] + break + + + if MAX_POINTS_PER_VOXEL == None: + logger.info('Could Not Parse Config... Exiting') + import sys + sys.exit() + + VOXEL_SIZE_X = abs(POINT_CLOUD_RANGE[0] - POINT_CLOUD_RANGE[3]) / VOXEL_SIZES[0] + VOXEL_SIZE_Y = abs(POINT_CLOUD_RANGE[1] - POINT_CLOUD_RANGE[4]) / VOXEL_SIZES[1] + + FEATURE_SIZE_X = VOXEL_SIZE_X / 2 #Is this number of bins? + FEATURE_SIZE_Y = VOXEL_SIZE_Y / 2 dummy_voxels = torch.zeros( - (MAX_VOXELS, 32, 4), + (MAX_VOXELS, MAX_POINTS_PER_VOXEL, 4), dtype=torch.float32, device='cuda:0') @@ -121,7 +144,7 @@ def main(): dummy_input['voxels'] = dummy_voxels dummy_input['voxel_num_points'] = dummy_voxel_num dummy_input['voxel_coords'] = dummy_voxel_idxs - dummy_input['batch_size'] = 1 + dummy_input['batch_size'] = torch.tensor(1) torch.onnx.export(model, # model being run dummy_input, # model input (or a tuple for multiple inputs) @@ -135,12 +158,12 @@ def main(): ) onnx_raw = onnx.load("./pointpillar_raw.onnx") # load onnx model - onnx_trim_post = simplify_postprocess(onnx_raw) + onnx_trim_post = simplify_postprocess(onnx_raw, FEATURE_SIZE_X, FEATURE_SIZE_Y, NUMBER_OF_CLASSES) onnx_simp, check = simplify(onnx_trim_post) assert check, "Simplified ONNX model could not be validated" - onnx_final = simplify_preprocess(onnx_simp) + onnx_final = simplify_preprocess(onnx_simp, VOXEL_SIZE_X, VOXEL_SIZE_Y, MAX_POINTS_PER_VOXEL) onnx.save(onnx_final, "pointpillar.onnx") print('finished exporting onnx') diff --git a/tool/simplifier_onnx.py b/tool/simplifier_onnx.py index d7b5457..76c94e4 100644 --- a/tool/simplifier_onnx.py +++ b/tool/simplifier_onnx.py @@ -18,7 +18,7 @@ import onnx_graphsurgeon as gs @gs.Graph.register() -def replace_with_clip(self, inputs, outputs): +def replace_with_clip(self, inputs, outputs, voxel_array): for inp in inputs: inp.outputs.clear() @@ -26,7 +26,7 @@ def replace_with_clip(self, inputs, outputs): out.inputs.clear() op_attrs = dict() - op_attrs["dense_shape"] = np.array([496,432]) + op_attrs["dense_shape"] = voxel_array return self.layer(name="PPScatter_0", op="PPScatterPlugin", inputs=inputs, outputs=outputs, attrs=op_attrs) @@ -36,13 +36,13 @@ def loop_node(graph, current_node, loop_time=0): current_node = next_node return next_node -def simplify_postprocess(onnx_model): +def simplify_postprocess(onnx_model, FEATURE_SIZE_X, FEATURE_SIZE_Y, NUMBER_OF_CLASSES): print("Use onnx_graphsurgeon to adjust postprocessing part in the onnx...") graph = gs.import_onnx(onnx_model) - cls_preds = gs.Variable(name="cls_preds", dtype=np.float32, shape=(1, 248, 216, 18)) - box_preds = gs.Variable(name="box_preds", dtype=np.float32, shape=(1, 248, 216, 42)) - dir_cls_preds = gs.Variable(name="dir_cls_preds", dtype=np.float32, shape=(1, 248, 216, 12)) + cls_preds = gs.Variable(name="cls_preds", dtype=np.float32, shape=(1, int(FEATURE_SIZE_Y), int(FEATURE_SIZE_X), 2*NUMBER_OF_CLASSES*NUMBER_OF_CLASSES)) + box_preds = gs.Variable(name="box_preds", dtype=np.float32, shape=(1, int(FEATURE_SIZE_Y), int(FEATURE_SIZE_X), 14*NUMBER_OF_CLASSES)) + dir_cls_preds = gs.Variable(name="dir_cls_preds", dtype=np.float32, shape=(1, int(FEATURE_SIZE_Y), int(FEATURE_SIZE_X), 4*NUMBER_OF_CLASSES)) tmap = graph.tensors() new_inputs = [tmap["voxels"], tmap["voxel_idxs"], tmap["voxel_num"]] @@ -73,18 +73,20 @@ def simplify_postprocess(onnx_model): return gs.export_onnx(graph) -def simplify_preprocess(onnx_model): +def simplify_preprocess(onnx_model, VOXEL_SIZE_X, VOXEL_SIZE_Y, MAX_POINTS_PER_VOXEL): print("Use onnx_graphsurgeon to modify onnx...") graph = gs.import_onnx(onnx_model) tmap = graph.tensors() MAX_VOXELS = tmap["voxels"].shape[0] + VOXEL_ARRAY = np.array([int(VOXEL_SIZE_X),int(VOXEL_SIZE_Y)]) + # voxels: [V, P, C'] # V is the maximum number of voxels per frame # P is the maximum number of points per voxel # C' is the number of channels(features) per point in voxels. - input_new = gs.Variable(name="voxels", dtype=np.float32, shape=(MAX_VOXELS, 32, 10)) + input_new = gs.Variable(name="voxels", dtype=np.float32, shape=(MAX_VOXELS, MAX_POINTS_PER_VOXEL, 10)) # voxel_idxs: [V, 4] # V is the maximum number of voxels per frame @@ -113,7 +115,7 @@ def simplify_preprocess(onnx_model): graph.inputs.append(Y) inputs = [last_node_pillarvfe.outputs[0], X, Y] outputs = [first_node_after_pillarscatter.inputs[0]] - graph.replace_with_clip(inputs, outputs) + graph.replace_with_clip(inputs, outputs, VOXEL_ARRAY) # Remove the now-dangling subgraph. graph.cleanup().toposort()