├── .gitignore
├── LICENSE
├── README.md
├── SECURITY.md
├── data
├── kitti
│ └── ImageSets
│ │ ├── test.txt
│ │ ├── train.txt
│ │ └── val.txt
└── waymo
│ └── ImageSets
│ ├── infer_time.txt
│ ├── train.txt
│ └── val.txt
├── docs
├── DEMO.md
├── GETTING_STARTED.md
├── INSTALL.md
└── latency.png
├── pcdet
├── __init__.py
├── config.py
├── datasets
│ ├── __init__.py
│ ├── augmentor
│ │ ├── augmentor_utils.py
│ │ ├── data_augmentor.py
│ │ └── database_sampler.py
│ ├── dataset.py
│ ├── kitti
│ │ ├── kitti_dataset.py
│ │ ├── kitti_object_eval_python
│ │ │ ├── LICENSE
│ │ │ ├── README.md
│ │ │ ├── eval.py
│ │ │ ├── evaluate.py
│ │ │ ├── kitti_common.py
│ │ │ └── rotate_iou.py
│ │ └── kitti_utils.py
│ ├── nuscenes
│ │ ├── nuscenes_dataset.py
│ │ └── nuscenes_utils.py
│ ├── processor
│ │ ├── data_processor.py
│ │ └── point_feature_encoder.py
│ └── waymo
│ │ ├── waymo_dataset.py
│ │ ├── waymo_eval.py
│ │ └── waymo_utils.py
├── models
│ ├── __init__.py
│ ├── backbones_2d
│ │ ├── __init__.py
│ │ ├── base_bev_backbone.py
│ │ └── map_to_bev
│ │ │ ├── __init__.py
│ │ │ ├── conv2d_collapse.py
│ │ │ ├── height_compression.py
│ │ │ ├── pillar_reencoding.py
│ │ │ └── pointpillar_scatter.py
│ ├── backbones_3d
│ │ ├── __init__.py
│ │ ├── pfe
│ │ │ ├── __init__.py
│ │ │ └── voxel_set_abstraction.py
│ │ ├── pillar_adaptor
│ │ │ ├── __init__.py
│ │ │ ├── pillar_adaptor_template.py
│ │ │ ├── sparse_interp_adaptor.py
│ │ │ └── sparse_naive_adaptor.py
│ │ ├── pointnet2_backbone.py
│ │ ├── spconv_backbone.py
│ │ ├── spconv_unet.py
│ │ └── vfe
│ │ │ ├── __init__.py
│ │ │ ├── dynamic_kp_vfe.py
│ │ │ ├── dynamic_mean_vfe.py
│ │ │ ├── dynamic_pillar_vfe.py
│ │ │ ├── image_vfe.py
│ │ │ ├── image_vfe_modules
│ │ │ ├── f2v
│ │ │ │ ├── __init__.py
│ │ │ │ ├── frustum_grid_generator.py
│ │ │ │ ├── frustum_to_voxel.py
│ │ │ │ └── sampler.py
│ │ │ └── ffn
│ │ │ │ ├── __init__.py
│ │ │ │ ├── ddn
│ │ │ │ ├── __init__.py
│ │ │ │ ├── ddn_deeplabv3.py
│ │ │ │ └── ddn_template.py
│ │ │ │ ├── ddn_loss
│ │ │ │ ├── __init__.py
│ │ │ │ ├── balancer.py
│ │ │ │ └── ddn_loss.py
│ │ │ │ └── depth_ffn.py
│ │ │ ├── mean_vfe.py
│ │ │ ├── pillar_vfe.py
│ │ │ └── vfe_template.py
│ ├── dense_heads
│ │ ├── __init__.py
│ │ ├── anchor_head_multi.py
│ │ ├── anchor_head_separate.py
│ │ ├── anchor_head_single.py
│ │ ├── anchor_head_template.py
│ │ ├── center_head.py
│ │ ├── point_head_box.py
│ │ ├── point_head_simple.py
│ │ ├── point_head_template.py
│ │ ├── point_intra_part_head.py
│ │ └── target_assigner
│ │ │ ├── anchor_generator.py
│ │ │ ├── atss_target_assigner.py
│ │ │ └── axis_aligned_target_assigner.py
│ ├── detectors
│ │ ├── PartA2_net.py
│ │ ├── __init__.py
│ │ ├── caddn.py
│ │ ├── centerpoint.py
│ │ ├── detector3d_template.py
│ │ ├── point_rcnn.py
│ │ ├── pointpillar.py
│ │ ├── pv_rcnn.py
│ │ ├── pv_rcnn_plusplus.py
│ │ ├── second_net.py
│ │ ├── second_net_iou.py
│ │ └── voxel_rcnn.py
│ ├── kd_adapt_block
│ │ ├── __init__.py
│ │ └── kd_adapt_block.py
│ ├── kd_heads
│ │ ├── __init__.py
│ │ ├── anchor_head
│ │ │ ├── __init__.py
│ │ │ ├── anchor_feature_kd_head.py
│ │ │ ├── anchor_kd_head.py
│ │ │ ├── anchor_label_kd_head.py
│ │ │ └── anchor_logit_kd_head.py
│ │ ├── center_head
│ │ │ ├── __init__.py
│ │ │ ├── center_feature_kd_head.py
│ │ │ ├── center_kd_head.py
│ │ │ ├── center_label_assign_kd_head.py
│ │ │ ├── center_logit_kd_head.py
│ │ │ ├── center_roi_kd_head.py
│ │ │ └── center_vfe_kd_head.py
│ │ └── kd_head.py
│ ├── model_utils
│ │ ├── basic_block_2d.py
│ │ ├── batch_norm_utils.py
│ │ ├── centernet_utils.py
│ │ ├── efficientnet_utils.py
│ │ ├── model_nms_utils.py
│ │ ├── positional_encoding.py
│ │ └── rotated_roi_grid_pool.py
│ └── roi_heads
│ │ ├── __init__.py
│ │ ├── partA2_head.py
│ │ ├── pointrcnn_head.py
│ │ ├── pvrcnn_head.py
│ │ ├── roi_head_template.py
│ │ ├── second_head.py
│ │ ├── target_assigner
│ │ └── proposal_target_layer.py
│ │ └── voxelrcnn_head.py
├── ops
│ ├── iou3d_nms
│ │ ├── iou3d_nms_utils.py
│ │ └── src
│ │ │ ├── iou3d_cpu.cpp
│ │ │ ├── iou3d_cpu.h
│ │ │ ├── iou3d_nms.cpp
│ │ │ ├── iou3d_nms.h
│ │ │ ├── iou3d_nms_api.cpp
│ │ │ └── iou3d_nms_kernel.cu
│ ├── pointnet2
│ │ ├── pointnet2_batch
│ │ │ ├── pointnet2_modules.py
│ │ │ ├── pointnet2_utils.py
│ │ │ └── src
│ │ │ │ ├── ball_query.cpp
│ │ │ │ ├── ball_query_gpu.cu
│ │ │ │ ├── ball_query_gpu.h
│ │ │ │ ├── cuda_utils.h
│ │ │ │ ├── group_points.cpp
│ │ │ │ ├── group_points_gpu.cu
│ │ │ │ ├── group_points_gpu.h
│ │ │ │ ├── interpolate.cpp
│ │ │ │ ├── interpolate_gpu.cu
│ │ │ │ ├── interpolate_gpu.h
│ │ │ │ ├── pointnet2_api.cpp
│ │ │ │ ├── sampling.cpp
│ │ │ │ ├── sampling_gpu.cu
│ │ │ │ └── sampling_gpu.h
│ │ └── pointnet2_stack
│ │ │ ├── pointnet2_modules.py
│ │ │ ├── pointnet2_utils.py
│ │ │ ├── src
│ │ │ ├── ball_query.cpp
│ │ │ ├── ball_query_gpu.cu
│ │ │ ├── ball_query_gpu.h
│ │ │ ├── cuda_utils.h
│ │ │ ├── group_points.cpp
│ │ │ ├── group_points_gpu.cu
│ │ │ ├── group_points_gpu.h
│ │ │ ├── interpolate.cpp
│ │ │ ├── interpolate_gpu.cu
│ │ │ ├── interpolate_gpu.h
│ │ │ ├── pointnet2_api.cpp
│ │ │ ├── sampling.cpp
│ │ │ ├── sampling_gpu.cu
│ │ │ ├── sampling_gpu.h
│ │ │ ├── vector_pool.cpp
│ │ │ ├── vector_pool_gpu.cu
│ │ │ ├── vector_pool_gpu.h
│ │ │ ├── voxel_query.cpp
│ │ │ ├── voxel_query_gpu.cu
│ │ │ └── voxel_query_gpu.h
│ │ │ ├── voxel_pool_modules.py
│ │ │ └── voxel_query_utils.py
│ ├── roiaware_pool3d
│ │ ├── roiaware_pool3d_utils.py
│ │ └── src
│ │ │ ├── roiaware_pool3d.cpp
│ │ │ └── roiaware_pool3d_kernel.cu
│ └── roipoint_pool3d
│ │ ├── roipoint_pool3d_utils.py
│ │ └── src
│ │ ├── roipoint_pool3d.cpp
│ │ └── roipoint_pool3d_kernel.cu
└── utils
│ ├── __init__.py
│ ├── box_coder_utils.py
│ ├── box_utils.py
│ ├── calibration_kitti.py
│ ├── common_utils.py
│ ├── commu_utils.py
│ ├── kd_utils
│ ├── __init__.py
│ ├── kd_forwad.py
│ ├── kd_tgi_utils.py
│ ├── kd_utils.py
│ └── kd_vis_utils.py
│ ├── loss_utils.py
│ ├── object3d_kitti.py
│ ├── spconv_utils.py
│ └── transform_utils.py
├── requirements.txt
├── setup.py
└── tools
├── _init_path.py
├── cfgs
├── dataset_configs
│ ├── kitti_dataset.yaml
│ ├── nuscenes_dataset.yaml
│ └── waymo_dataset.yaml
├── kitti_models
│ ├── CaDDN.yaml
│ ├── PartA2.yaml
│ ├── PartA2_free.yaml
│ ├── pointpillar.yaml
│ ├── pointrcnn.yaml
│ ├── pointrcnn_iou.yaml
│ ├── pv_rcnn.yaml
│ ├── second.yaml
│ ├── second_iou.yaml
│ ├── second_multihead.yaml
│ └── voxel_rcnn_car.yaml
├── nuscenes_models
│ ├── cbgs_pp_multihead.yaml
│ └── cbgs_second_multihead.yaml
└── waymo_models
│ ├── PartA2.yaml
│ ├── PartA2_free.yaml
│ ├── centerpoint.yaml
│ ├── centerpoint_pillar_1x.yaml
│ ├── centerpoint_without_resnet.yaml
│ ├── cp-pillar
│ ├── cp-pillar-v0.4.yaml
│ ├── cp-pillar-v0.48.yaml
│ ├── cp-pillar-v0.48_sparsekd.yaml
│ ├── cp-pillar-v0.4_sparsekd.yaml
│ ├── cp-pillar-v0.64.yaml
│ ├── cp-pillar-v0.64_sparsekd.yaml
│ └── cp-pillar.yaml
│ ├── cp-pp.yaml
│ ├── cp-voxel
│ ├── cp-voxel-s.yaml
│ ├── cp-voxel-s_sparsekd.yaml
│ ├── cp-voxel-xs.yaml
│ ├── cp-voxel-xs_sparsekd.yaml
│ ├── cp-voxel-xxs.yaml
│ ├── cp-voxel-xxs_sparsekd.yaml
│ ├── cp-voxel.yaml
│ └── cp-voxel_sparsekd_crossstage.yaml
│ ├── pointpillar_1x.yaml
│ ├── pointrcnn.yaml
│ ├── pv_rcnn.yaml
│ ├── pv_rcnn_plusplus.yaml
│ ├── pv_rcnn_plusplus_resnet.yaml
│ ├── pv_rcnn_with_centerhead_rpn.yaml
│ ├── second.yaml
│ └── voxel_rcnn_with_centerhead_dyn_voxel.yaml
├── demo.py
├── eval_utils
└── eval_utils.py
├── extra_files
└── conv.py
├── scripts
├── dist_test.sh
├── dist_train.sh
├── slurm_test_mgpu.sh
├── slurm_test_single.sh
├── slurm_train.sh
├── slurm_train_single.sh
└── torch_train.sh
├── test.py
├── train.py
├── train_utils
├── optimization
│ ├── __init__.py
│ ├── fastai_optim.py
│ └── learning_schedules_fastai.py
├── train_kd_utils.py
└── train_utils.py
└── visual_utils
├── open3d_vis_utils.py
└── visualize_utils.py
/.gitignore:
--------------------------------------------------------------------------------
1 | **__pycache__**
2 | **build**
3 | **egg-info**
4 | **dist**
5 | data/
6 | *.pyc
7 | venv/
8 | *.idea/
9 | *.so
10 | *.sh
11 | *.pth
12 | *.pkl
13 | *.zip
14 | *.bin
15 | output
16 | version.py
17 | *.pdf
18 | *.ipynb
19 | *.png
--------------------------------------------------------------------------------
/SECURITY.md:
--------------------------------------------------------------------------------
1 | # Security Policy
2 |
3 | ## Supported Versions
4 |
5 | Use this section to tell people about which versions of your project are
6 | currently being supported with security updates.
7 |
8 | | Version | Supported |
9 | | ------- | ------------------ |
10 | | 5.1.x | :white_check_mark: |
11 | | 5.0.x | :x: |
12 | | 4.0.x | :white_check_mark: |
13 | | < 4.0 | :x: |
14 |
15 | ## Reporting a Vulnerability
16 |
17 | Use this section to tell people how to report a vulnerability.
18 |
19 | Tell them where to go, how often they can expect to get an update on a
20 | reported vulnerability, what to expect if the vulnerability is accepted or
21 | declined, etc.
22 |
--------------------------------------------------------------------------------
/data/waymo/ImageSets/infer_time.txt:
--------------------------------------------------------------------------------
1 | segment-10017090168044687777_6380_000_6400_000_with_camera_labels.tfrecord
--------------------------------------------------------------------------------
/docs/DEMO.md:
--------------------------------------------------------------------------------
1 | # Quick Demo
2 |
3 | Here we provide a quick demo to test a pretrained model on the custom point cloud data and visualize the predicted results.
4 |
5 | We suppose you already followed the [INSTALL.md](INSTALL.md) to install the `OpenPCDet` repo successfully.
6 |
7 | 1. Download the provided pretrained models as shown in the [README.md](../README.md).
8 |
9 | 2. Make sure you have already installed the [`Open3D`](https://github.com/isl-org/Open3D) (faster) or `mayavi` visualization tools.
10 | If not, you could install it as follows:
11 | ```
12 | pip install open3d
13 | # or
14 | pip install mayavi
15 | ```
16 |
17 | 3. Prepare your custom point cloud data (skip this step if you use the original KITTI data).
18 | * You need to transform the coordinate of your custom point cloud to
19 | the unified normative coordinate of `OpenPCDet`, that is, x-axis points towards to front direction,
20 | y-axis points towards to the left direction, and z-axis points towards to the top direction.
21 | * (Optional) the z-axis origin of your point cloud coordinate should be about 1.6m above the ground surface,
22 | since currently the provided models are trained on the KITTI dataset.
23 | * Set the intensity information, and save your transformed custom data to `numpy file`:
24 | ```python
25 | # Transform your point cloud data
26 | ...
27 |
28 | # Save it to the file.
29 | # The shape of points should be (num_points, 4), that is [x, y, z, intensity] (Only for KITTI dataset).
30 | # If you doesn't have the intensity information, just set them to zeros.
31 | # If you have the intensity information, you should normalize them to [0, 1].
32 | points[:, 3] = 0
33 | np.save(`my_data.npy`, points)
34 | ```
35 |
36 | 4. Run the demo with a pretrained model (e.g. PV-RCNN) and your custom point cloud data as follows:
37 | ```shell
38 | python demo.py --cfg_file cfgs/kitti_models/pv_rcnn.yaml \
39 | --ckpt pv_rcnn_8369.pth \
40 | --data_path ${POINT_CLOUD_DATA}
41 | ```
42 | Here `${POINT_CLOUD_DATA}` could be in any of the following format:
43 | * Your transformed custom data with a single numpy file like `my_data.npy`.
44 | * Your transformed custom data with a directory to test with multiple point cloud data.
45 | * The original KITTI `.bin` data within `data/kitti`, like `data/kitti/training/velodyne/000008.bin`.
46 |
47 | Then you could see the predicted results with visualized point cloud as follows:
48 |
49 |
50 | 
51 |
52 |
--------------------------------------------------------------------------------
/docs/GETTING_STARTED.md:
--------------------------------------------------------------------------------
1 | # Getting Started
2 | The dataset configs are located within [tools/cfgs/dataset_configs](../tools/cfgs/dataset_configs),
3 | and the model configs are located within [tools/cfgs](../tools/cfgs) for different datasets.
4 |
5 |
6 | ## Dataset Preparation
7 | Please follow the OpenPCDet [tutorial](https://github.com/open-mmlab/OpenPCDet/blob/master/docs/GETTING_STARTED.md) to
8 | prepare needed datasets.
9 |
10 | ## Training & Testing
11 | [//]: # ( TODO)
12 | ### Step 1: Train a teacher model (CP-Pillar as example)
13 | ```shell
14 | sh scripts/dist_train.sh ${NUM_GPUS} --cfg_file cfgs/waymo_models/cp-pillar/cp-pillar.yaml
15 | ```
16 |
17 | ### Step 2: Distillation (CP-Pillar-v0.4 as example)
18 | Modify following keys in the student distillation config
19 | ```yaml
20 | # cfgs/waymo_models/cp-pillar/cp-pillar-v0.4_sparsekd.yaml
21 | TEACHER_CKPT: ${PATH_TO_TEACHER_CKPT}
22 | PRETRAINED_MODEL: ${PATH_TO_TEACHER_CKPT}
23 | ```
24 | Run the training config
25 | ```shell
26 | sh scripts/dist_train.sh ${NUM_GPUS} --cfg_file cfgs/waymo_models/cp-pillar/cp-pillar-v0.4_sparsekd.yaml
27 | ```
28 |
29 | ## Calculate Efficiency Metrics
30 |
31 | ### Prepare
32 | Make sure you have installed our customized Thop as [INSTALL.md](./INSTALL.md).
33 | To calculate the Flops and Acts for spconv-based models, you also need to replace original `conv.py` in spconv
34 | with our modified one.
35 | ```shell
36 | # replace our modified conv file for
37 | # make sure your spconv is at least 2.1.20
38 | cp extra_files/conv.py ${CONDA_PATH}/envs/${ENV_NAME}/lib/${PYTHON_VERSION}/site-packages/spconv/pytorch/
39 | ```
40 |
41 | ### Command
42 | ```shell
43 | # Take Waymo as an example
44 | # This command have to be executed on single gpu only
45 | python test.py --cfg_file ${CONFIG_PATH} --batch_size 1 --ckpt ${CKPT_PATH} --infer_time --cal_params \
46 | --set DATA_CONFIG.DATA_SPLIT.test infer_time DATA_CONFIG.SAMPLED_INTERAVL.test 2
47 | ```
48 |
49 |
--------------------------------------------------------------------------------
/docs/INSTALL.md:
--------------------------------------------------------------------------------
1 | # Installation
2 |
3 | ### Requirements
4 | All the codes are tested in the following environment:
5 | * Linux (tested on Ubuntu 14.04/16.04/18.04/20.04/21.04)
6 | * Python 3.6+
7 | * PyTorch 1.1 or higher (tested on PyTorch 1.1, 1,3, 1,5~1.10)
8 | * CUDA 9.0 or higher (PyTorch 1.3+ needs CUDA 9.2+)
9 | * [`spconv v1.x`](https://github.com/traveller59/spconv/tree/v1.2.1) or [`spconv v2.x`](https://github.com/traveller59/spconv)
10 |
11 | Notice that all results in the paper is running in with Spconv 1.2.1, but the flops calculation for spconv-based network
12 | (i.e. CenterPoint-Voxel or SECOND) is running with Spconv 2.x.
13 |
14 | ### Install `pcdet v0.5.2`
15 | NOTE: Please re-install `pcdet v0.5.2` by running `python setup.py develop` even if you have already installed previous version.
16 |
17 | a. Clone this repository.
18 | ```shell
19 | git clone https://github.com/jihanyang/SparseKD.git
20 | ```
21 |
22 | b. Install the dependent libraries as follows:
23 |
24 | * Install the dependent python libraries:
25 |
26 | ```
27 | pip install -r requirements.txt
28 | ```
29 |
30 | * Install the SparseConv library, we use the implementation from [`[spconv]`](https://github.com/traveller59/spconv).
31 | * If you use PyTorch 1.1, then make sure you install the `spconv v1.0` with ([commit 8da6f96](https://github.com/traveller59/spconv/tree/8da6f967fb9a054d8870c3515b1b44eca2103634)) instead of the latest one.
32 | * If you use PyTorch 1.3+, then you need to install the `spconv v1.2.1`. As mentioned by the author of [`spconv`](https://github.com/traveller59/spconv/tree/v1.2.1), you need to use their docker if you use PyTorch 1.4+.
33 | * You could also install latest `spconv v2.x` with pip, see the official documents of [spconv](https://github.com/traveller59/spconv).
34 |
35 | c. Install this `pcdet` library and its dependent libraries by running the following command:
36 | ```shell
37 | python setup.py develop
38 | ```
39 |
40 | ### Install Thop
41 | (If you don't need to calculate **Flops**, **Parameters** and **Acts**, just skip this part)
42 | To calculate **Flops** and **Acts**, we leverage the popular [pytorch-OpCounter](https://github.com/Lyken17/pytorch-OpCounter).
43 |
44 | Note: that the **Flops** in our paper actually means **Macs** in [pytorch-OpCounter](https://github.com/Lyken17/pytorch-OpCounter).
45 | We use the term `Flops` in order to follow the same formulation as [RegNetX](https://github.com/facebookresearch/pycls), we name it as flops in
46 | the paper.
47 |
48 | a. Clone our customized Thop.
49 | ```shell
50 | git clone https://github.com/jihanyang/pytorch-OpCounter.git
51 | ```
52 |
53 | b. Install.
54 | ```shell
55 | cd pytorch-OpCounter && pip install .
56 | ```
57 |
--------------------------------------------------------------------------------
/docs/latency.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/CVMI-Lab/SparseKD/2e00edef6b94fb208ba64193940bc856e5effbcb/docs/latency.png
--------------------------------------------------------------------------------
/pcdet/__init__.py:
--------------------------------------------------------------------------------
1 | import subprocess
2 | from pathlib import Path
3 |
4 | from .version import __version__
5 |
6 | __all__ = [
7 | '__version__'
8 | ]
9 |
10 |
11 | def get_git_commit_number():
12 | if not (Path(__file__).parent / '../.git').exists():
13 | return '0000000'
14 |
15 | cmd_out = subprocess.run(['git', 'rev-parse', 'HEAD'], stdout=subprocess.PIPE)
16 | git_commit_number = cmd_out.stdout.decode('utf-8')[:7]
17 | return git_commit_number
18 |
19 |
20 | script_version = get_git_commit_number()
21 |
22 |
23 | if script_version not in __version__:
24 | __version__ = __version__ + '+py%s' % script_version
25 |
--------------------------------------------------------------------------------
/pcdet/config.py:
--------------------------------------------------------------------------------
1 | from pathlib import Path
2 |
3 | import yaml
4 | from easydict import EasyDict
5 |
6 |
7 | def log_config_to_file(cfg, pre='cfg', logger=None):
8 | for key, val in cfg.items():
9 | if isinstance(cfg[key], EasyDict):
10 | logger.info('\n%s.%s = edict()' % (pre, key))
11 | log_config_to_file(cfg[key], pre=pre + '.' + key, logger=logger)
12 | continue
13 | logger.info('%s.%s: %s' % (pre, key, val))
14 |
15 |
16 | def cfg_from_list(cfg_list, config):
17 | """Set config keys via list (e.g., from command line)."""
18 | from ast import literal_eval
19 | assert len(cfg_list) % 2 == 0
20 | for k, v in zip(cfg_list[0::2], cfg_list[1::2]):
21 | key_list = k.split('.')
22 | d = config
23 | for subkey in key_list[:-1]:
24 | assert subkey in d, 'NotFoundKey: %s' % subkey
25 | d = d[subkey]
26 | subkey = key_list[-1]
27 | assert subkey in d, 'NotFoundKey: %s' % subkey
28 | try:
29 | value = literal_eval(v)
30 | except:
31 | value = v
32 |
33 | if type(value) != type(d[subkey]) and isinstance(d[subkey], EasyDict):
34 | key_val_list = value.split(',')
35 | for src in key_val_list:
36 | cur_key, cur_val = src.split(':')
37 | val_type = type(d[subkey][cur_key])
38 | cur_val = val_type(cur_val)
39 | d[subkey][cur_key] = cur_val
40 | elif type(value) != type(d[subkey]) and isinstance(d[subkey], list):
41 | val_list = value.split(',')
42 | for k, x in enumerate(val_list):
43 | val_list[k] = type(d[subkey][0])(x)
44 | d[subkey] = val_list
45 | else:
46 | assert type(value) == type(d[subkey]), \
47 | 'type {} does not match original type {}'.format(type(value), type(d[subkey]))
48 | d[subkey] = value
49 |
50 |
51 | def merge_new_config(config, new_config):
52 | if '_BASE_CONFIG_' in new_config:
53 | with open(new_config['_BASE_CONFIG_'], 'r') as f:
54 | try:
55 | yaml_config = yaml.safe_load(f, Loader=yaml.FullLoader)
56 | except:
57 | yaml_config = yaml.safe_load(f)
58 | config.update(EasyDict(yaml_config))
59 |
60 | for key, val in new_config.items():
61 | if not isinstance(val, dict):
62 | config[key] = val
63 | continue
64 | if key not in config:
65 | config[key] = EasyDict()
66 | merge_new_config(config[key], val)
67 |
68 | return config
69 |
70 |
71 | def cfg_from_yaml_file(cfg_file, config):
72 | with open(cfg_file, 'r') as f:
73 | try:
74 | new_config = yaml.safe_load(f, Loader=yaml.FullLoader)
75 | except:
76 | new_config = yaml.safe_load(f)
77 |
78 | merge_new_config(config=config, new_config=new_config)
79 |
80 | return config
81 |
82 |
83 | cfg = EasyDict()
84 | cfg.ROOT_DIR = (Path(__file__).resolve().parent / '../').resolve()
85 | cfg.LOCAL_RANK = 0
86 |
--------------------------------------------------------------------------------
/pcdet/datasets/__init__.py:
--------------------------------------------------------------------------------
1 | from functools import partial
2 | import torch
3 | from torch.utils.data import DataLoader
4 | from torch.utils.data import DistributedSampler as _DistributedSampler
5 |
6 | from pcdet.utils import common_utils
7 |
8 | from .dataset import DatasetTemplate
9 | from .kitti.kitti_dataset import KittiDataset
10 | from .nuscenes.nuscenes_dataset import NuScenesDataset
11 | from .waymo.waymo_dataset import WaymoDataset
12 |
13 | __all__ = {
14 | 'DatasetTemplate': DatasetTemplate,
15 | 'KittiDataset': KittiDataset,
16 | 'NuScenesDataset': NuScenesDataset,
17 | 'WaymoDataset': WaymoDataset
18 | }
19 |
20 |
21 | class DistributedSampler(_DistributedSampler):
22 |
23 | def __init__(self, dataset, num_replicas=None, rank=None, shuffle=True):
24 | super().__init__(dataset, num_replicas=num_replicas, rank=rank)
25 | self.shuffle = shuffle
26 |
27 | def __iter__(self):
28 | if self.shuffle:
29 | g = torch.Generator()
30 | g.manual_seed(self.epoch)
31 | indices = torch.randperm(len(self.dataset), generator=g).tolist()
32 | else:
33 | indices = torch.arange(len(self.dataset)).tolist()
34 |
35 | indices += indices[:(self.total_size - len(indices))]
36 | assert len(indices) == self.total_size
37 |
38 | indices = indices[self.rank:self.total_size:self.num_replicas]
39 | assert len(indices) == self.num_samples
40 |
41 | return iter(indices)
42 |
43 |
44 | def build_dataloader(dataset_cfg, class_names, batch_size, dist, root_path=None, workers=4,
45 | logger=None, training=True, merge_all_iters_to_one_epoch=False, total_epochs=0):
46 |
47 | dataset = __all__[dataset_cfg.DATASET](
48 | dataset_cfg=dataset_cfg,
49 | class_names=class_names,
50 | root_path=root_path,
51 | training=training,
52 | logger=logger,
53 | )
54 |
55 | if merge_all_iters_to_one_epoch:
56 | assert hasattr(dataset, 'merge_all_iters_to_one_epoch')
57 | dataset.merge_all_iters_to_one_epoch(merge=True, epochs=total_epochs)
58 |
59 | if dist:
60 | if training:
61 | sampler = torch.utils.data.distributed.DistributedSampler(dataset)
62 | else:
63 | rank, world_size = common_utils.get_dist_info()
64 | sampler = DistributedSampler(dataset, world_size, rank, shuffle=False)
65 | else:
66 | sampler = None
67 | dataloader = DataLoader(
68 | dataset, batch_size=batch_size, pin_memory=True, num_workers=workers,
69 | shuffle=(sampler is None) and training, collate_fn=dataset.collate_batch,
70 | drop_last=False, sampler=sampler, timeout=0, worker_init_fn=partial(common_utils.worker_init_fn, seed=666)
71 | )
72 |
73 | return dataset, dataloader, sampler
74 |
--------------------------------------------------------------------------------
/pcdet/datasets/kitti/kitti_object_eval_python/LICENSE:
--------------------------------------------------------------------------------
1 | MIT License
2 |
3 | Copyright (c) 2018
4 |
5 | Permission is hereby granted, free of charge, to any person obtaining a copy
6 | of this software and associated documentation files (the "Software"), to deal
7 | in the Software without restriction, including without limitation the rights
8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9 | copies of the Software, and to permit persons to whom the Software is
10 | furnished to do so, subject to the following conditions:
11 |
12 | The above copyright notice and this permission notice shall be included in all
13 | copies or substantial portions of the Software.
14 |
15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21 | SOFTWARE.
22 |
--------------------------------------------------------------------------------
/pcdet/datasets/kitti/kitti_object_eval_python/README.md:
--------------------------------------------------------------------------------
1 | # kitti-object-eval-python
2 | **Note**: This is borrowed from [traveller59/kitti-object-eval-python](https://github.com/traveller59/kitti-object-eval-python)
3 |
4 | Fast kitti object detection eval in python(finish eval in less than 10 second), support 2d/bev/3d/aos. , support coco-style AP. If you use command line interface, numba need some time to compile jit functions.
5 | ## Dependencies
6 | Only support python 3.6+, need `numpy`, `skimage`, `numba`, `fire`. If you have Anaconda, just install `cudatoolkit` in anaconda. Otherwise, please reference to this [page](https://github.com/numba/numba#custom-python-environments) to set up llvm and cuda for numba.
7 | * Install by conda:
8 | ```
9 | conda install -c numba cudatoolkit=x.x (8.0, 9.0, 9.1, depend on your environment)
10 | ```
11 | ## Usage
12 | * commandline interface:
13 | ```
14 | python evaluate.py evaluate --label_path=/path/to/your_gt_label_folder --result_path=/path/to/your_result_folder --label_split_file=/path/to/val.txt --current_class=0 --coco=False
15 | ```
16 | * python interface:
17 | ```Python
18 | import kitti_common as kitti
19 | from eval import get_official_eval_result, get_coco_eval_result
20 | def _read_imageset_file(path):
21 | with open(path, 'r') as f:
22 | lines = f.readlines()
23 | return [int(line) for line in lines]
24 | det_path = "/path/to/your_result_folder"
25 | dt_annos = kitti.get_label_annos(det_path)
26 | gt_path = "/path/to/your_gt_label_folder"
27 | gt_split_file = "/path/to/val.txt" # from https://xiaozhichen.github.io/files/mv3d/imagesets.tar.gz
28 | val_image_ids = _read_imageset_file(gt_split_file)
29 | gt_annos = kitti.get_label_annos(gt_path, val_image_ids)
30 | print(get_official_eval_result(gt_annos, dt_annos, 0)) # 6s in my computer
31 | print(get_coco_eval_result(gt_annos, dt_annos, 0)) # 18s in my computer
32 | ```
33 |
--------------------------------------------------------------------------------
/pcdet/datasets/kitti/kitti_object_eval_python/evaluate.py:
--------------------------------------------------------------------------------
1 | import time
2 |
3 | import fire
4 |
5 | import .kitti_common as kitti
6 | from .eval import get_coco_eval_result, get_official_eval_result
7 |
8 |
9 | def _read_imageset_file(path):
10 | with open(path, 'r') as f:
11 | lines = f.readlines()
12 | return [int(line) for line in lines]
13 |
14 |
15 | def evaluate(label_path,
16 | result_path,
17 | label_split_file,
18 | current_class=0,
19 | coco=False,
20 | score_thresh=-1):
21 | dt_annos = kitti.get_label_annos(result_path)
22 | if score_thresh > 0:
23 | dt_annos = kitti.filter_annos_low_score(dt_annos, score_thresh)
24 | val_image_ids = _read_imageset_file(label_split_file)
25 | gt_annos = kitti.get_label_annos(label_path, val_image_ids)
26 | if coco:
27 | return get_coco_eval_result(gt_annos, dt_annos, current_class)
28 | else:
29 | return get_official_eval_result(gt_annos, dt_annos, current_class)
30 |
31 |
32 | if __name__ == '__main__':
33 | fire.Fire()
34 |
--------------------------------------------------------------------------------
/pcdet/datasets/kitti/kitti_utils.py:
--------------------------------------------------------------------------------
1 | import numpy as np
2 | from ...utils import box_utils
3 |
4 |
5 | def transform_annotations_to_kitti_format(annos, map_name_to_kitti=None, info_with_fakelidar=False):
6 | """
7 | Args:
8 | annos:
9 | map_name_to_kitti: dict, map name to KITTI names (Car, Pedestrian, Cyclist)
10 | info_with_fakelidar:
11 | Returns:
12 |
13 | """
14 | for anno in annos:
15 | for k in range(anno['name'].shape[0]):
16 | anno['name'][k] = map_name_to_kitti[anno['name'][k]]
17 |
18 | anno['bbox'] = np.zeros((len(anno['name']), 4))
19 | anno['bbox'][:, 2:4] = 50 # [0, 0, 50, 50]
20 | anno['truncated'] = np.zeros(len(anno['name']))
21 | anno['occluded'] = np.zeros(len(anno['name']))
22 | if 'boxes_lidar' in anno:
23 | gt_boxes_lidar = anno['boxes_lidar'].copy()
24 | else:
25 | gt_boxes_lidar = anno['gt_boxes_lidar'].copy()
26 |
27 | if len(gt_boxes_lidar) > 0:
28 | if info_with_fakelidar:
29 | gt_boxes_lidar = box_utils.boxes3d_kitti_fakelidar_to_lidar(gt_boxes_lidar)
30 |
31 | gt_boxes_lidar[:, 2] -= gt_boxes_lidar[:, 5] / 2
32 | anno['location'] = np.zeros((gt_boxes_lidar.shape[0], 3))
33 | anno['location'][:, 0] = -gt_boxes_lidar[:, 1] # x = -y_lidar
34 | anno['location'][:, 1] = -gt_boxes_lidar[:, 2] # y = -z_lidar
35 | anno['location'][:, 2] = gt_boxes_lidar[:, 0] # z = x_lidar
36 | dxdydz = gt_boxes_lidar[:, 3:6]
37 | anno['dimensions'] = dxdydz[:, [0, 2, 1]] # lwh ==> lhw
38 | anno['rotation_y'] = -gt_boxes_lidar[:, 6] - np.pi / 2.0
39 | anno['alpha'] = -np.arctan2(-gt_boxes_lidar[:, 1], gt_boxes_lidar[:, 0]) + anno['rotation_y']
40 | else:
41 | anno['location'] = anno['dimensions'] = np.zeros((0, 3))
42 | anno['rotation_y'] = anno['alpha'] = np.zeros(0)
43 |
44 | return annos
45 |
46 |
47 | def calib_to_matricies(calib):
48 | """
49 | Converts calibration object to transformation matricies
50 | Args:
51 | calib: calibration.Calibration, Calibration object
52 | Returns
53 | V2R: (4, 4), Lidar to rectified camera transformation matrix
54 | P2: (3, 4), Camera projection matrix
55 | """
56 | V2C = np.vstack((calib.V2C, np.array([0, 0, 0, 1], dtype=np.float32))) # (4, 4)
57 | R0 = np.hstack((calib.R0, np.zeros((3, 1), dtype=np.float32))) # (3, 4)
58 | R0 = np.vstack((R0, np.array([0, 0, 0, 1], dtype=np.float32))) # (4, 4)
59 | V2R = R0 @ V2C
60 | P2 = calib.P2
61 | return V2R, P2
--------------------------------------------------------------------------------
/pcdet/datasets/processor/point_feature_encoder.py:
--------------------------------------------------------------------------------
1 | import numpy as np
2 |
3 |
4 | class PointFeatureEncoder(object):
5 | def __init__(self, config, point_cloud_range=None):
6 | super().__init__()
7 | self.point_encoding_config = config
8 | assert list(self.point_encoding_config.src_feature_list[0:3]) == ['x', 'y', 'z']
9 | self.used_feature_list = self.point_encoding_config.used_feature_list
10 | self.src_feature_list = self.point_encoding_config.src_feature_list
11 | self.point_cloud_range = point_cloud_range
12 |
13 | @property
14 | def num_point_features(self):
15 | return getattr(self, self.point_encoding_config.encoding_type)(points=None)
16 |
17 | def forward(self, data_dict):
18 | """
19 | Args:
20 | data_dict:
21 | points: (N, 3 + C_in)
22 | ...
23 | Returns:
24 | data_dict:
25 | points: (N, 3 + C_out),
26 | use_lead_xyz: whether to use xyz as point-wise features
27 | ...
28 | """
29 | data_dict['points'], use_lead_xyz = getattr(self, self.point_encoding_config.encoding_type)(
30 | data_dict['points']
31 | )
32 | data_dict['use_lead_xyz'] = use_lead_xyz
33 |
34 | if self.point_encoding_config.get('filter_sweeps', False) and 'timestamp' in self.src_feature_list:
35 | max_sweeps = self.point_encoding_config.max_sweeps
36 | idx = self.src_feature_list.index('timestamp')
37 | dt = np.round(data_dict['points'][:, idx], 2)
38 | max_dt = sorted(np.unique(dt))[min(len(np.unique(dt))-1, max_sweeps-1)]
39 | data_dict['points'] = data_dict['points'][dt <= max_dt]
40 |
41 | return data_dict
42 |
43 | def absolute_coordinates_encoding(self, points=None):
44 | if points is None:
45 | num_output_features = len(self.used_feature_list)
46 | return num_output_features
47 |
48 | point_feature_list = [points[:, 0:3]]
49 | for x in self.used_feature_list:
50 | if x in ['x', 'y', 'z']:
51 | continue
52 | idx = self.src_feature_list.index(x)
53 | point_feature_list.append(points[:, idx:idx+1])
54 | point_features = np.concatenate(point_feature_list, axis=1)
55 | return point_features, True
56 |
--------------------------------------------------------------------------------
/pcdet/models/__init__.py:
--------------------------------------------------------------------------------
1 | from collections import namedtuple
2 |
3 | import numpy as np
4 | import torch
5 |
6 | from .detectors import build_detector
7 | from pcdet.utils import commu_utils
8 |
9 | try:
10 | import kornia
11 | except:
12 | pass
13 | # print('Warning: kornia is not installed. This package is only required by CaDDN')
14 |
15 |
16 | def build_teacher_network(cfg, args, train_set, dist, logger):
17 | teacher_model = build_network(model_cfg=cfg.MODEL_TEACHER, num_class=len(cfg.CLASS_NAMES), dataset=train_set)
18 | teacher_model.cuda()
19 |
20 | for param_k in teacher_model.parameters():
21 | param_k.requires_grad = False # not update by gradient
22 |
23 | teacher_model.train()
24 |
25 | if args.teacher_ckpt is not None:
26 | logger.info('Loading teacher parameters >>>>>>')
27 | teacher_model.load_params_from_file(filename=args.teacher_ckpt, to_cpu=dist, logger=logger)
28 |
29 | teacher_model.is_teacher = True
30 | for cur_module in teacher_model.module_list:
31 | cur_module.is_teacher = True
32 | cur_module.kd = True
33 |
34 | return teacher_model
35 |
36 |
37 | def build_network(model_cfg, num_class, dataset):
38 | model = build_detector(
39 | model_cfg=model_cfg, num_class=num_class, dataset=dataset
40 | )
41 | return model
42 |
43 |
44 | def load_data_to_gpu(batch_dict):
45 | for key, val in batch_dict.items():
46 | if not isinstance(val, np.ndarray):
47 | continue
48 | elif key in ['frame_id', 'metadata', 'calib']:
49 | continue
50 | elif key in ['images']:
51 | batch_dict[key] = kornia.image_to_tensor(val).float().cuda().contiguous()
52 | elif key in ['image_shape']:
53 | batch_dict[key] = torch.from_numpy(val).int().cuda()
54 | else:
55 | batch_dict[key] = torch.from_numpy(val).float().cuda()
56 |
57 |
58 | def model_fn_decorator():
59 | ModelReturn = namedtuple('ModelReturn', ['loss', 'tb_dict', 'disp_dict'])
60 |
61 | def model_func(model, batch_dict):
62 | load_data_to_gpu(batch_dict)
63 | ret_dict, tb_dict, disp_dict = model(batch_dict)
64 |
65 | loss = ret_dict['loss'].mean()
66 | if hasattr(model, 'update_global_step'):
67 | model.update_global_step()
68 | else:
69 | model.module.update_global_step()
70 |
71 | return ModelReturn(loss, tb_dict, disp_dict)
72 |
73 | return model_func
74 |
--------------------------------------------------------------------------------
/pcdet/models/backbones_2d/__init__.py:
--------------------------------------------------------------------------------
1 | from .base_bev_backbone import BaseBEVBackbone, BaseBEVBackboneSkipPost
2 |
3 | __all__ = {
4 | 'BaseBEVBackbone': BaseBEVBackbone,
5 | 'BaseBEVBackboneSkipPost': BaseBEVBackboneSkipPost
6 | }
7 |
--------------------------------------------------------------------------------
/pcdet/models/backbones_2d/map_to_bev/__init__.py:
--------------------------------------------------------------------------------
1 | from .height_compression import HeightCompression
2 | from .pointpillar_scatter import PointPillarScatter
3 | from .conv2d_collapse import Conv2DCollapse
4 | from .pillar_reencoding import PillarReencoding
5 |
6 | __all__ = {
7 | 'HeightCompression': HeightCompression,
8 | 'PointPillarScatter': PointPillarScatter,
9 | 'Conv2DCollapse': Conv2DCollapse,
10 | 'PillarReencoding': PillarReencoding
11 | }
12 |
--------------------------------------------------------------------------------
/pcdet/models/backbones_2d/map_to_bev/conv2d_collapse.py:
--------------------------------------------------------------------------------
1 | import torch
2 | import torch.nn as nn
3 |
4 | from pcdet.models.model_utils.basic_block_2d import BasicBlock2D
5 |
6 |
7 | class Conv2DCollapse(nn.Module):
8 |
9 | def __init__(self, model_cfg, grid_size):
10 | """
11 | Initializes 2D convolution collapse module
12 | Args:
13 | model_cfg: EasyDict, Model configuration
14 | grid_size: (X, Y, Z) Voxel grid size
15 | """
16 | super().__init__()
17 | self.model_cfg = model_cfg
18 | self.num_heights = grid_size[-1]
19 | self.num_bev_features = self.model_cfg.NUM_BEV_FEATURES
20 | self.block = BasicBlock2D(in_channels=self.num_bev_features * self.num_heights,
21 | out_channels=self.num_bev_features,
22 | **self.model_cfg.ARGS)
23 |
24 | def forward(self, batch_dict):
25 | """
26 | Collapses voxel features to BEV via concatenation and channel reduction
27 | Args:
28 | batch_dict:
29 | voxel_features: (B, C, Z, Y, X), Voxel feature representation
30 | Returns:
31 | batch_dict:
32 | spatial_features: (B, C, Y, X), BEV feature representation
33 | """
34 | voxel_features = batch_dict["voxel_features"]
35 | bev_features = voxel_features.flatten(start_dim=1, end_dim=2) # (B, C, Z, Y, X) -> (B, C*Z, Y, X)
36 | bev_features = self.block(bev_features) # (B, C*Z, Y, X) -> (B, C, Y, X)
37 | batch_dict["spatial_features"] = bev_features
38 | return batch_dict
39 |
--------------------------------------------------------------------------------
/pcdet/models/backbones_2d/map_to_bev/height_compression.py:
--------------------------------------------------------------------------------
1 | import torch.nn as nn
2 |
3 |
4 | class HeightCompression(nn.Module):
5 | def __init__(self, model_cfg, **kwargs):
6 | super().__init__()
7 | self.model_cfg = model_cfg
8 | self.num_bev_features = self.model_cfg.NUM_BEV_FEATURES
9 |
10 | def forward(self, batch_dict):
11 | """
12 | Args:
13 | batch_dict:
14 | encoded_spconv_tensor: sparse tensor
15 | Returns:
16 | batch_dict:
17 | spatial_features:
18 |
19 | """
20 | encoded_spconv_tensor = batch_dict['encoded_spconv_tensor']
21 | spatial_features = encoded_spconv_tensor.dense()
22 | N, C, D, H, W = spatial_features.shape
23 | spatial_features = spatial_features.view(N, C * D, H, W)
24 | batch_dict['spatial_features'] = spatial_features
25 | batch_dict['spatial_features_stride'] = batch_dict['encoded_spconv_tensor_stride']
26 | return batch_dict
27 |
--------------------------------------------------------------------------------
/pcdet/models/backbones_2d/map_to_bev/pillar_reencoding.py:
--------------------------------------------------------------------------------
1 | import torch
2 | import torch.nn as nn
3 |
4 | from .pointpillar_scatter import PointPillarScatter
5 | from pcdet.models.model_utils.efficientnet_utils import get_act_layer
6 | from pcdet.models.model_utils.basic_block_2d import build_block
7 |
8 |
9 | class PillarReencoding(nn.Module):
10 | def __init__(self, model_cfg, grid_size, **kwargs):
11 | super().__init__()
12 |
13 | self.model_cfg = model_cfg
14 | input_channels = model_cfg.IN_CHANNEL
15 | self.num_bev_features = self.model_cfg.NUM_BEV_FEATURES
16 | self.pillar_scatter = PointPillarScatter(model_cfg, grid_size, in_channel=input_channels, **kwargs)
17 |
18 | act_fn = get_act_layer(self.model_cfg.get('ACT_FN', 'ReLU'))
19 |
20 | if self.model_cfg.get('LAYER_NUMS', None) is not None:
21 | assert len(self.model_cfg.LAYER_NUMS) == len(self.model_cfg.LAYER_STRIDES) == len(self.model_cfg.NUM_FILTERS)
22 | layer_nums = self.model_cfg.LAYER_NUMS
23 | layer_strides = self.model_cfg.LAYER_STRIDES
24 | num_filters = self.model_cfg.NUM_FILTERS
25 | else:
26 | layer_nums = layer_strides = num_filters = []
27 |
28 | num_levels = len(layer_nums)
29 | c_in_list = [input_channels, *num_filters[:-1]]
30 | self.blocks = nn.ModuleList()
31 | for idx in range(num_levels):
32 | cur_layers = [
33 | nn.ZeroPad2d(1),
34 | nn.Conv2d(
35 | c_in_list[idx], num_filters[idx], kernel_size=3,
36 | stride=layer_strides[idx], padding=0, bias=False
37 | ),
38 | nn.BatchNorm2d(num_filters[idx], eps=1e-3, momentum=0.01),
39 | act_fn()
40 | ]
41 | for k in range(layer_nums[idx]):
42 | cur_layers.extend(
43 | build_block(
44 | self.model_cfg.get('CONV_BLOCK', 'BasicBlock2D'), num_filters[idx],
45 | num_filters[idx], act_fn=act_fn, kernel_size=3, padding=1, bias=False
46 | ),
47 | )
48 | self.blocks.append(nn.Sequential(*cur_layers))
49 |
50 | def forward(self, batch_dict, **kwargs):
51 | data_dict = self.pillar_scatter(batch_dict, **kwargs)
52 |
53 | spatial_features = data_dict['spatial_features']
54 | x = spatial_features
55 | for i in range(len(self.blocks)):
56 | x = self.blocks[i](x)
57 |
58 | data_dict['spatial_features'] = x
59 |
60 | return data_dict
61 |
--------------------------------------------------------------------------------
/pcdet/models/backbones_2d/map_to_bev/pointpillar_scatter.py:
--------------------------------------------------------------------------------
1 | import torch
2 | import torch.nn as nn
3 |
4 | from pcdet.models.model_utils.basic_block_2d import Focus
5 |
6 |
7 | class PointPillarScatter(nn.Module):
8 | def __init__(self, model_cfg, grid_size, **kwargs):
9 | super().__init__()
10 |
11 | self.model_cfg = model_cfg
12 | self.num_bev_features = self.model_cfg.NUM_BEV_FEATURES if 'in_channel' not in kwargs else kwargs['in_channel']
13 | self.nx, self.ny, self.nz = grid_size
14 | assert self.nz == 1
15 |
16 | if kwargs.get('focus', None):
17 | self.focus = Focus()
18 |
19 | def forward(self, batch_dict, **kwargs):
20 | if kwargs.get('pillar_feature_tea', None):
21 | pillar_features, coords = batch_dict['pillar_features_tea'], batch_dict['voxel_coords_tea']
22 | else:
23 | pillar_features, coords = batch_dict['pillar_features'], batch_dict['voxel_coords']
24 |
25 | batch_spatial_features = []
26 | batch_size = coords[:, 0].max().int().item() + 1
27 |
28 | num_bev_features = self.num_bev_features // 4 if getattr(self, 'focus', None) else self.num_bev_features
29 | for batch_idx in range(batch_size):
30 | spatial_feature = torch.zeros(
31 | num_bev_features,
32 | self.nz * self.nx * self.ny,
33 | dtype=pillar_features.dtype,
34 | device=pillar_features.device)
35 |
36 | batch_mask = coords[:, 0] == batch_idx
37 | this_coords = coords[batch_mask, :]
38 | indices = this_coords[:, 1] + this_coords[:, 2] * self.nx + this_coords[:, 3]
39 | indices = indices.type(torch.long)
40 | pillars = pillar_features[batch_mask, :]
41 | pillars = pillars.t()
42 | spatial_feature[:, indices] = pillars
43 | batch_spatial_features.append(spatial_feature)
44 |
45 | batch_spatial_features = torch.stack(batch_spatial_features, 0)
46 | batch_spatial_features = batch_spatial_features.view(batch_size, num_bev_features * self.nz, self.ny, self.nx)
47 |
48 | if getattr(self, 'focus', None):
49 | batch_spatial_features = self.focus(batch_spatial_features)
50 |
51 | if kwargs.get('out_feature_name', None):
52 | batch_dict[kwargs['out_feature_name']] = batch_spatial_features
53 | else:
54 | batch_dict['spatial_features'] = batch_spatial_features
55 | return batch_dict
56 |
57 |
58 | def point_pillar_scatter(num_bev_features, grid_size, pillar_features, coords):
59 | nx, ny, nz = grid_size
60 | batch_spatial_features = []
61 | batch_size = coords[:, 0].max().int().item() + 1
62 |
63 | for batch_idx in range(batch_size):
64 | spatial_feature = torch.zeros(
65 | num_bev_features,
66 | nz * nx * ny,
67 | dtype=pillar_features.dtype,
68 | device=pillar_features.device)
69 |
70 | batch_mask = coords[:, 0] == batch_idx
71 | this_coords = coords[batch_mask, :]
72 | indices = this_coords[:, 1] + this_coords[:, 2] * nx + this_coords[:, 3]
73 | indices = indices.type(torch.long)
74 | pillars = pillar_features[batch_mask, :]
75 | pillars = pillars.t()
76 | spatial_feature[:, indices] = pillars
77 | batch_spatial_features.append(spatial_feature)
78 |
79 | batch_spatial_features = torch.stack(batch_spatial_features, 0)
80 | batch_spatial_features = batch_spatial_features.view(batch_size, num_bev_features * nz, ny, nx)
81 |
82 | return batch_spatial_features
83 |
--------------------------------------------------------------------------------
/pcdet/models/backbones_3d/__init__.py:
--------------------------------------------------------------------------------
1 | from .pointnet2_backbone import PointNet2Backbone, PointNet2MSG
2 | from .spconv_backbone import VoxelBackBone8x, VoxelResBackBone8x
3 | from .spconv_unet import UNetV2
4 |
5 |
6 | __all__ = {
7 | 'VoxelBackBone8x': VoxelBackBone8x,
8 | 'UNetV2': UNetV2,
9 | 'PointNet2Backbone': PointNet2Backbone,
10 | 'PointNet2MSG': PointNet2MSG,
11 | 'VoxelResBackBone8x': VoxelResBackBone8x,
12 | }
13 |
--------------------------------------------------------------------------------
/pcdet/models/backbones_3d/pfe/__init__.py:
--------------------------------------------------------------------------------
1 | from .voxel_set_abstraction import VoxelSetAbstraction
2 |
3 | __all__ = {
4 | 'VoxelSetAbstraction': VoxelSetAbstraction
5 | }
6 |
--------------------------------------------------------------------------------
/pcdet/models/backbones_3d/pillar_adaptor/__init__.py:
--------------------------------------------------------------------------------
1 | from .sparse_interp_adaptor import SparseInterpAdaptor
2 | from .sparse_naive_adaptor import SparseNaiveAdaptor
3 |
4 | __all__ = {
5 | 'SparseInterpAdaptor': SparseInterpAdaptor,
6 | 'SparseNaiveAdaptor': SparseNaiveAdaptor
7 | }
8 |
--------------------------------------------------------------------------------
/pcdet/models/backbones_3d/pillar_adaptor/pillar_adaptor_template.py:
--------------------------------------------------------------------------------
1 | import torch.nn as nn
2 |
3 |
4 | class PillarAdaptorTemplate(nn.Module):
5 | def __init__(self, model_cfg, in_channel, point_cloud_range):
6 | super(PillarAdaptorTemplate, self).__init__()
7 | self.model_cfg = model_cfg
8 | self.kd_only = self.model_cfg.KD_ONLY
9 | self.cal_loss = self.model_cfg.CAL_LOSS
10 |
11 | self.position = self.model_cfg.POSITION
12 | self.groups = self.model_cfg.CONV.get('GROUPS', None)
13 | self.point_could_range = point_cloud_range
14 |
15 | def init_weights(self):
16 | for m in self.modules():
17 | if isinstance(m, nn.Conv2d) or isinstance(m, nn.Conv1d):
18 | nn.init.kaiming_normal_(m.weight)
19 | if m.bias is not None:
20 | nn.init.constant_(m.bias, 0)
21 | if isinstance(m, nn.BatchNorm2d) or isinstance(m, nn.BatchNorm1d):
22 | nn.init.constant_(m.weight, 1.0)
23 | nn.init.constant_(m.bias, 0)
24 |
25 | def build_loss(self):
26 | self.add_module('loss_func', getattr(nn, self.model_cfg.LOSS_CONFIG.NAME)())
27 |
28 | def group_teacher_voxel_coord_by_z(self, voxel_coord):
29 | """
30 | Args:
31 | voxel_coord: [N, 3] (z, y, x)
32 |
33 | Returns:
34 |
35 | """
36 | assert voxel_coord[:, 0].max() <= (self.groups - 1)
37 |
38 | group_mask_list = []
39 | for idx in range(self.groups):
40 | mask = voxel_coord[:, 0] == idx
41 | group_mask_list.append(mask)
42 |
43 | return group_mask_list
44 |
--------------------------------------------------------------------------------
/pcdet/models/backbones_3d/vfe/__init__.py:
--------------------------------------------------------------------------------
1 | from .mean_vfe import MeanVFE
2 | from .pillar_vfe import PillarVFE
3 | from .dynamic_mean_vfe import DynamicMeanVFE
4 | from .dynamic_pillar_vfe import DynamicPillarVFE
5 | from .dynamic_kp_vfe import DynamicKPVFE
6 | from .image_vfe import ImageVFE
7 | from .vfe_template import VFETemplate
8 | from .dynamic_pillar_vfe import DynamicPillarVFETea
9 |
10 |
11 | __all__ = {
12 | 'VFETemplate': VFETemplate,
13 | 'MeanVFE': MeanVFE,
14 | 'PillarVFE': PillarVFE,
15 | 'ImageVFE': ImageVFE,
16 | 'DynMeanVFE': DynamicMeanVFE,
17 | 'DynPillarVFE': DynamicPillarVFE,
18 | 'DynKPVFE': DynamicKPVFE,
19 | 'DynPillarVFETea': DynamicPillarVFETea,
20 | }
21 |
--------------------------------------------------------------------------------
/pcdet/models/backbones_3d/vfe/dynamic_mean_vfe.py:
--------------------------------------------------------------------------------
1 | import torch
2 |
3 | from .vfe_template import VFETemplate
4 |
5 | try:
6 | import torch_scatter
7 | except Exception as e:
8 | # Incase someone doesn't want to use dynamic pillar vfe and hasn't installed torch_scatter
9 | pass
10 |
11 | from .vfe_template import VFETemplate
12 |
13 |
14 | class DynamicMeanVFE(VFETemplate):
15 | def __init__(self, model_cfg, num_point_features, voxel_size, grid_size, point_cloud_range, **kwargs):
16 | super().__init__(model_cfg=model_cfg)
17 | self.num_point_features = num_point_features
18 |
19 | self.grid_size = torch.tensor(grid_size).cuda().int()
20 | self.voxel_size = torch.tensor(voxel_size).cuda()
21 | self.point_cloud_range = torch.tensor(point_cloud_range).cuda()
22 |
23 | self.voxel_x = voxel_size[0]
24 | self.voxel_y = voxel_size[1]
25 | self.voxel_z = voxel_size[2]
26 | self.x_offset = self.voxel_x / 2 + point_cloud_range[0]
27 | self.y_offset = self.voxel_y / 2 + point_cloud_range[1]
28 | self.z_offset = self.voxel_z / 2 + point_cloud_range[2]
29 |
30 | self.scale_xyz = grid_size[0] * grid_size[1] * grid_size[2]
31 | self.scale_yz = grid_size[1] * grid_size[2]
32 | self.scale_z = grid_size[2]
33 |
34 | def get_output_feature_dim(self):
35 | return self.num_point_features
36 |
37 | @torch.no_grad()
38 | def forward(self, batch_dict, **kwargs):
39 | """
40 | Args:
41 | batch_dict:
42 | voxels: (num_voxels, max_points_per_voxel, C)
43 | voxel_num_points: optional (num_voxels)
44 | **kwargs:
45 |
46 | Returns:
47 | vfe_features: (num_voxels, C)
48 | """
49 | batch_size = batch_dict['batch_size']
50 | points = batch_dict['points'] # (batch_idx, x, y, z, i, e)
51 |
52 | # # debug
53 | point_coords = torch.floor((points[:, 1:4] - self.point_cloud_range[0:3]) / self.voxel_size).int()
54 | mask = ((point_coords >= 0) & (point_coords < self.grid_size)).all(dim=1)
55 | points = points[mask]
56 | point_coords = point_coords[mask]
57 | merge_coords = points[:, 0].int() * self.scale_xyz + \
58 | point_coords[:, 0] * self.scale_yz + \
59 | point_coords[:, 1] * self.scale_z + \
60 | point_coords[:, 2]
61 | points_data = points[:, 1:].contiguous()
62 |
63 | unq_coords, unq_inv, unq_cnt = torch.unique(merge_coords, return_inverse=True, return_counts=True)
64 |
65 | points_mean = torch_scatter.scatter_mean(points_data, unq_inv, dim=0)
66 |
67 | unq_coords = unq_coords.int()
68 | voxel_coords = torch.stack((unq_coords // self.scale_xyz,
69 | (unq_coords % self.scale_xyz) // self.scale_yz,
70 | (unq_coords % self.scale_yz) // self.scale_z,
71 | unq_coords % self.scale_z), dim=1)
72 | voxel_coords = voxel_coords[:, [0, 3, 2, 1]]
73 |
74 | batch_dict['voxel_features'] = points_mean.contiguous()
75 | batch_dict['voxel_coords'] = voxel_coords.contiguous()
76 | return batch_dict
77 |
--------------------------------------------------------------------------------
/pcdet/models/backbones_3d/vfe/image_vfe.py:
--------------------------------------------------------------------------------
1 | import torch
2 |
3 | from .vfe_template import VFETemplate
4 | from .image_vfe_modules import ffn, f2v
5 |
6 |
7 | class ImageVFE(VFETemplate):
8 | def __init__(self, model_cfg, grid_size, point_cloud_range, depth_downsample_factor, **kwargs):
9 | super().__init__(model_cfg=model_cfg)
10 | self.grid_size = grid_size
11 | self.pc_range = point_cloud_range
12 | self.downsample_factor = depth_downsample_factor
13 | self.module_topology = [
14 | 'ffn', 'f2v'
15 | ]
16 | self.build_modules()
17 |
18 | def build_modules(self):
19 | """
20 | Builds modules
21 | """
22 | for module_name in self.module_topology:
23 | module = getattr(self, 'build_%s' % module_name)()
24 | self.add_module(module_name, module)
25 |
26 | def build_ffn(self):
27 | """
28 | Builds frustum feature network
29 | Returns:
30 | ffn_module: nn.Module, Frustum feature network
31 | """
32 | ffn_module = ffn.__all__[self.model_cfg.FFN.NAME](
33 | model_cfg=self.model_cfg.FFN,
34 | downsample_factor=self.downsample_factor
35 | )
36 | self.disc_cfg = ffn_module.disc_cfg
37 | return ffn_module
38 |
39 | def build_f2v(self):
40 | """
41 | Builds frustum to voxel transformation
42 | Returns:
43 | f2v_module: nn.Module, Frustum to voxel transformation
44 | """
45 | f2v_module = f2v.__all__[self.model_cfg.F2V.NAME](
46 | model_cfg=self.model_cfg.F2V,
47 | grid_size=self.grid_size,
48 | pc_range=self.pc_range,
49 | disc_cfg=self.disc_cfg
50 | )
51 | return f2v_module
52 |
53 | def get_output_feature_dim(self):
54 | """
55 | Gets number of output channels
56 | Returns:
57 | out_feature_dim: int, Number of output channels
58 | """
59 | out_feature_dim = self.ffn.get_output_feature_dim()
60 | return out_feature_dim
61 |
62 | def forward(self, batch_dict, **kwargs):
63 | """
64 | Args:
65 | batch_dict:
66 | images: (N, 3, H_in, W_in), Input images
67 | **kwargs:
68 | Returns:
69 | batch_dict:
70 | voxel_features: (B, C, Z, Y, X), Image voxel features
71 | """
72 | batch_dict = self.ffn(batch_dict)
73 | batch_dict = self.f2v(batch_dict)
74 | return batch_dict
75 |
76 | def get_loss(self):
77 | """
78 | Gets DDN loss
79 | Returns:
80 | loss: (1), Depth distribution network loss
81 | tb_dict: dict[float], All losses to log in tensorboard
82 | """
83 |
84 | loss, tb_dict = self.ffn.get_loss()
85 | return loss, tb_dict
86 |
--------------------------------------------------------------------------------
/pcdet/models/backbones_3d/vfe/image_vfe_modules/f2v/__init__.py:
--------------------------------------------------------------------------------
1 | from .frustum_to_voxel import FrustumToVoxel
2 |
3 | __all__ = {
4 | 'FrustumToVoxel': FrustumToVoxel
5 | }
6 |
--------------------------------------------------------------------------------
/pcdet/models/backbones_3d/vfe/image_vfe_modules/f2v/frustum_to_voxel.py:
--------------------------------------------------------------------------------
1 | import torch
2 | import torch.nn as nn
3 |
4 | from .frustum_grid_generator import FrustumGridGenerator
5 | from .sampler import Sampler
6 |
7 |
8 | class FrustumToVoxel(nn.Module):
9 |
10 | def __init__(self, model_cfg, grid_size, pc_range, disc_cfg):
11 | """
12 | Initializes module to transform frustum features to voxel features via 3D transformation and sampling
13 | Args:
14 | model_cfg: EasyDict, Module configuration
15 | grid_size: [X, Y, Z], Voxel grid size
16 | pc_range: [x_min, y_min, z_min, x_max, y_max, z_max], Voxelization point cloud range (m)
17 | disc_cfg: EasyDict, Depth discretiziation configuration
18 | """
19 | super().__init__()
20 | self.model_cfg = model_cfg
21 | self.grid_size = grid_size
22 | self.pc_range = pc_range
23 | self.disc_cfg = disc_cfg
24 | self.grid_generator = FrustumGridGenerator(grid_size=grid_size,
25 | pc_range=pc_range,
26 | disc_cfg=disc_cfg)
27 | self.sampler = Sampler(**model_cfg.SAMPLER)
28 |
29 | def forward(self, batch_dict):
30 | """
31 | Generates voxel features via 3D transformation and sampling
32 | Args:
33 | batch_dict:
34 | frustum_features: (B, C, D, H_image, W_image), Image frustum features
35 | lidar_to_cam: (B, 4, 4), LiDAR to camera frame transformation
36 | cam_to_img: (B, 3, 4), Camera projection matrix
37 | image_shape: (B, 2), Image shape [H, W]
38 | Returns:
39 | batch_dict:
40 | voxel_features: (B, C, Z, Y, X), Image voxel features
41 | """
42 | # Generate sampling grid for frustum volume
43 | grid = self.grid_generator(lidar_to_cam=batch_dict["trans_lidar_to_cam"],
44 | cam_to_img=batch_dict["trans_cam_to_img"],
45 | image_shape=batch_dict["image_shape"]) # (B, X, Y, Z, 3)
46 |
47 | # Sample frustum volume to generate voxel volume
48 | voxel_features = self.sampler(input_features=batch_dict["frustum_features"],
49 | grid=grid) # (B, C, X, Y, Z)
50 |
51 | # (B, C, X, Y, Z) -> (B, C, Z, Y, X)
52 | voxel_features = voxel_features.permute(0, 1, 4, 3, 2)
53 | batch_dict["voxel_features"] = voxel_features
54 | return batch_dict
55 |
--------------------------------------------------------------------------------
/pcdet/models/backbones_3d/vfe/image_vfe_modules/f2v/sampler.py:
--------------------------------------------------------------------------------
1 | import torch
2 | import torch.nn as nn
3 | import torch.nn.functional as F
4 |
5 |
6 | class Sampler(nn.Module):
7 |
8 | def __init__(self, mode="bilinear", padding_mode="zeros"):
9 | """
10 | Initializes module
11 | Args:
12 | mode: string, Sampling mode [bilinear/nearest]
13 | padding_mode: string, Padding mode for outside grid values [zeros/border/reflection]
14 | """
15 | super().__init__()
16 | self.mode = mode
17 | self.padding_mode = padding_mode
18 |
19 | def forward(self, input_features, grid):
20 | """
21 | Samples input using sampling grid
22 | Args:
23 | input_features: (B, C, D, H, W), Input frustum features
24 | grid: (B, X, Y, Z, 3), Sampling grids for input features
25 | Returns
26 | output_features: (B, C, X, Y, Z) Output voxel features
27 | """
28 | # Sample from grid
29 | output = F.grid_sample(input=input_features, grid=grid, mode=self.mode, padding_mode=self.padding_mode)
30 | return output
31 |
--------------------------------------------------------------------------------
/pcdet/models/backbones_3d/vfe/image_vfe_modules/ffn/__init__.py:
--------------------------------------------------------------------------------
1 | from .depth_ffn import DepthFFN
2 |
3 | __all__ = {
4 | 'DepthFFN': DepthFFN
5 | }
6 |
--------------------------------------------------------------------------------
/pcdet/models/backbones_3d/vfe/image_vfe_modules/ffn/ddn/__init__.py:
--------------------------------------------------------------------------------
1 | from .ddn_deeplabv3 import DDNDeepLabV3
2 |
3 | __all__ = {
4 | 'DDNDeepLabV3': DDNDeepLabV3
5 | }
6 |
--------------------------------------------------------------------------------
/pcdet/models/backbones_3d/vfe/image_vfe_modules/ffn/ddn/ddn_deeplabv3.py:
--------------------------------------------------------------------------------
1 | from .ddn_template import DDNTemplate
2 |
3 | try:
4 | import torchvision
5 | except:
6 | pass
7 |
8 |
9 | class DDNDeepLabV3(DDNTemplate):
10 |
11 | def __init__(self, backbone_name, **kwargs):
12 | """
13 | Initializes DDNDeepLabV3 model
14 | Args:
15 | backbone_name: string, ResNet Backbone Name [ResNet50/ResNet101]
16 | """
17 | if backbone_name == "ResNet50":
18 | constructor = torchvision.models.segmentation.deeplabv3_resnet50
19 | elif backbone_name == "ResNet101":
20 | constructor = torchvision.models.segmentation.deeplabv3_resnet101
21 | else:
22 | raise NotImplementedError
23 |
24 | super().__init__(constructor=constructor, **kwargs)
25 |
--------------------------------------------------------------------------------
/pcdet/models/backbones_3d/vfe/image_vfe_modules/ffn/ddn_loss/__init__.py:
--------------------------------------------------------------------------------
1 | from .ddn_loss import DDNLoss
2 |
3 | __all__ = {
4 | "DDNLoss": DDNLoss
5 | }
6 |
--------------------------------------------------------------------------------
/pcdet/models/backbones_3d/vfe/image_vfe_modules/ffn/ddn_loss/balancer.py:
--------------------------------------------------------------------------------
1 | import torch
2 | import torch.nn as nn
3 |
4 | from pcdet.utils import loss_utils
5 |
6 |
7 | class Balancer(nn.Module):
8 | def __init__(self, fg_weight, bg_weight, downsample_factor=1):
9 | """
10 | Initialize fixed foreground/background loss balancer
11 | Args:
12 | fg_weight: float, Foreground loss weight
13 | bg_weight: float, Background loss weight
14 | downsample_factor: int, Depth map downsample factor
15 | """
16 | super().__init__()
17 | self.fg_weight = fg_weight
18 | self.bg_weight = bg_weight
19 | self.downsample_factor = downsample_factor
20 |
21 | def forward(self, loss, gt_boxes2d):
22 | """
23 | Forward pass
24 | Args:
25 | loss: (B, H, W), Pixel-wise loss
26 | gt_boxes2d: (B, N, 4), 2D box labels for foreground/background balancing
27 | Returns:
28 | loss: (1), Total loss after foreground/background balancing
29 | tb_dict: dict[float], All losses to log in tensorboard
30 | """
31 | # Compute masks
32 | fg_mask = loss_utils.compute_fg_mask(gt_boxes2d=gt_boxes2d,
33 | shape=loss.shape,
34 | downsample_factor=self.downsample_factor,
35 | device=loss.device)
36 | bg_mask = ~fg_mask
37 |
38 | # Compute balancing weights
39 | weights = self.fg_weight * fg_mask + self.bg_weight * bg_mask
40 | num_pixels = fg_mask.sum() + bg_mask.sum()
41 |
42 | # Compute losses
43 | loss *= weights
44 | fg_loss = loss[fg_mask].sum() / num_pixels
45 | bg_loss = loss[bg_mask].sum() / num_pixels
46 |
47 | # Get total loss
48 | loss = fg_loss + bg_loss
49 | tb_dict = {"balancer_loss": loss.item(), "fg_loss": fg_loss.item(), "bg_loss": bg_loss.item()}
50 | return loss, tb_dict
51 |
--------------------------------------------------------------------------------
/pcdet/models/backbones_3d/vfe/image_vfe_modules/ffn/ddn_loss/ddn_loss.py:
--------------------------------------------------------------------------------
1 | import torch
2 | import torch.nn as nn
3 |
4 |
5 | from .balancer import Balancer
6 | from pcdet.utils import transform_utils
7 |
8 | try:
9 | from kornia.losses.focal import FocalLoss
10 | except:
11 | pass
12 | # print('Warning: kornia is not installed. This package is only required by CaDDN')
13 |
14 |
15 | class DDNLoss(nn.Module):
16 |
17 | def __init__(self,
18 | weight,
19 | alpha,
20 | gamma,
21 | disc_cfg,
22 | fg_weight,
23 | bg_weight,
24 | downsample_factor):
25 | """
26 | Initializes DDNLoss module
27 | Args:
28 | weight: float, Loss function weight
29 | alpha: float, Alpha value for Focal Loss
30 | gamma: float, Gamma value for Focal Loss
31 | disc_cfg: dict, Depth discretiziation configuration
32 | fg_weight: float, Foreground loss weight
33 | bg_weight: float, Background loss weight
34 | downsample_factor: int, Depth map downsample factor
35 | """
36 | super().__init__()
37 | self.device = torch.cuda.current_device()
38 | self.disc_cfg = disc_cfg
39 | self.balancer = Balancer(downsample_factor=downsample_factor,
40 | fg_weight=fg_weight,
41 | bg_weight=bg_weight)
42 |
43 | # Set loss function
44 | self.alpha = alpha
45 | self.gamma = gamma
46 | self.loss_func = FocalLoss(alpha=self.alpha, gamma=self.gamma, reduction="none")
47 | self.weight = weight
48 |
49 | def forward(self, depth_logits, depth_maps, gt_boxes2d):
50 | """
51 | Gets DDN loss
52 | Args:
53 | depth_logits: (B, D+1, H, W), Predicted depth logits
54 | depth_maps: (B, H, W), Depth map [m]
55 | gt_boxes2d: torch.Tensor (B, N, 4), 2D box labels for foreground/background balancing
56 | Returns:
57 | loss: (1), Depth distribution network loss
58 | tb_dict: dict[float], All losses to log in tensorboard
59 | """
60 | tb_dict = {}
61 |
62 | # Bin depth map to create target
63 | depth_target = transform_utils.bin_depths(depth_maps, **self.disc_cfg, target=True)
64 |
65 | # Compute loss
66 | loss = self.loss_func(depth_logits, depth_target)
67 |
68 | # Compute foreground/background balancing
69 | loss, tb_dict = self.balancer(loss=loss, gt_boxes2d=gt_boxes2d)
70 |
71 | # Final loss
72 | loss *= self.weight
73 | tb_dict.update({"ddn_loss": loss.item()})
74 |
75 | return loss, tb_dict
76 |
--------------------------------------------------------------------------------
/pcdet/models/backbones_3d/vfe/mean_vfe.py:
--------------------------------------------------------------------------------
1 | import torch
2 |
3 | from .vfe_template import VFETemplate
4 |
5 |
6 | class MeanVFE(VFETemplate):
7 | def __init__(self, model_cfg, num_point_features, **kwargs):
8 | super().__init__(model_cfg=model_cfg)
9 | self.num_point_features = num_point_features
10 |
11 | def get_output_feature_dim(self):
12 | return self.num_point_features
13 |
14 | def forward(self, batch_dict, **kwargs):
15 | """
16 | Args:
17 | batch_dict:
18 | voxels: (num_voxels, max_points_per_voxel, C)
19 | voxel_num_points: optional (num_voxels)
20 | **kwargs:
21 |
22 | Returns:
23 | vfe_features: (num_voxels, C)
24 | """
25 | voxel_features, voxel_num_points = batch_dict['voxels'], batch_dict['voxel_num_points']
26 | points_mean = voxel_features[:, :, :].sum(dim=1, keepdim=False)
27 | normalizer = torch.clamp_min(voxel_num_points.view(-1, 1), min=1.0).type_as(voxel_features)
28 | points_mean = points_mean / normalizer
29 | batch_dict['voxel_features'] = points_mean.contiguous()
30 |
31 | return batch_dict
32 |
--------------------------------------------------------------------------------
/pcdet/models/dense_heads/__init__.py:
--------------------------------------------------------------------------------
1 | from .anchor_head_multi import AnchorHeadMulti
2 | from .anchor_head_single import AnchorHeadSingle
3 | from .anchor_head_template import AnchorHeadTemplate
4 | from .point_head_box import PointHeadBox
5 | from .point_head_simple import PointHeadSimple
6 | from .point_intra_part_head import PointIntraPartOffsetHead
7 | from .center_head import CenterHead
8 | from .anchor_head_separate import AnchorHeadSeparate
9 |
10 | __all__ = {
11 | 'AnchorHeadTemplate': AnchorHeadTemplate,
12 | 'AnchorHeadSingle': AnchorHeadSingle,
13 | 'PointIntraPartOffsetHead': PointIntraPartOffsetHead,
14 | 'PointHeadSimple': PointHeadSimple,
15 | 'PointHeadBox': PointHeadBox,
16 | 'AnchorHeadMulti': AnchorHeadMulti,
17 | 'CenterHead': CenterHead,
18 | 'AnchorHeadSeparate': AnchorHeadSeparate
19 | }
20 |
--------------------------------------------------------------------------------
/pcdet/models/dense_heads/point_head_simple.py:
--------------------------------------------------------------------------------
1 | import torch
2 |
3 | from ...utils import box_utils
4 | from .point_head_template import PointHeadTemplate
5 |
6 |
7 | class PointHeadSimple(PointHeadTemplate):
8 | """
9 | A simple point-based segmentation head, which are used for PV-RCNN keypoint segmentaion.
10 | Reference Paper: https://arxiv.org/abs/1912.13192
11 | PV-RCNN: Point-Voxel Feature Set Abstraction for 3D Object Detection
12 | """
13 | def __init__(self, num_class, input_channels, model_cfg, **kwargs):
14 | super().__init__(model_cfg=model_cfg, num_class=num_class)
15 | self.cls_layers = self.make_fc_layers(
16 | fc_cfg=self.model_cfg.CLS_FC,
17 | input_channels=input_channels,
18 | output_channels=num_class
19 | )
20 |
21 | def assign_targets(self, input_dict):
22 | """
23 | Args:
24 | input_dict:
25 | point_features: (N1 + N2 + N3 + ..., C)
26 | batch_size:
27 | point_coords: (N1 + N2 + N3 + ..., 4) [bs_idx, x, y, z]
28 | gt_boxes (optional): (B, M, 8)
29 | Returns:
30 | point_cls_labels: (N1 + N2 + N3 + ...), long type, 0:background, -1:ignored
31 | point_part_labels: (N1 + N2 + N3 + ..., 3)
32 | """
33 | point_coords = input_dict['point_coords']
34 | gt_boxes = input_dict['gt_boxes']
35 | assert gt_boxes.shape.__len__() == 3, 'gt_boxes.shape=%s' % str(gt_boxes.shape)
36 | assert point_coords.shape.__len__() in [2], 'points.shape=%s' % str(point_coords.shape)
37 |
38 | batch_size = gt_boxes.shape[0]
39 | extend_gt_boxes = box_utils.enlarge_box3d(
40 | gt_boxes.view(-1, gt_boxes.shape[-1]), extra_width=self.model_cfg.TARGET_CONFIG.GT_EXTRA_WIDTH
41 | ).view(batch_size, -1, gt_boxes.shape[-1])
42 | targets_dict = self.assign_stack_targets(
43 | points=point_coords, gt_boxes=gt_boxes, extend_gt_boxes=extend_gt_boxes,
44 | set_ignore_flag=True, use_ball_constraint=False,
45 | ret_part_labels=False
46 | )
47 |
48 | return targets_dict
49 |
50 | def get_loss(self, tb_dict=None):
51 | tb_dict = {} if tb_dict is None else tb_dict
52 | point_loss_cls, tb_dict_1 = self.get_cls_layer_loss()
53 |
54 | point_loss = point_loss_cls
55 | tb_dict.update(tb_dict_1)
56 | return point_loss, tb_dict
57 |
58 | def forward(self, batch_dict):
59 | """
60 | Args:
61 | batch_dict:
62 | batch_size:
63 | point_features: (N1 + N2 + N3 + ..., C) or (B, N, C)
64 | point_features_before_fusion: (N1 + N2 + N3 + ..., C)
65 | point_coords: (N1 + N2 + N3 + ..., 4) [bs_idx, x, y, z]
66 | point_labels (optional): (N1 + N2 + N3 + ...)
67 | gt_boxes (optional): (B, M, 8)
68 | Returns:
69 | batch_dict:
70 | point_cls_scores: (N1 + N2 + N3 + ..., 1)
71 | point_part_offset: (N1 + N2 + N3 + ..., 3)
72 | """
73 | if self.model_cfg.get('USE_POINT_FEATURES_BEFORE_FUSION', False):
74 | point_features = batch_dict['point_features_before_fusion']
75 | else:
76 | point_features = batch_dict['point_features']
77 | point_cls_preds = self.cls_layers(point_features) # (total_points, num_class)
78 |
79 | ret_dict = {
80 | 'point_cls_preds': point_cls_preds,
81 | }
82 |
83 | point_cls_scores = torch.sigmoid(point_cls_preds)
84 | batch_dict['point_cls_scores'], _ = point_cls_scores.max(dim=-1)
85 |
86 | if self.training:
87 | targets_dict = self.assign_targets(batch_dict)
88 | ret_dict['point_cls_labels'] = targets_dict['point_cls_labels']
89 | self.forward_ret_dict = ret_dict
90 |
91 | return batch_dict
92 |
--------------------------------------------------------------------------------
/pcdet/models/detectors/PartA2_net.py:
--------------------------------------------------------------------------------
1 | from .detector3d_template import Detector3DTemplate
2 |
3 |
4 | class PartA2Net(Detector3DTemplate):
5 | def __init__(self, model_cfg, num_class, dataset):
6 | super().__init__(model_cfg=model_cfg, num_class=num_class, dataset=dataset)
7 | self.module_list = self.build_networks()
8 |
9 | def forward(self, batch_dict, **kwargs):
10 | for cur_module in self.module_list:
11 | batch_dict = cur_module(batch_dict)
12 |
13 | if self.training:
14 | loss, tb_dict, disp_dict = self.get_training_loss()
15 |
16 | ret_dict = {
17 | 'loss': loss
18 | }
19 | return ret_dict, tb_dict, disp_dict
20 | else:
21 | pred_dicts, recall_dicts = self.post_processing(batch_dict)
22 | return pred_dicts, recall_dicts
23 |
24 | def get_training_loss(self):
25 | disp_dict = {}
26 | loss_rpn, tb_dict = self.dense_head.get_loss()
27 | loss_point, tb_dict = self.point_head.get_loss(tb_dict)
28 | loss_rcnn, tb_dict = self.roi_head.get_loss(tb_dict)
29 |
30 | loss = loss_rpn + loss_point + loss_rcnn
31 | return loss, tb_dict, disp_dict
32 |
--------------------------------------------------------------------------------
/pcdet/models/detectors/__init__.py:
--------------------------------------------------------------------------------
1 | from .detector3d_template import Detector3DTemplate
2 | from .PartA2_net import PartA2Net
3 | from .point_rcnn import PointRCNN
4 | from .pointpillar import PointPillar
5 | from .pv_rcnn import PVRCNN
6 | from .second_net import SECONDNet
7 | from .second_net_iou import SECONDNetIoU
8 | from .caddn import CaDDN
9 | from .voxel_rcnn import VoxelRCNN
10 | from .centerpoint import CenterPoint
11 | from .pv_rcnn_plusplus import PVRCNNPlusPlus
12 |
13 |
14 | __all__ = {
15 | 'Detector3DTemplate': Detector3DTemplate,
16 | 'SECONDNet': SECONDNet,
17 | 'PartA2Net': PartA2Net,
18 | 'PVRCNN': PVRCNN,
19 | 'PointPillar': PointPillar,
20 | 'PointRCNN': PointRCNN,
21 | 'SECONDNetIoU': SECONDNetIoU,
22 | 'CaDDN': CaDDN,
23 | 'VoxelRCNN': VoxelRCNN,
24 | 'CenterPoint': CenterPoint,
25 | 'PVRCNNPlusPlus': PVRCNNPlusPlus
26 | }
27 |
28 |
29 | def build_detector(model_cfg, num_class, dataset):
30 | model = __all__[model_cfg.NAME](
31 | model_cfg=model_cfg, num_class=num_class, dataset=dataset
32 | )
33 |
34 | return model
35 |
--------------------------------------------------------------------------------
/pcdet/models/detectors/caddn.py:
--------------------------------------------------------------------------------
1 | from .detector3d_template import Detector3DTemplate
2 |
3 |
4 | class CaDDN(Detector3DTemplate):
5 | def __init__(self, model_cfg, num_class, dataset):
6 | super().__init__(model_cfg=model_cfg, num_class=num_class, dataset=dataset)
7 | self.module_list = self.build_networks()
8 |
9 | def forward(self, batch_dict):
10 | for cur_module in self.module_list:
11 | batch_dict = cur_module(batch_dict)
12 |
13 | if self.training:
14 | loss, tb_dict, disp_dict = self.get_training_loss()
15 |
16 | ret_dict = {
17 | 'loss': loss
18 | }
19 | return ret_dict, tb_dict, disp_dict
20 | else:
21 | pred_dicts, recall_dicts = self.post_processing(batch_dict)
22 | return pred_dicts, recall_dicts
23 |
24 | def get_training_loss(self):
25 | disp_dict = {}
26 |
27 | loss_rpn, tb_dict_rpn = self.dense_head.get_loss()
28 | loss_depth, tb_dict_depth = self.vfe.get_loss()
29 |
30 | tb_dict = {
31 | 'loss_rpn': loss_rpn.item(),
32 | 'loss_depth': loss_depth.item(),
33 | **tb_dict_rpn,
34 | **tb_dict_depth
35 | }
36 |
37 | loss = loss_rpn + loss_depth
38 | return loss, tb_dict, disp_dict
39 |
--------------------------------------------------------------------------------
/pcdet/models/detectors/point_rcnn.py:
--------------------------------------------------------------------------------
1 | from .detector3d_template import Detector3DTemplate
2 |
3 |
4 | class PointRCNN(Detector3DTemplate):
5 | def __init__(self, model_cfg, num_class, dataset):
6 | super().__init__(model_cfg=model_cfg, num_class=num_class, dataset=dataset)
7 | self.module_list = self.build_networks()
8 |
9 | def forward(self, batch_dict, **kwargs):
10 | for cur_module in self.module_list:
11 | batch_dict = cur_module(batch_dict)
12 |
13 | if self.training:
14 | loss, tb_dict, disp_dict = self.get_training_loss()
15 |
16 | ret_dict = {
17 | 'loss': loss
18 | }
19 | return ret_dict, tb_dict, disp_dict
20 | else:
21 | pred_dicts, recall_dicts = self.post_processing(batch_dict)
22 | return pred_dicts, recall_dicts
23 |
24 | def get_training_loss(self):
25 | disp_dict = {}
26 | loss_point, tb_dict = self.point_head.get_loss()
27 | loss_rcnn, tb_dict = self.roi_head.get_loss(tb_dict)
28 |
29 | loss = loss_point + loss_rcnn
30 | return loss, tb_dict, disp_dict
31 |
--------------------------------------------------------------------------------
/pcdet/models/detectors/pointpillar.py:
--------------------------------------------------------------------------------
1 | import time
2 | from pcdet.utils import common_utils
3 | from .detector3d_template import Detector3DTemplate
4 | from pcdet.models.kd_heads.anchor_head.anchor_kd_head import AnchorHeadKD
5 |
6 |
7 | class PointPillar(Detector3DTemplate):
8 | def __init__(self, model_cfg, num_class, dataset):
9 | super().__init__(model_cfg=model_cfg, num_class=num_class, dataset=dataset)
10 | self.module_list = self.build_networks()
11 |
12 | # for KD only
13 | self.kd_head = AnchorHeadKD(self.model_cfg, self.dense_head) if model_cfg.get('KD', None) else None
14 | self.dense_head.kd_head = self.kd_head
15 |
16 | # for measure time only
17 | self.module_time_meter = common_utils.DictAverageMeter()
18 |
19 | def forward(self, batch_dict, record_time=False, **kwargs):
20 | for cur_module in self.module_list:
21 | if record_time:
22 | end = time.time()
23 | batch_dict = cur_module(batch_dict)
24 | if record_time:
25 | module_name = str(type(cur_module)).split('.')[-1][:-2]
26 | self.module_time_meter.update(module_name, (time.time() - end)*1000)
27 |
28 | if self.training:
29 | loss, tb_dict, disp_dict = self.get_training_loss()
30 |
31 | # only student model has KD_LOSS cfg
32 | if self.model_cfg.get('KD_LOSS', None) and self.model_cfg.KD_LOSS.ENABLED:
33 | kd_loss, tb_dict, disp_dict = self.get_kd_loss(batch_dict, tb_dict, disp_dict)
34 | loss += kd_loss
35 |
36 | ret_dict = {
37 | 'loss': loss
38 | }
39 | return ret_dict, tb_dict, disp_dict
40 | else:
41 | if record_time:
42 | end = time.time()
43 | pred_dicts, recall_dicts = self.post_processing(batch_dict)
44 | if record_time:
45 | self.module_time_meter.update('post_processing', (time.time() - end)*1000)
46 | return pred_dicts, recall_dicts
47 |
48 | def get_training_loss(self):
49 | disp_dict = {}
50 |
51 | loss_rpn, tb_dict = self.dense_head.get_loss()
52 | tb_dict = {
53 | 'loss_rpn': loss_rpn.item(),
54 | **tb_dict
55 | }
56 |
57 | loss = loss_rpn
58 | return loss, tb_dict, disp_dict
59 |
--------------------------------------------------------------------------------
/pcdet/models/detectors/pv_rcnn.py:
--------------------------------------------------------------------------------
1 | import time
2 | from .detector3d_template import Detector3DTemplate
3 |
4 |
5 | class PVRCNN(Detector3DTemplate):
6 | def __init__(self, model_cfg, num_class, dataset):
7 | super().__init__(model_cfg=model_cfg, num_class=num_class, dataset=dataset)
8 | self.module_list = self.build_networks()
9 |
10 | def forward(self, batch_dict, **kwargs):
11 | for cur_module in self.module_list:
12 | # end = time.time()
13 | batch_dict = cur_module(batch_dict)
14 | # module_name = str(type(cur_module)).split('.')[-1][:-2]
15 | # print('{}: {:.2f}'.format(module_name, (time.time() - end) * 1000))
16 |
17 | if self.training:
18 | loss, tb_dict, disp_dict = self.get_training_loss()
19 |
20 | ret_dict = {
21 | 'loss': loss
22 | }
23 | return ret_dict, tb_dict, disp_dict
24 | else:
25 | pred_dicts, recall_dicts = self.post_processing(batch_dict)
26 | return pred_dicts, recall_dicts
27 |
28 | def get_training_loss(self):
29 | disp_dict = {}
30 | loss_rpn, tb_dict = self.dense_head.get_loss()
31 | loss_point, tb_dict = self.point_head.get_loss(tb_dict)
32 | loss_rcnn, tb_dict = self.roi_head.get_loss(tb_dict)
33 |
34 | loss = loss_rpn + loss_point + loss_rcnn
35 | return loss, tb_dict, disp_dict
36 |
--------------------------------------------------------------------------------
/pcdet/models/detectors/pv_rcnn_plusplus.py:
--------------------------------------------------------------------------------
1 | from .detector3d_template import Detector3DTemplate
2 | from pcdet.models.kd_heads.center_head.center_kd_head import CenterHeadKD
3 |
4 |
5 | class PVRCNNPlusPlus(Detector3DTemplate):
6 | def __init__(self, model_cfg, num_class, dataset):
7 | super().__init__(model_cfg=model_cfg, num_class=num_class, dataset=dataset)
8 | self.module_list = self.build_networks()
9 |
10 | if self.dense_head is None and self.dense_head_aux is not None:
11 | self.dense_head = self.dense_head_aux
12 |
13 | self.kd_head = CenterHeadKD(self.model_cfg, self.dense_head) if model_cfg.get('KD', None) else None
14 |
15 | self.dense_head.kd_head = self.kd_head
16 |
17 | def forward(self, batch_dict, **kwargs):
18 | batch_dict = self.vfe(batch_dict)
19 | batch_dict = self.backbone_3d(batch_dict)
20 | batch_dict = self.map_to_bev_module(batch_dict)
21 | batch_dict = self.backbone_2d(batch_dict)
22 | batch_dict = self.dense_head(batch_dict)
23 |
24 | batch_dict = self.roi_head.proposal_layer(
25 | batch_dict, nms_config=self.roi_head.model_cfg.NMS_CONFIG['TRAIN' if self.training else 'TEST']
26 | )
27 | if self.training:
28 | targets_dict = self.roi_head.assign_targets(batch_dict)
29 | batch_dict['rois'] = targets_dict['rois']
30 | batch_dict['roi_labels'] = targets_dict['roi_labels']
31 | batch_dict['roi_targets_dict'] = targets_dict
32 | num_rois_per_scene = targets_dict['rois'].shape[1]
33 | if 'roi_valid_num' in batch_dict:
34 | batch_dict['roi_valid_num'] = [num_rois_per_scene for _ in range(batch_dict['batch_size'])]
35 |
36 | batch_dict = self.pfe(batch_dict)
37 | batch_dict = self.point_head(batch_dict)
38 | batch_dict = self.roi_head(batch_dict)
39 |
40 | if self.is_teacher and self.training:
41 | return batch_dict
42 |
43 | if self.training:
44 | loss, tb_dict, disp_dict = self.get_training_loss()
45 |
46 | ret_dict = {
47 | 'loss': loss
48 | }
49 | return ret_dict, tb_dict, disp_dict
50 | else:
51 | pred_dicts, recall_dicts = self.post_processing(batch_dict)
52 | return pred_dicts, recall_dicts
53 |
54 | def get_training_loss(self):
55 | disp_dict = {}
56 | loss_rpn, tb_dict = self.dense_head.get_loss()
57 | if self.point_head is not None:
58 | loss_point, tb_dict = self.point_head.get_loss(tb_dict)
59 | else:
60 | loss_point = 0
61 | loss_rcnn, tb_dict = self.roi_head.get_loss(tb_dict)
62 |
63 | loss = loss_rpn + loss_point + loss_rcnn
64 | return loss, tb_dict, disp_dict
65 |
--------------------------------------------------------------------------------
/pcdet/models/detectors/second_net.py:
--------------------------------------------------------------------------------
1 | import time
2 | from pcdet.utils import common_utils
3 | from .detector3d_template import Detector3DTemplate
4 | from pcdet.models.kd_heads.anchor_head.anchor_kd_head import AnchorHeadKD
5 |
6 |
7 | class SECONDNet(Detector3DTemplate):
8 | def __init__(self, model_cfg, num_class, dataset):
9 | super().__init__(model_cfg=model_cfg, num_class=num_class, dataset=dataset)
10 | self.module_list = self.build_networks()
11 |
12 | # for KD only
13 | self.kd_head = AnchorHeadKD(self.model_cfg, self.dense_head) if model_cfg.get('KD', None) else None
14 | self.dense_head.kd_head = self.kd_head
15 |
16 | # for measure time only
17 | self.module_time_meter = common_utils.DictAverageMeter()
18 |
19 | def forward(self, batch_dict, record_time=None, **kwargs):
20 | for cur_module in self.module_list:
21 | if record_time:
22 | end = time.time()
23 | batch_dict = cur_module(batch_dict)
24 | if record_time:
25 | module_name = str(type(cur_module)).split('.')[-1][:-2]
26 | self.module_time_meter.update(module_name, (time.time() - end)*1000)
27 |
28 | if self.training:
29 | loss, tb_dict, disp_dict = self.get_training_loss()
30 |
31 | # only student model has KD_LOSS cfg
32 | if self.model_cfg.get('KD_LOSS', None) and self.model_cfg.KD_LOSS.ENABLED:
33 | kd_loss, tb_dict, disp_dict = self.get_kd_loss(batch_dict, tb_dict, disp_dict)
34 | loss += kd_loss
35 |
36 | ret_dict = {
37 | 'loss': loss
38 | }
39 | return ret_dict, tb_dict, disp_dict
40 | else:
41 | if record_time:
42 | end = time.time()
43 | pred_dicts, recall_dicts = self.post_processing(batch_dict)
44 | # import ipdb;
45 | # ipdb.set_trace(context=20)
46 | # self.dataset.__vis_open3d__(points=batch_dict['points'][:, 1:].cpu().numpy(),
47 | # gt_boxes=batch_dict['gt_boxes'][0].detach().cpu().numpy(),
48 | # ref_boxes=pred_dicts[0]['pred_boxes'].cpu().numpy())
49 | if record_time:
50 | self.module_time_meter.update('post_processing', (time.time() - end)*1000)
51 | return pred_dicts, recall_dicts
52 |
53 | def get_training_loss(self):
54 | disp_dict = {}
55 |
56 | loss_rpn, tb_dict = self.dense_head.get_loss()
57 | tb_dict = {
58 | 'loss_rpn': loss_rpn.item(),
59 | **tb_dict
60 | }
61 |
62 | loss = loss_rpn
63 | return loss, tb_dict, disp_dict
64 |
--------------------------------------------------------------------------------
/pcdet/models/detectors/voxel_rcnn.py:
--------------------------------------------------------------------------------
1 | from .detector3d_template import Detector3DTemplate
2 | from pcdet.models.kd_heads.center_head.center_kd_head import CenterHeadKD
3 |
4 |
5 | class VoxelRCNN(Detector3DTemplate):
6 | def __init__(self, model_cfg, num_class, dataset):
7 | super().__init__(model_cfg=model_cfg, num_class=num_class, dataset=dataset)
8 | self.module_list = self.build_networks()
9 |
10 | if self.dense_head is None and self.dense_head_aux is not None:
11 | self.dense_head = self.dense_head_aux
12 |
13 | self.kd_head = CenterHeadKD(self.model_cfg, self.dense_head) if model_cfg.get('KD', None) else None
14 | self.dense_head.kd_head = self.kd_head
15 | if self.kd_head is not None:
16 | self.kd_head.roi_head = self.roi_head
17 |
18 | def forward(self, batch_dict, **kwargs):
19 | for cur_module in self.module_list:
20 | batch_dict = cur_module(batch_dict)
21 |
22 | if self.is_teacher and self.training:
23 | return batch_dict
24 |
25 | if self.training:
26 | loss, tb_dict, disp_dict = self.get_training_loss()
27 |
28 | # only student model has KD_LOSS cfg
29 | if self.model_cfg.get('KD_LOSS', None) and self.model_cfg.KD_LOSS.ENABLED:
30 | kd_loss, tb_dict, disp_dict = self.get_kd_loss(batch_dict, tb_dict, disp_dict)
31 | loss += kd_loss
32 |
33 | ret_dict = {
34 | 'loss': loss
35 | }
36 | return ret_dict, tb_dict, disp_dict
37 | else:
38 | pred_dicts, recall_dicts = self.post_processing(batch_dict)
39 | return pred_dicts, recall_dicts
40 |
41 | def get_training_loss(self):
42 | disp_dict = {}
43 | loss = 0
44 |
45 | loss_rpn, tb_dict = self.dense_head.get_loss()
46 | loss_rcnn, tb_dict = self.roi_head.get_loss(tb_dict)
47 |
48 | loss = loss + loss_rpn + loss_rcnn
49 | return loss, tb_dict, disp_dict
50 |
--------------------------------------------------------------------------------
/pcdet/models/kd_adapt_block/__init__.py:
--------------------------------------------------------------------------------
1 | from .kd_adapt_block import KDAdaptBlock
2 |
3 | __all__ = {
4 | 'KDAdaptBlock': KDAdaptBlock,
5 | }
6 |
--------------------------------------------------------------------------------
/pcdet/models/kd_adapt_block/kd_adapt_block.py:
--------------------------------------------------------------------------------
1 | import torch
2 | import torch.nn as nn
3 | from torch.nn.init import kaiming_normal_
4 |
5 | from pcdet.models.model_utils.basic_block_2d import build_block
6 |
7 |
8 | class KDAdaptBlock(nn.Module):
9 | def __init__(self, model_cfg, point_cloud_range):
10 | super().__init__()
11 | self.model_cfg = model_cfg
12 | self.point_cloud_range = point_cloud_range
13 | self.align_module_list = []
14 |
15 | for adapt_layer_name, adapt_layer_cfg in self.model_cfg.MODULE.items():
16 | self.add_module(adapt_layer_name.lower(), BasicAdaptLayer(adapt_layer_cfg))
17 | self.align_module_list.append(getattr(self, adapt_layer_name.lower()))
18 |
19 | def forward(self, batch_dict):
20 | if self.training:
21 | for adapt_layer in self.align_module_list:
22 | batch_dict = adapt_layer(batch_dict)
23 |
24 | return batch_dict
25 |
26 |
27 | class BasicAdaptLayer(nn.Module):
28 | def __init__(self, block_cfg):
29 | super().__init__()
30 | self.block_cfg = block_cfg
31 | self.in_feature_name = block_cfg.in_feature_name
32 | self.out_feature_name = block_cfg.out_feature_name
33 |
34 | self.build_adaptation_layer(block_cfg)
35 |
36 | self.init_weights(weight_init='xavier')
37 |
38 | def init_weights(self, weight_init='xavier'):
39 | if weight_init == 'kaiming':
40 | init_func = nn.init.kaiming_normal_
41 | elif weight_init == 'xavier':
42 | init_func = nn.init.xavier_normal_
43 | elif weight_init == 'normal':
44 | init_func = nn.init.normal_
45 | else:
46 | raise NotImplementedError
47 |
48 | for m in self.modules():
49 | if isinstance(m, nn.Conv2d) or isinstance(m, nn.Conv1d):
50 | if weight_init == 'normal':
51 | init_func(m.weight, mean=0, std=0.001)
52 | else:
53 | init_func(m.weight)
54 | if m.bias is not None:
55 | nn.init.constant_(m.bias, 0)
56 |
57 | def build_adaptation_layer(self, block_cfg):
58 | align_block = []
59 |
60 | in_channel = block_cfg.in_channel
61 | block_types = block_cfg.block_type
62 | num_filters = block_cfg.num_filters
63 | kernel_sizes = block_cfg.kernel_size
64 | num_strides = block_cfg.strides
65 | paddings = block_cfg.padding
66 |
67 | for i in range(len(num_filters)):
68 | align_block.extend(build_block(
69 | block_types[i], in_channel, num_filters[i], kernel_size=kernel_sizes[i],
70 | stride=num_strides[i], padding=paddings[i], bias=False
71 | ))
72 |
73 | self.adapt_layer = nn.Sequential(*align_block)
74 |
75 | def forward(self, batch_dict):
76 | in_feature = batch_dict[self.in_feature_name]
77 |
78 | out_feature = self.adapt_layer(in_feature)
79 |
80 | batch_dict[self.out_feature_name] = out_feature
81 |
82 | return batch_dict
83 |
--------------------------------------------------------------------------------
/pcdet/models/kd_heads/__init__.py:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/CVMI-Lab/SparseKD/2e00edef6b94fb208ba64193940bc856e5effbcb/pcdet/models/kd_heads/__init__.py
--------------------------------------------------------------------------------
/pcdet/models/kd_heads/anchor_head/__init__.py:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/CVMI-Lab/SparseKD/2e00edef6b94fb208ba64193940bc856e5effbcb/pcdet/models/kd_heads/anchor_head/__init__.py
--------------------------------------------------------------------------------
/pcdet/models/kd_heads/anchor_head/anchor_kd_head.py:
--------------------------------------------------------------------------------
1 | import torch
2 | import torch.nn as nn
3 | import numpy as np
4 |
5 | from .anchor_logit_kd_head import AnchorLogitKDHead
6 | from .anchor_feature_kd_head import AnchorFeatureKDHead
7 | from .anchor_label_kd_head import AnchorLabelAssignKDHead
8 |
9 |
10 | class AnchorHeadKD(AnchorLogitKDHead, AnchorFeatureKDHead, AnchorLabelAssignKDHead):
11 | def __init__(self, model_cfg, dense_head):
12 | super(AnchorHeadKD, self).__init__(model_cfg, dense_head)
13 | self.build_loss(dense_head)
14 |
15 | def get_kd_loss(self, batch_dict, tb_dict):
16 | kd_loss = 0.0
17 | if self.model_cfg.get('LOGIT_KD', None) and self.model_cfg.LOGIT_KD.ENABLED:
18 | kd_logit_loss, tb_dict = self.get_logit_kd_loss(batch_dict, tb_dict)
19 | kd_loss += kd_logit_loss
20 |
21 | if self.model_cfg.get('FEATURE_KD', None) and self.model_cfg.FEATURE_KD.ENABLED:
22 | kd_feature_loss, tb_dict = self.get_feature_kd_loss(
23 | batch_dict, tb_dict, self.model_cfg.KD_LOSS.FEATURE_LOSS
24 | )
25 | kd_loss += kd_feature_loss
26 |
27 | return kd_loss, tb_dict
28 |
--------------------------------------------------------------------------------
/pcdet/models/kd_heads/anchor_head/anchor_label_kd_head.py:
--------------------------------------------------------------------------------
1 | import torch
2 | import numpy as np
3 |
4 | from pcdet.ops.iou3d_nms import iou3d_nms_utils
5 | from pcdet.models.kd_heads.kd_head import KDHeadTemplate
6 |
7 |
8 | class AnchorLabelAssignKDHead(KDHeadTemplate):
9 | def __init__(self, model_cfg, dense_head):
10 | super(AnchorLabelAssignKDHead, self).__init__(model_cfg, dense_head)
11 |
--------------------------------------------------------------------------------
/pcdet/models/kd_heads/center_head/__init__.py:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/CVMI-Lab/SparseKD/2e00edef6b94fb208ba64193940bc856e5effbcb/pcdet/models/kd_heads/center_head/__init__.py
--------------------------------------------------------------------------------
/pcdet/models/kd_heads/center_head/center_roi_kd_head.py:
--------------------------------------------------------------------------------
1 | import torch
2 |
3 | from pcdet.models.kd_heads.kd_head import KDHeadTemplate
4 | from pcdet.utils import loss_utils
5 |
6 |
7 | class CenterRoIKDHead(KDHeadTemplate):
8 | def __init__(self, model_cfg, dense_head):
9 | super().__init__(model_cfg, dense_head)
10 |
11 | def build_roi_kd_loss(self):
12 | # logit kd hm loss
13 | if self.model_cfg.KD_LOSS.ROI_CLS_LOSS.type in ['SmoothL1Loss', 'MSELoss']:
14 | self.kd_roi_cls_loss_func = getattr(torch.nn, self.model_cfg.KD_LOSS.ROI_CLS_LOSS.type)(reduction='none')
15 | else:
16 | raise NotImplementedError
17 |
18 | # logit kd regression loss
19 | if self.model_cfg.KD_LOSS.ROI_REG_LOSS.type == 'WeightedSmoothL1Loss':
20 | self.kd_roi_reg_loss_func = getattr(loss_utils, self.model_cfg.KD_LOSS.ROI_REG_LOSS.type)(
21 | code_weights=self.model_cfg.KD_LOSS.ROI_REG_LOSS.code_weights
22 | )
23 | else:
24 | raise NotImplementedError
25 |
26 | def get_roi_kd_loss(self, batch_dict, tb_dict):
27 | loss_cfg = self.model_cfg.KD_LOSS
28 |
29 | rcnn_cls_stu = self.roi_head.forward_ret_dict['rcnn_cls']
30 | rcnn_cls_tea = batch_dict['rcnn_cls_tea']
31 |
32 | # cls loss
33 | if loss_cfg.ROI_CLS_LOSS.weight == 0:
34 | kd_roi_cls_loss_raw = 0
35 | elif loss_cfg.ROI_CLS_LOSS.type in ['SmoothL1Loss', 'MSELoss']:
36 | kd_roi_cls_loss_raw = self.kd_roi_cls_loss_func(rcnn_cls_stu, rcnn_cls_tea).mean()
37 | else:
38 | raise NotImplementedError
39 | kd_roi_cls_loss = kd_roi_cls_loss_raw * loss_cfg.ROI_CLS_LOSS.weight
40 |
41 | reg_valid_mask = self.roi_head.forward_ret_dict['reg_valid_mask'].view(-1)
42 | rcnn_reg_stu = self.roi_head.forward_ret_dict['rcnn_reg']
43 | rcnn_reg_tea = batch_dict['rcnn_reg_tea']
44 | fg_mask = (reg_valid_mask > 0)
45 | fg_sum = fg_mask.long().sum().item()
46 | rcnn_batch_size = reg_valid_mask.shape[0]
47 |
48 | # reg loss
49 | if loss_cfg.ROI_REG_LOSS.weight == 0:
50 | kd_roi_reg_loss_raw = 0
51 | elif loss_cfg.ROI_CLS_LOSS.type in ['SmoothL1Loss', 'MSELoss']:
52 | kd_roi_reg_loss_raw = self.kd_roi_cls_loss_func(
53 | rcnn_reg_stu.unsqueeze(0), rcnn_reg_tea.unsqueeze(0)
54 | )
55 | kd_roi_reg_loss_raw = (kd_roi_reg_loss_raw.view(rcnn_batch_size, -1) * \
56 | fg_mask.unsqueeze(dim=-1).float()).sum() / max(fg_sum, 1)
57 | else:
58 | raise NotImplementedError
59 | kd_roi_reg_loss = kd_roi_reg_loss_raw * loss_cfg.ROI_REG_LOSS.weight
60 |
61 | kd_roi_loss = kd_roi_cls_loss + kd_roi_reg_loss
62 |
63 | tb_dict['kd_r-c_ls'] = kd_roi_cls_loss if isinstance(kd_roi_cls_loss, float) else kd_roi_cls_loss.item()
64 | tb_dict['kd_r-r_ls'] = kd_roi_reg_loss if isinstance(kd_roi_reg_loss, float) else kd_roi_reg_loss.item()
65 |
66 | return kd_roi_loss, tb_dict
67 |
--------------------------------------------------------------------------------
/pcdet/models/model_utils/batch_norm_utils.py:
--------------------------------------------------------------------------------
1 | import torch
2 | import torch.nn as nn
3 | import time
4 |
5 |
6 | class SparseBatchNorm1d(nn.Module):
7 | def __init__(self, channel, **kwargs):
8 | super(SparseBatchNorm1d, self).__init__()
9 |
10 | self.batch_norm = nn.BatchNorm1d(channel, **kwargs)
11 |
12 | def forward(self, input):
13 | """
14 | Args:
15 | input:
16 |
17 | Returns:
18 |
19 | """
20 | if input.dim() != 4:
21 | raise ValueError(
22 | "expected 4D input (got {}D input)".format(input.dim())
23 | )
24 |
25 | # torch.backends.cudnn.enabled = False
26 | input_bhwc = input.permute(0, 2, 3, 1)
27 | mask = torch.any(input_bhwc > 0, dim=-1)
28 | valid_input = input_bhwc[mask, :]
29 | valid_out = self.batch_norm(valid_input)
30 | input_bhwc[mask] = valid_out
31 | out = input_bhwc.permute(0, 3, 1, 2)
32 | # torch.backends.cudnn.enabled = True
33 |
34 | return out
35 |
36 |
37 | class FastSparseBatchNorm1d(nn.Module):
38 | def __init__(self, channel, **kwargs):
39 | super(FastSparseBatchNorm1d, self).__init__()
40 |
41 | self.batch_norm = nn.BatchNorm1d(channel, **kwargs)
42 |
43 | def forward(self, input):
44 | """
45 | Args:
46 | input:
47 |
48 | Returns:
49 |
50 | """
51 | if input.dim() != 4:
52 | raise ValueError(
53 | "expected 4D input (got {}D input)".format(input.dim())
54 | )
55 |
56 | torch.backends.cudnn.enabled = False
57 | input_bhwc = input.permute(0, 2, 3, 1)
58 | mask = torch.any(input_bhwc > 0, dim=-1)
59 | valid_input = input_bhwc[mask, :]
60 | valid_out = self.batch_norm(valid_input)
61 | input_bhwc[mask] = valid_out
62 | out = input_bhwc.permute(0, 3, 1, 2)
63 | torch.backends.cudnn.enabled = True
64 |
65 | return out
66 |
67 |
68 | def get_norm_layer(name):
69 | if name == 'BatchNorm2d':
70 | return nn.BatchNorm2d
71 | elif name == 'SparseBatchNorm1d':
72 | return SparseBatchNorm1d
73 | elif name == 'FastSparseBatchNorm1d':
74 | return FastSparseBatchNorm1d
75 | else:
76 | raise ValueError(name)
77 |
--------------------------------------------------------------------------------
/pcdet/models/model_utils/model_nms_utils.py:
--------------------------------------------------------------------------------
1 | import torch
2 |
3 | from ...ops.iou3d_nms import iou3d_nms_utils
4 |
5 |
6 | def class_agnostic_nms(box_scores, box_preds, nms_config, score_thresh=None):
7 | src_box_scores = box_scores
8 | if score_thresh is not None:
9 | scores_mask = (box_scores >= score_thresh)
10 | box_scores = box_scores[scores_mask]
11 | box_preds = box_preds[scores_mask]
12 |
13 | selected = []
14 | if nms_config.NMS_THRESH != 1.1 and box_scores.shape[0] > 0:
15 | box_scores_nms, indices = torch.topk(box_scores, k=min(nms_config.NMS_PRE_MAXSIZE, box_scores.shape[0]))
16 | boxes_for_nms = box_preds[indices]
17 | keep_idx, selected_scores = getattr(iou3d_nms_utils, nms_config.NMS_TYPE)(
18 | boxes_for_nms[:, 0:7], box_scores_nms, nms_config.NMS_THRESH, **nms_config
19 | )
20 | selected = indices[keep_idx[:nms_config.NMS_POST_MAXSIZE]]
21 | else:
22 | selected = torch.arange(box_scores.shape[0])
23 |
24 | if score_thresh is not None:
25 | original_idxs = scores_mask.nonzero().view(-1)
26 | selected = original_idxs[selected]
27 | return selected, src_box_scores[selected]
28 |
29 |
30 | def multi_classes_nms(cls_scores, box_preds, nms_config, score_thresh=None):
31 | """
32 | Args:
33 | cls_scores: (N, num_class)
34 | box_preds: (N, 7 + C)
35 | nms_config:
36 | score_thresh:
37 |
38 | Returns:
39 |
40 | """
41 | pred_scores, pred_labels, pred_boxes = [], [], []
42 | for k in range(cls_scores.shape[1]):
43 | if score_thresh is not None:
44 | scores_mask = (cls_scores[:, k] >= score_thresh)
45 | box_scores = cls_scores[scores_mask, k]
46 | cur_box_preds = box_preds[scores_mask]
47 | else:
48 | box_scores = cls_scores[:, k]
49 | cur_box_preds = box_preds
50 |
51 | selected = []
52 | if box_scores.shape[0] > 0:
53 | box_scores_nms, indices = torch.topk(box_scores, k=min(nms_config.NMS_PRE_MAXSIZE, box_scores.shape[0]))
54 | boxes_for_nms = cur_box_preds[indices]
55 | keep_idx, selected_scores = getattr(iou3d_nms_utils, nms_config.NMS_TYPE)(
56 | boxes_for_nms[:, 0:7], box_scores_nms, nms_config.NMS_THRESH, **nms_config
57 | )
58 | selected = indices[keep_idx[:nms_config.NMS_POST_MAXSIZE]]
59 |
60 | pred_scores.append(box_scores[selected])
61 | pred_labels.append(box_scores.new_ones(len(selected)).long() * k)
62 | pred_boxes.append(cur_box_preds[selected])
63 |
64 | pred_scores = torch.cat(pred_scores, dim=0)
65 | pred_labels = torch.cat(pred_labels, dim=0)
66 | pred_boxes = torch.cat(pred_boxes, dim=0)
67 |
68 | return pred_scores, pred_labels, pred_boxes
69 |
--------------------------------------------------------------------------------
/pcdet/models/model_utils/rotated_roi_grid_pool.py:
--------------------------------------------------------------------------------
1 | import torch
2 | import torch.nn as nn
3 |
4 | from pcdet.utils import common_utils
5 |
6 |
7 | class RotatedGridPool(nn.Module):
8 | def __init__(self, point_cloud_range, pool_cfg):
9 | super(RotatedGridPool, self).__init__()
10 | self.min_x = point_cloud_range[0]
11 | self.min_y = point_cloud_range[1]
12 | self.grid_size = pool_cfg.GRID_SIZE
13 |
14 | def forward(self, features_2d, rois, voxel_size, feature_map_stride):
15 | """
16 | Args:
17 | features_2d: (B, C, H, W)
18 | rois: (B, num_rois, 7 + C) tensor or list [num_rois, 7 + C]
19 | voxel_size
20 | Returns:
21 | """
22 | batch_size = features_2d.shape[0]
23 | height, width = features_2d.size(2), features_2d.size(3)
24 |
25 | pooled_features_list = []
26 | torch.backends.cudnn.enabled = False
27 | for b_id in range(batch_size):
28 | batch_rois = rois[b_id]
29 | if batch_rois.shape[0] == 0:
30 | continue
31 |
32 | valid_mask = batch_rois[:, 3] != 0
33 | # no valid box in the scene
34 | if valid_mask.sum() == 0:
35 | continue
36 | batch_rois = batch_rois[valid_mask]
37 |
38 | voxel_size_x = voxel_size[0]
39 | voxel_size_y = voxel_size[1]
40 |
41 | # Map global boxes coordinates to feature map coordinates
42 | x1 = (batch_rois[:, 0] - batch_rois[:, 3] / 2 - self.min_x) / (voxel_size_x * feature_map_stride)
43 | x2 = (batch_rois[:, 0] + batch_rois[:, 3] / 2 - self.min_x) / (voxel_size_x * feature_map_stride)
44 | y1 = (batch_rois[:, 1] - batch_rois[:, 4] / 2 - self.min_y) / (voxel_size_y * feature_map_stride)
45 | y2 = (batch_rois[:, 1] + batch_rois[:, 4] / 2 - self.min_y) / (voxel_size_y * feature_map_stride)
46 |
47 | angle, _ = common_utils.check_numpy_to_torch(batch_rois[:, 6])
48 |
49 | cosa = torch.cos(angle)
50 | sina = torch.sin(angle)
51 |
52 | theta = torch.stack((
53 | (x2 - x1) / (width - 1) * cosa, (x2 - x1) / (width - 1) * (-sina), (x1 + x2 - width + 1) / (width - 1),
54 | (y2 - y1) / (height - 1) * sina, (y2 - y1) / (height - 1) * cosa, (y1 + y2 - height + 1) / (height - 1)
55 | ), dim=1).view(-1, 2, 3).float()
56 |
57 | # Correct grid
58 | scale1 = (y2 - y1) / torch.clamp(x2 - x1, min=0.01)
59 | scale2 = (x2 - x1) / torch.clamp(y2 - y1, min=0.01)
60 |
61 | theta[:, 0, 1] *= scale1
62 | theta[:, 1, 0] *= scale2
63 |
64 | grid = nn.functional.affine_grid(
65 | theta,
66 | torch.Size((batch_rois.size(0), features_2d.size(1), self.grid_size, self.grid_size))
67 | )
68 |
69 | new_grid = grid.view(1, batch_rois.size(0), self.grid_size * self.grid_size, 2)
70 |
71 | pooled_features = nn.functional.grid_sample(
72 | features_2d[b_id].unsqueeze(0), new_grid
73 | ).squeeze(0)
74 |
75 | pooled_features = pooled_features.permute(1, 0, 2)
76 | pooled_features = pooled_features.view(
77 | batch_rois.size(0), features_2d.size(1), self.grid_size, self.grid_size
78 | )
79 | pooled_features_list.append(pooled_features)
80 |
81 | torch.backends.cudnn.enabled = True
82 | pooled_features = torch.cat(pooled_features_list, dim=0)
83 |
84 | return pooled_features
85 |
--------------------------------------------------------------------------------
/pcdet/models/roi_heads/__init__.py:
--------------------------------------------------------------------------------
1 | from .partA2_head import PartA2FCHead
2 | from .pointrcnn_head import PointRCNNHead
3 | from .pvrcnn_head import PVRCNNHead
4 | from .second_head import SECONDHead
5 | from .voxelrcnn_head import VoxelRCNNHead
6 | from .roi_head_template import RoIHeadTemplate
7 |
8 |
9 | __all__ = {
10 | 'RoIHeadTemplate': RoIHeadTemplate,
11 | 'PartA2FCHead': PartA2FCHead,
12 | 'PVRCNNHead': PVRCNNHead,
13 | 'SECONDHead': SECONDHead,
14 | 'PointRCNNHead': PointRCNNHead,
15 | 'VoxelRCNNHead': VoxelRCNNHead
16 | }
17 |
--------------------------------------------------------------------------------
/pcdet/ops/iou3d_nms/src/iou3d_cpu.h:
--------------------------------------------------------------------------------
1 | #ifndef IOU3D_CPU_H
2 | #define IOU3D_CPU_H
3 |
4 | #include
5 | #include
6 | #include
7 | #include
8 |
9 | int boxes_iou_bev_cpu(at::Tensor boxes_a_tensor, at::Tensor boxes_b_tensor, at::Tensor ans_iou_tensor);
10 |
11 | #endif
12 |
--------------------------------------------------------------------------------
/pcdet/ops/iou3d_nms/src/iou3d_nms.h:
--------------------------------------------------------------------------------
1 | #ifndef IOU3D_NMS_H
2 | #define IOU3D_NMS_H
3 |
4 | #include
5 | #include
6 | #include
7 | #include
8 |
9 | int boxes_overlap_bev_gpu(at::Tensor boxes_a, at::Tensor boxes_b, at::Tensor ans_overlap);
10 | int boxes_iou_bev_gpu(at::Tensor boxes_a, at::Tensor boxes_b, at::Tensor ans_iou);
11 | int nms_gpu(at::Tensor boxes, at::Tensor keep, float nms_overlap_thresh);
12 | int nms_normal_gpu(at::Tensor boxes, at::Tensor keep, float nms_overlap_thresh);
13 |
14 | #endif
15 |
--------------------------------------------------------------------------------
/pcdet/ops/iou3d_nms/src/iou3d_nms_api.cpp:
--------------------------------------------------------------------------------
1 | #include
2 | #include
3 | #include
4 | #include
5 | #include
6 |
7 | #include "iou3d_cpu.h"
8 | #include "iou3d_nms.h"
9 |
10 |
11 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
12 | m.def("boxes_overlap_bev_gpu", &boxes_overlap_bev_gpu, "oriented boxes overlap");
13 | m.def("boxes_iou_bev_gpu", &boxes_iou_bev_gpu, "oriented boxes iou");
14 | m.def("nms_gpu", &nms_gpu, "oriented nms gpu");
15 | m.def("nms_normal_gpu", &nms_normal_gpu, "nms gpu");
16 | m.def("boxes_iou_bev_cpu", &boxes_iou_bev_cpu, "oriented boxes iou");
17 | }
18 |
--------------------------------------------------------------------------------
/pcdet/ops/pointnet2/pointnet2_batch/src/ball_query.cpp:
--------------------------------------------------------------------------------
1 | /*
2 | batch version of ball query, modified from the original implementation of official PointNet++ codes.
3 | Written by Shaoshuai Shi
4 | All Rights Reserved 2018.
5 | */
6 |
7 |
8 | #include
9 | #include
10 | #include
11 | #include
12 | #include
13 | #include "ball_query_gpu.h"
14 |
15 | extern THCState *state;
16 |
17 | #define CHECK_CUDA(x) do { \
18 | if (!x.type().is_cuda()) { \
19 | fprintf(stderr, "%s must be CUDA tensor at %s:%d\n", #x, __FILE__, __LINE__); \
20 | exit(-1); \
21 | } \
22 | } while (0)
23 | #define CHECK_CONTIGUOUS(x) do { \
24 | if (!x.is_contiguous()) { \
25 | fprintf(stderr, "%s must be contiguous tensor at %s:%d\n", #x, __FILE__, __LINE__); \
26 | exit(-1); \
27 | } \
28 | } while (0)
29 | #define CHECK_INPUT(x) CHECK_CUDA(x);CHECK_CONTIGUOUS(x)
30 |
31 |
32 | int ball_query_wrapper_fast(int b, int n, int m, float radius, int nsample,
33 | at::Tensor new_xyz_tensor, at::Tensor xyz_tensor, at::Tensor idx_tensor) {
34 | CHECK_INPUT(new_xyz_tensor);
35 | CHECK_INPUT(xyz_tensor);
36 | const float *new_xyz = new_xyz_tensor.data();
37 | const float *xyz = xyz_tensor.data();
38 | int *idx = idx_tensor.data();
39 |
40 | ball_query_kernel_launcher_fast(b, n, m, radius, nsample, new_xyz, xyz, idx);
41 | return 1;
42 | }
43 |
--------------------------------------------------------------------------------
/pcdet/ops/pointnet2/pointnet2_batch/src/ball_query_gpu.cu:
--------------------------------------------------------------------------------
1 | /*
2 | batch version of ball query, modified from the original implementation of official PointNet++ codes.
3 | Written by Shaoshuai Shi
4 | All Rights Reserved 2018.
5 | */
6 |
7 | #include
8 | #include
9 | #include
10 |
11 | #include "ball_query_gpu.h"
12 | #include "cuda_utils.h"
13 |
14 |
15 | __global__ void ball_query_kernel_fast(int b, int n, int m, float radius, int nsample,
16 | const float *__restrict__ new_xyz, const float *__restrict__ xyz, int *__restrict__ idx) {
17 | // new_xyz: (B, M, 3)
18 | // xyz: (B, N, 3)
19 | // output:
20 | // idx: (B, M, nsample)
21 | int bs_idx = blockIdx.y;
22 | int pt_idx = blockIdx.x * blockDim.x + threadIdx.x;
23 | if (bs_idx >= b || pt_idx >= m) return;
24 |
25 | new_xyz += bs_idx * m * 3 + pt_idx * 3;
26 | xyz += bs_idx * n * 3;
27 | idx += bs_idx * m * nsample + pt_idx * nsample;
28 |
29 | float radius2 = radius * radius;
30 | float new_x = new_xyz[0];
31 | float new_y = new_xyz[1];
32 | float new_z = new_xyz[2];
33 |
34 | int cnt = 0;
35 | for (int k = 0; k < n; ++k) {
36 | float x = xyz[k * 3 + 0];
37 | float y = xyz[k * 3 + 1];
38 | float z = xyz[k * 3 + 2];
39 | float d2 = (new_x - x) * (new_x - x) + (new_y - y) * (new_y - y) + (new_z - z) * (new_z - z);
40 | if (d2 < radius2){
41 | if (cnt == 0){
42 | for (int l = 0; l < nsample; ++l) {
43 | idx[l] = k;
44 | }
45 | }
46 | idx[cnt] = k;
47 | ++cnt;
48 | if (cnt >= nsample) break;
49 | }
50 | }
51 | }
52 |
53 |
54 | void ball_query_kernel_launcher_fast(int b, int n, int m, float radius, int nsample, \
55 | const float *new_xyz, const float *xyz, int *idx) {
56 | // new_xyz: (B, M, 3)
57 | // xyz: (B, N, 3)
58 | // output:
59 | // idx: (B, M, nsample)
60 |
61 | cudaError_t err;
62 |
63 | dim3 blocks(DIVUP(m, THREADS_PER_BLOCK), b); // blockIdx.x(col), blockIdx.y(row)
64 | dim3 threads(THREADS_PER_BLOCK);
65 |
66 | ball_query_kernel_fast<<>>(b, n, m, radius, nsample, new_xyz, xyz, idx);
67 | // cudaDeviceSynchronize(); // for using printf in kernel function
68 | err = cudaGetLastError();
69 | if (cudaSuccess != err) {
70 | fprintf(stderr, "CUDA kernel failed : %s\n", cudaGetErrorString(err));
71 | exit(-1);
72 | }
73 | }
74 |
--------------------------------------------------------------------------------
/pcdet/ops/pointnet2/pointnet2_batch/src/ball_query_gpu.h:
--------------------------------------------------------------------------------
1 | #ifndef _BALL_QUERY_GPU_H
2 | #define _BALL_QUERY_GPU_H
3 |
4 | #include
5 | #include
6 | #include
7 | #include
8 |
9 | int ball_query_wrapper_fast(int b, int n, int m, float radius, int nsample,
10 | at::Tensor new_xyz_tensor, at::Tensor xyz_tensor, at::Tensor idx_tensor);
11 |
12 | void ball_query_kernel_launcher_fast(int b, int n, int m, float radius, int nsample,
13 | const float *xyz, const float *new_xyz, int *idx);
14 |
15 | #endif
16 |
--------------------------------------------------------------------------------
/pcdet/ops/pointnet2/pointnet2_batch/src/cuda_utils.h:
--------------------------------------------------------------------------------
1 | #ifndef _CUDA_UTILS_H
2 | #define _CUDA_UTILS_H
3 |
4 | #include
5 |
6 | #define TOTAL_THREADS 1024
7 | #define THREADS_PER_BLOCK 256
8 | #define DIVUP(m,n) ((m) / (n) + ((m) % (n) > 0))
9 |
10 | inline int opt_n_threads(int work_size) {
11 | const int pow_2 = std::log(static_cast(work_size)) / std::log(2.0);
12 |
13 | return max(min(1 << pow_2, TOTAL_THREADS), 1);
14 | }
15 | #endif
16 |
--------------------------------------------------------------------------------
/pcdet/ops/pointnet2/pointnet2_batch/src/group_points.cpp:
--------------------------------------------------------------------------------
1 | /*
2 | batch version of point grouping, modified from the original implementation of official PointNet++ codes.
3 | Written by Shaoshuai Shi
4 | All Rights Reserved 2018.
5 | */
6 |
7 |
8 | #include
9 | #include
10 | #include
11 | #include
12 | #include
13 | #include "group_points_gpu.h"
14 |
15 | extern THCState *state;
16 |
17 |
18 | int group_points_grad_wrapper_fast(int b, int c, int n, int npoints, int nsample,
19 | at::Tensor grad_out_tensor, at::Tensor idx_tensor, at::Tensor grad_points_tensor) {
20 |
21 | float *grad_points = grad_points_tensor.data();
22 | const int *idx = idx_tensor.data();
23 | const float *grad_out = grad_out_tensor.data();
24 |
25 | group_points_grad_kernel_launcher_fast(b, c, n, npoints, nsample, grad_out, idx, grad_points);
26 | return 1;
27 | }
28 |
29 |
30 | int group_points_wrapper_fast(int b, int c, int n, int npoints, int nsample,
31 | at::Tensor points_tensor, at::Tensor idx_tensor, at::Tensor out_tensor) {
32 |
33 | const float *points = points_tensor.data();
34 | const int *idx = idx_tensor.data();
35 | float *out = out_tensor.data();
36 |
37 | group_points_kernel_launcher_fast(b, c, n, npoints, nsample, points, idx, out);
38 | return 1;
39 | }
40 |
--------------------------------------------------------------------------------
/pcdet/ops/pointnet2/pointnet2_batch/src/group_points_gpu.cu:
--------------------------------------------------------------------------------
1 | /*
2 | batch version of point grouping, modified from the original implementation of official PointNet++ codes.
3 | Written by Shaoshuai Shi
4 | All Rights Reserved 2018.
5 | */
6 |
7 | #include
8 | #include
9 |
10 | #include "cuda_utils.h"
11 | #include "group_points_gpu.h"
12 |
13 |
14 | __global__ void group_points_grad_kernel_fast(int b, int c, int n, int npoints, int nsample,
15 | const float *__restrict__ grad_out, const int *__restrict__ idx, float *__restrict__ grad_points) {
16 | // grad_out: (B, C, npoints, nsample)
17 | // idx: (B, npoints, nsample)
18 | // output:
19 | // grad_points: (B, C, N)
20 | int bs_idx = blockIdx.z;
21 | int c_idx = blockIdx.y;
22 | int index = blockIdx.x * blockDim.x + threadIdx.x;
23 | int pt_idx = index / nsample;
24 | if (bs_idx >= b || c_idx >= c || pt_idx >= npoints) return;
25 |
26 | int sample_idx = index % nsample;
27 | grad_out += bs_idx * c * npoints * nsample + c_idx * npoints * nsample + pt_idx * nsample + sample_idx;
28 | idx += bs_idx * npoints * nsample + pt_idx * nsample + sample_idx;
29 |
30 | atomicAdd(grad_points + bs_idx * c * n + c_idx * n + idx[0] , grad_out[0]);
31 | }
32 |
33 | void group_points_grad_kernel_launcher_fast(int b, int c, int n, int npoints, int nsample,
34 | const float *grad_out, const int *idx, float *grad_points) {
35 | // grad_out: (B, C, npoints, nsample)
36 | // idx: (B, npoints, nsample)
37 | // output:
38 | // grad_points: (B, C, N)
39 | cudaError_t err;
40 | dim3 blocks(DIVUP(npoints * nsample, THREADS_PER_BLOCK), c, b); // blockIdx.x(col), blockIdx.y(row)
41 | dim3 threads(THREADS_PER_BLOCK);
42 |
43 | group_points_grad_kernel_fast<<>>(b, c, n, npoints, nsample, grad_out, idx, grad_points);
44 |
45 | err = cudaGetLastError();
46 | if (cudaSuccess != err) {
47 | fprintf(stderr, "CUDA kernel failed : %s\n", cudaGetErrorString(err));
48 | exit(-1);
49 | }
50 | }
51 |
52 |
53 | __global__ void group_points_kernel_fast(int b, int c, int n, int npoints, int nsample,
54 | const float *__restrict__ points, const int *__restrict__ idx, float *__restrict__ out) {
55 | // points: (B, C, N)
56 | // idx: (B, npoints, nsample)
57 | // output:
58 | // out: (B, C, npoints, nsample)
59 | int bs_idx = blockIdx.z;
60 | int c_idx = blockIdx.y;
61 | int index = blockIdx.x * blockDim.x + threadIdx.x;
62 | int pt_idx = index / nsample;
63 | if (bs_idx >= b || c_idx >= c || pt_idx >= npoints) return;
64 |
65 | int sample_idx = index % nsample;
66 |
67 | idx += bs_idx * npoints * nsample + pt_idx * nsample + sample_idx;
68 | int in_idx = bs_idx * c * n + c_idx * n + idx[0];
69 | int out_idx = bs_idx * c * npoints * nsample + c_idx * npoints * nsample + pt_idx * nsample + sample_idx;
70 |
71 | out[out_idx] = points[in_idx];
72 | }
73 |
74 |
75 | void group_points_kernel_launcher_fast(int b, int c, int n, int npoints, int nsample,
76 | const float *points, const int *idx, float *out) {
77 | // points: (B, C, N)
78 | // idx: (B, npoints, nsample)
79 | // output:
80 | // out: (B, C, npoints, nsample)
81 | cudaError_t err;
82 | dim3 blocks(DIVUP(npoints * nsample, THREADS_PER_BLOCK), c, b); // blockIdx.x(col), blockIdx.y(row)
83 | dim3 threads(THREADS_PER_BLOCK);
84 |
85 | group_points_kernel_fast<<>>(b, c, n, npoints, nsample, points, idx, out);
86 | // cudaDeviceSynchronize(); // for using printf in kernel function
87 | err = cudaGetLastError();
88 | if (cudaSuccess != err) {
89 | fprintf(stderr, "CUDA kernel failed : %s\n", cudaGetErrorString(err));
90 | exit(-1);
91 | }
92 | }
93 |
--------------------------------------------------------------------------------
/pcdet/ops/pointnet2/pointnet2_batch/src/group_points_gpu.h:
--------------------------------------------------------------------------------
1 | #ifndef _GROUP_POINTS_GPU_H
2 | #define _GROUP_POINTS_GPU_H
3 |
4 | #include
5 | #include
6 | #include
7 | #include
8 |
9 |
10 | int group_points_wrapper_fast(int b, int c, int n, int npoints, int nsample,
11 | at::Tensor points_tensor, at::Tensor idx_tensor, at::Tensor out_tensor);
12 |
13 | void group_points_kernel_launcher_fast(int b, int c, int n, int npoints, int nsample,
14 | const float *points, const int *idx, float *out);
15 |
16 | int group_points_grad_wrapper_fast(int b, int c, int n, int npoints, int nsample,
17 | at::Tensor grad_out_tensor, at::Tensor idx_tensor, at::Tensor grad_points_tensor);
18 |
19 | void group_points_grad_kernel_launcher_fast(int b, int c, int n, int npoints, int nsample,
20 | const float *grad_out, const int *idx, float *grad_points);
21 |
22 | #endif
23 |
--------------------------------------------------------------------------------
/pcdet/ops/pointnet2/pointnet2_batch/src/interpolate.cpp:
--------------------------------------------------------------------------------
1 | /*
2 | batch version of point interpolation, modified from the original implementation of official PointNet++ codes.
3 | Written by Shaoshuai Shi
4 | All Rights Reserved 2018.
5 | */
6 |
7 |
8 | #include
9 | #include
10 | #include
11 | #include
12 | #include
13 | #include
14 | #include
15 | #include
16 | #include "interpolate_gpu.h"
17 |
18 | extern THCState *state;
19 |
20 |
21 | void three_nn_wrapper_fast(int b, int n, int m, at::Tensor unknown_tensor,
22 | at::Tensor known_tensor, at::Tensor dist2_tensor, at::Tensor idx_tensor) {
23 | const float *unknown = unknown_tensor.data();
24 | const float *known = known_tensor.data();
25 | float *dist2 = dist2_tensor.data();
26 | int *idx = idx_tensor.data();
27 |
28 | three_nn_kernel_launcher_fast(b, n, m, unknown, known, dist2, idx);
29 | }
30 |
31 |
32 | void three_interpolate_wrapper_fast(int b, int c, int m, int n,
33 | at::Tensor points_tensor,
34 | at::Tensor idx_tensor,
35 | at::Tensor weight_tensor,
36 | at::Tensor out_tensor) {
37 |
38 | const float *points = points_tensor.data();
39 | const float *weight = weight_tensor.data();
40 | float *out = out_tensor.data();
41 | const int *idx = idx_tensor.data();
42 |
43 | three_interpolate_kernel_launcher_fast(b, c, m, n, points, idx, weight, out);
44 | }
45 |
46 | void three_interpolate_grad_wrapper_fast(int b, int c, int n, int m,
47 | at::Tensor grad_out_tensor,
48 | at::Tensor idx_tensor,
49 | at::Tensor weight_tensor,
50 | at::Tensor grad_points_tensor) {
51 |
52 | const float *grad_out = grad_out_tensor.data();
53 | const float *weight = weight_tensor.data();
54 | float *grad_points = grad_points_tensor.data();
55 | const int *idx = idx_tensor.data();
56 |
57 | three_interpolate_grad_kernel_launcher_fast(b, c, n, m, grad_out, idx, weight, grad_points);
58 | }
59 |
--------------------------------------------------------------------------------
/pcdet/ops/pointnet2/pointnet2_batch/src/interpolate_gpu.h:
--------------------------------------------------------------------------------
1 | #ifndef _INTERPOLATE_GPU_H
2 | #define _INTERPOLATE_GPU_H
3 |
4 | #include
5 | #include
6 | #include
7 | #include
8 |
9 |
10 | void three_nn_wrapper_fast(int b, int n, int m, at::Tensor unknown_tensor,
11 | at::Tensor known_tensor, at::Tensor dist2_tensor, at::Tensor idx_tensor);
12 |
13 | void three_nn_kernel_launcher_fast(int b, int n, int m, const float *unknown,
14 | const float *known, float *dist2, int *idx);
15 |
16 |
17 | void three_interpolate_wrapper_fast(int b, int c, int m, int n, at::Tensor points_tensor,
18 | at::Tensor idx_tensor, at::Tensor weight_tensor, at::Tensor out_tensor);
19 |
20 | void three_interpolate_kernel_launcher_fast(int b, int c, int m, int n,
21 | const float *points, const int *idx, const float *weight, float *out);
22 |
23 |
24 | void three_interpolate_grad_wrapper_fast(int b, int c, int n, int m, at::Tensor grad_out_tensor,
25 | at::Tensor idx_tensor, at::Tensor weight_tensor, at::Tensor grad_points_tensor);
26 |
27 | void three_interpolate_grad_kernel_launcher_fast(int b, int c, int n, int m, const float *grad_out,
28 | const int *idx, const float *weight, float *grad_points);
29 |
30 | #endif
31 |
--------------------------------------------------------------------------------
/pcdet/ops/pointnet2/pointnet2_batch/src/pointnet2_api.cpp:
--------------------------------------------------------------------------------
1 | #include
2 | #include
3 |
4 | #include "ball_query_gpu.h"
5 | #include "group_points_gpu.h"
6 | #include "sampling_gpu.h"
7 | #include "interpolate_gpu.h"
8 |
9 |
10 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
11 | m.def("ball_query_wrapper", &ball_query_wrapper_fast, "ball_query_wrapper_fast");
12 |
13 | m.def("group_points_wrapper", &group_points_wrapper_fast, "group_points_wrapper_fast");
14 | m.def("group_points_grad_wrapper", &group_points_grad_wrapper_fast, "group_points_grad_wrapper_fast");
15 |
16 | m.def("gather_points_wrapper", &gather_points_wrapper_fast, "gather_points_wrapper_fast");
17 | m.def("gather_points_grad_wrapper", &gather_points_grad_wrapper_fast, "gather_points_grad_wrapper_fast");
18 |
19 | m.def("farthest_point_sampling_wrapper", &farthest_point_sampling_wrapper, "farthest_point_sampling_wrapper");
20 |
21 | m.def("three_nn_wrapper", &three_nn_wrapper_fast, "three_nn_wrapper_fast");
22 | m.def("three_interpolate_wrapper", &three_interpolate_wrapper_fast, "three_interpolate_wrapper_fast");
23 | m.def("three_interpolate_grad_wrapper", &three_interpolate_grad_wrapper_fast, "three_interpolate_grad_wrapper_fast");
24 | }
25 |
--------------------------------------------------------------------------------
/pcdet/ops/pointnet2/pointnet2_batch/src/sampling.cpp:
--------------------------------------------------------------------------------
1 | /*
2 | batch version of point sampling and gathering, modified from the original implementation of official PointNet++ codes.
3 | Written by Shaoshuai Shi
4 | All Rights Reserved 2018.
5 | */
6 |
7 |
8 | #include
9 | #include
10 | #include
11 | #include
12 |
13 | #include "sampling_gpu.h"
14 |
15 | extern THCState *state;
16 |
17 |
18 | int gather_points_wrapper_fast(int b, int c, int n, int npoints,
19 | at::Tensor points_tensor, at::Tensor idx_tensor, at::Tensor out_tensor){
20 | const float *points = points_tensor.data();
21 | const int *idx = idx_tensor.data();
22 | float *out = out_tensor.data();
23 |
24 | gather_points_kernel_launcher_fast(b, c, n, npoints, points, idx, out);
25 | return 1;
26 | }
27 |
28 |
29 | int gather_points_grad_wrapper_fast(int b, int c, int n, int npoints,
30 | at::Tensor grad_out_tensor, at::Tensor idx_tensor, at::Tensor grad_points_tensor) {
31 |
32 | const float *grad_out = grad_out_tensor.data();
33 | const int *idx = idx_tensor.data();
34 | float *grad_points = grad_points_tensor.data();
35 |
36 | gather_points_grad_kernel_launcher_fast(b, c, n, npoints, grad_out, idx, grad_points);
37 | return 1;
38 | }
39 |
40 |
41 | int farthest_point_sampling_wrapper(int b, int n, int m,
42 | at::Tensor points_tensor, at::Tensor temp_tensor, at::Tensor idx_tensor) {
43 |
44 | const float *points = points_tensor.data();
45 | float *temp = temp_tensor.data();
46 | int *idx = idx_tensor.data();
47 |
48 | farthest_point_sampling_kernel_launcher(b, n, m, points, temp, idx);
49 | return 1;
50 | }
51 |
--------------------------------------------------------------------------------
/pcdet/ops/pointnet2/pointnet2_batch/src/sampling_gpu.h:
--------------------------------------------------------------------------------
1 | #ifndef _SAMPLING_GPU_H
2 | #define _SAMPLING_GPU_H
3 |
4 | #include
5 | #include
6 | #include
7 |
8 |
9 | int gather_points_wrapper_fast(int b, int c, int n, int npoints,
10 | at::Tensor points_tensor, at::Tensor idx_tensor, at::Tensor out_tensor);
11 |
12 | void gather_points_kernel_launcher_fast(int b, int c, int n, int npoints,
13 | const float *points, const int *idx, float *out);
14 |
15 |
16 | int gather_points_grad_wrapper_fast(int b, int c, int n, int npoints,
17 | at::Tensor grad_out_tensor, at::Tensor idx_tensor, at::Tensor grad_points_tensor);
18 |
19 | void gather_points_grad_kernel_launcher_fast(int b, int c, int n, int npoints,
20 | const float *grad_out, const int *idx, float *grad_points);
21 |
22 |
23 | int farthest_point_sampling_wrapper(int b, int n, int m,
24 | at::Tensor points_tensor, at::Tensor temp_tensor, at::Tensor idx_tensor);
25 |
26 | void farthest_point_sampling_kernel_launcher(int b, int n, int m,
27 | const float *dataset, float *temp, int *idxs);
28 |
29 | #endif
30 |
--------------------------------------------------------------------------------
/pcdet/ops/pointnet2/pointnet2_stack/src/ball_query.cpp:
--------------------------------------------------------------------------------
1 | /*
2 | Stacked-batch-data version of ball query, modified from the original implementation of official PointNet++ codes.
3 | Written by Shaoshuai Shi
4 | All Rights Reserved 2019-2020.
5 | */
6 |
7 |
8 | #include
9 | #include
10 | #include
11 | #include
12 | #include
13 | #include "ball_query_gpu.h"
14 |
15 | extern THCState *state;
16 |
17 | #define CHECK_CUDA(x) do { \
18 | if (!x.type().is_cuda()) { \
19 | fprintf(stderr, "%s must be CUDA tensor at %s:%d\n", #x, __FILE__, __LINE__); \
20 | exit(-1); \
21 | } \
22 | } while (0)
23 | #define CHECK_CONTIGUOUS(x) do { \
24 | if (!x.is_contiguous()) { \
25 | fprintf(stderr, "%s must be contiguous tensor at %s:%d\n", #x, __FILE__, __LINE__); \
26 | exit(-1); \
27 | } \
28 | } while (0)
29 | #define CHECK_INPUT(x) CHECK_CUDA(x);CHECK_CONTIGUOUS(x)
30 |
31 | int ball_query_wrapper_stack(int B, int M, float radius, int nsample,
32 | at::Tensor new_xyz_tensor, at::Tensor new_xyz_batch_cnt_tensor,
33 | at::Tensor xyz_tensor, at::Tensor xyz_batch_cnt_tensor, at::Tensor idx_tensor) {
34 | CHECK_INPUT(new_xyz_tensor);
35 | CHECK_INPUT(xyz_tensor);
36 | CHECK_INPUT(new_xyz_batch_cnt_tensor);
37 | CHECK_INPUT(xyz_batch_cnt_tensor);
38 |
39 | const float *new_xyz = new_xyz_tensor.data();
40 | const float *xyz = xyz_tensor.data();
41 | const int *new_xyz_batch_cnt = new_xyz_batch_cnt_tensor.data();
42 | const int *xyz_batch_cnt = xyz_batch_cnt_tensor.data();
43 | int *idx = idx_tensor.data();
44 |
45 | ball_query_kernel_launcher_stack(B, M, radius, nsample, new_xyz, new_xyz_batch_cnt, xyz, xyz_batch_cnt, idx);
46 | return 1;
47 | }
48 |
--------------------------------------------------------------------------------
/pcdet/ops/pointnet2/pointnet2_stack/src/ball_query_gpu.cu:
--------------------------------------------------------------------------------
1 | /*
2 | Stacked-batch-data version of ball query, modified from the original implementation of official PointNet++ codes.
3 | Written by Shaoshuai Shi
4 | All Rights Reserved 2019-2020.
5 | */
6 |
7 |
8 | #include
9 | #include
10 | #include
11 |
12 | #include "ball_query_gpu.h"
13 | #include "cuda_utils.h"
14 |
15 |
16 | __global__ void ball_query_kernel_stack(int B, int M, float radius, int nsample, \
17 | const float *new_xyz, const int *new_xyz_batch_cnt, const float *xyz, const int *xyz_batch_cnt, int *idx) {
18 | // :param xyz: (N1 + N2 ..., 3) xyz coordinates of the features
19 | // :param xyz_batch_cnt: (batch_size), [N1, N2, ...]
20 | // :param new_xyz: (M1 + M2 ..., 3) centers of the ball query
21 | // :param new_xyz_batch_cnt: (batch_size), [M1, M2, ...]
22 | // output:
23 | // idx: (M, nsample)
24 | int pt_idx = blockIdx.x * blockDim.x + threadIdx.x;
25 | if (pt_idx >= M) return;
26 |
27 | int bs_idx = 0, pt_cnt = new_xyz_batch_cnt[0];
28 | for (int k = 1; k < B; k++){
29 | if (pt_idx < pt_cnt) break;
30 | pt_cnt += new_xyz_batch_cnt[k];
31 | bs_idx = k;
32 | }
33 |
34 | int xyz_batch_start_idx = 0;
35 | for (int k = 0; k < bs_idx; k++) xyz_batch_start_idx += xyz_batch_cnt[k];
36 | // for (int k = 0; k < bs_idx; k++) new_xyz_batch_start_idx += new_xyz_batch_cnt[k];
37 |
38 | new_xyz += pt_idx * 3;
39 | xyz += xyz_batch_start_idx * 3;
40 | idx += pt_idx * nsample;
41 |
42 | float radius2 = radius * radius;
43 | float new_x = new_xyz[0];
44 | float new_y = new_xyz[1];
45 | float new_z = new_xyz[2];
46 | int n = xyz_batch_cnt[bs_idx];
47 |
48 | int cnt = 0;
49 | for (int k = 0; k < n; ++k) {
50 | float x = xyz[k * 3 + 0];
51 | float y = xyz[k * 3 + 1];
52 | float z = xyz[k * 3 + 2];
53 | float d2 = (new_x - x) * (new_x - x) + (new_y - y) * (new_y - y) + (new_z - z) * (new_z - z);
54 | if (d2 < radius2){
55 | if (cnt == 0){
56 | for (int l = 0; l < nsample; ++l) {
57 | idx[l] = k;
58 | }
59 | }
60 | idx[cnt] = k;
61 | ++cnt;
62 | if (cnt >= nsample) break;
63 | }
64 | }
65 | if (cnt == 0) idx[0] = -1;
66 | }
67 |
68 |
69 | void ball_query_kernel_launcher_stack(int B, int M, float radius, int nsample,
70 | const float *new_xyz, const int *new_xyz_batch_cnt, const float *xyz, const int *xyz_batch_cnt, int *idx){
71 | // :param xyz: (N1 + N2 ..., 3) xyz coordinates of the features
72 | // :param xyz_batch_cnt: (batch_size), [N1, N2, ...]
73 | // :param new_xyz: (M1 + M2 ..., 3) centers of the ball query
74 | // :param new_xyz_batch_cnt: (batch_size), [M1, M2, ...]
75 | // output:
76 | // idx: (M, nsample)
77 |
78 | cudaError_t err;
79 |
80 | dim3 blocks(DIVUP(M, THREADS_PER_BLOCK)); // blockIdx.x(col), blockIdx.y(row)
81 | dim3 threads(THREADS_PER_BLOCK);
82 |
83 | ball_query_kernel_stack<<>>(B, M, radius, nsample, new_xyz, new_xyz_batch_cnt, xyz, xyz_batch_cnt, idx);
84 | // cudaDeviceSynchronize(); // for using printf in kernel function
85 | err = cudaGetLastError();
86 | if (cudaSuccess != err) {
87 | fprintf(stderr, "CUDA kernel failed : %s\n", cudaGetErrorString(err));
88 | exit(-1);
89 | }
90 | }
91 |
--------------------------------------------------------------------------------
/pcdet/ops/pointnet2/pointnet2_stack/src/ball_query_gpu.h:
--------------------------------------------------------------------------------
1 | /*
2 | Stacked-batch-data version of ball query, modified from the original implementation of official PointNet++ codes.
3 | Written by Shaoshuai Shi
4 | All Rights Reserved 2019-2020.
5 | */
6 |
7 |
8 | #ifndef _STACK_BALL_QUERY_GPU_H
9 | #define _STACK_BALL_QUERY_GPU_H
10 |
11 | #include
12 | #include
13 | #include
14 | #include
15 |
16 | int ball_query_wrapper_stack(int B, int M, float radius, int nsample,
17 | at::Tensor new_xyz_tensor, at::Tensor new_xyz_batch_cnt_tensor,
18 | at::Tensor xyz_tensor, at::Tensor xyz_batch_cnt_tensor, at::Tensor idx_tensor);
19 |
20 |
21 | void ball_query_kernel_launcher_stack(int B, int M, float radius, int nsample,
22 | const float *new_xyz, const int *new_xyz_batch_cnt, const float *xyz, const int *xyz_batch_cnt, int *idx);
23 |
24 |
25 | #endif
26 |
--------------------------------------------------------------------------------
/pcdet/ops/pointnet2/pointnet2_stack/src/cuda_utils.h:
--------------------------------------------------------------------------------
1 | #ifndef _STACK_CUDA_UTILS_H
2 | #define _STACK_CUDA_UTILS_H
3 |
4 | #include
5 |
6 | #define THREADS_PER_BLOCK 256
7 | #define DIVUP(m,n) ((m) / (n) + ((m) % (n) > 0))
8 |
9 | #endif
10 |
--------------------------------------------------------------------------------
/pcdet/ops/pointnet2/pointnet2_stack/src/group_points.cpp:
--------------------------------------------------------------------------------
1 | /*
2 | Stacked-batch-data version of point grouping, modified from the original implementation of official PointNet++ codes.
3 | Written by Shaoshuai Shi
4 | All Rights Reserved 2019-2020.
5 | */
6 |
7 |
8 | #include
9 | #include
10 | #include
11 | #include
12 | #include
13 | #include "group_points_gpu.h"
14 |
15 | extern THCState *state;
16 | #define CHECK_CUDA(x) do { \
17 | if (!x.type().is_cuda()) { \
18 | fprintf(stderr, "%s must be CUDA tensor at %s:%d\n", #x, __FILE__, __LINE__); \
19 | exit(-1); \
20 | } \
21 | } while (0)
22 | #define CHECK_CONTIGUOUS(x) do { \
23 | if (!x.is_contiguous()) { \
24 | fprintf(stderr, "%s must be contiguous tensor at %s:%d\n", #x, __FILE__, __LINE__); \
25 | exit(-1); \
26 | } \
27 | } while (0)
28 | #define CHECK_INPUT(x) CHECK_CUDA(x);CHECK_CONTIGUOUS(x)
29 |
30 |
31 | int group_points_grad_wrapper_stack(int B, int M, int C, int N, int nsample,
32 | at::Tensor grad_out_tensor, at::Tensor idx_tensor, at::Tensor idx_batch_cnt_tensor,
33 | at::Tensor features_batch_cnt_tensor, at::Tensor grad_features_tensor) {
34 |
35 | CHECK_INPUT(grad_out_tensor);
36 | CHECK_INPUT(idx_tensor);
37 | CHECK_INPUT(idx_batch_cnt_tensor);
38 | CHECK_INPUT(features_batch_cnt_tensor);
39 | CHECK_INPUT(grad_features_tensor);
40 |
41 | const float *grad_out = grad_out_tensor.data();
42 | const int *idx = idx_tensor.data();
43 | const int *idx_batch_cnt = idx_batch_cnt_tensor.data();
44 | const int *features_batch_cnt = features_batch_cnt_tensor.data();
45 | float *grad_features = grad_features_tensor.data();
46 |
47 | group_points_grad_kernel_launcher_stack(B, M, C, N, nsample, grad_out, idx, idx_batch_cnt, features_batch_cnt, grad_features);
48 | return 1;
49 | }
50 |
51 |
52 | int group_points_wrapper_stack(int B, int M, int C, int nsample,
53 | at::Tensor features_tensor, at::Tensor features_batch_cnt_tensor,
54 | at::Tensor idx_tensor, at::Tensor idx_batch_cnt_tensor, at::Tensor out_tensor) {
55 |
56 | CHECK_INPUT(features_tensor);
57 | CHECK_INPUT(features_batch_cnt_tensor);
58 | CHECK_INPUT(idx_tensor);
59 | CHECK_INPUT(idx_batch_cnt_tensor);
60 | CHECK_INPUT(out_tensor);
61 |
62 | const float *features = features_tensor.data();
63 | const int *idx = idx_tensor.data();
64 | const int *features_batch_cnt = features_batch_cnt_tensor.data();
65 | const int *idx_batch_cnt = idx_batch_cnt_tensor.data();
66 | float *out = out_tensor.data();
67 |
68 | group_points_kernel_launcher_stack(B, M, C, nsample, features, features_batch_cnt, idx, idx_batch_cnt, out);
69 | return 1;
70 | }
--------------------------------------------------------------------------------
/pcdet/ops/pointnet2/pointnet2_stack/src/group_points_gpu.h:
--------------------------------------------------------------------------------
1 | /*
2 | Stacked-batch-data version of point grouping, modified from the original implementation of official PointNet++ codes.
3 | Written by Shaoshuai Shi
4 | All Rights Reserved 2019-2020.
5 | */
6 |
7 |
8 | #ifndef _STACK_GROUP_POINTS_GPU_H
9 | #define _STACK_GROUP_POINTS_GPU_H
10 |
11 | #include
12 | #include
13 | #include
14 | #include
15 |
16 |
17 | int group_points_wrapper_stack(int B, int M, int C, int nsample,
18 | at::Tensor features_tensor, at::Tensor features_batch_cnt_tensor,
19 | at::Tensor idx_tensor, at::Tensor idx_batch_cnt_tensor, at::Tensor out_tensor);
20 |
21 | void group_points_kernel_launcher_stack(int B, int M, int C, int nsample,
22 | const float *features, const int *features_batch_cnt, const int *idx, const int *idx_batch_cnt, float *out);
23 |
24 | int group_points_grad_wrapper_stack(int B, int M, int C, int N, int nsample,
25 | at::Tensor grad_out_tensor, at::Tensor idx_tensor, at::Tensor idx_batch_cnt_tensor,
26 | at::Tensor features_batch_cnt_tensor, at::Tensor grad_features_tensor);
27 |
28 | void group_points_grad_kernel_launcher_stack(int B, int M, int C, int N, int nsample,
29 | const float *grad_out, const int *idx, const int *idx_batch_cnt, const int *features_batch_cnt, float *grad_features);
30 |
31 | #endif
32 |
--------------------------------------------------------------------------------
/pcdet/ops/pointnet2/pointnet2_stack/src/interpolate_gpu.h:
--------------------------------------------------------------------------------
1 | #ifndef _INTERPOLATE_GPU_H
2 | #define _INTERPOLATE_GPU_H
3 |
4 | #include
5 | #include
6 | #include
7 | #include
8 |
9 |
10 | void three_nn_wrapper_stack(at::Tensor unknown_tensor,
11 | at::Tensor unknown_batch_cnt_tensor, at::Tensor known_tensor,
12 | at::Tensor known_batch_cnt_tensor, at::Tensor dist2_tensor, at::Tensor idx_tensor);
13 |
14 |
15 | void three_interpolate_wrapper_stack(at::Tensor features_tensor,
16 | at::Tensor idx_tensor, at::Tensor weight_tensor, at::Tensor out_tensor);
17 |
18 |
19 |
20 | void three_interpolate_grad_wrapper_stack(at::Tensor grad_out_tensor, at::Tensor idx_tensor,
21 | at::Tensor weight_tensor, at::Tensor grad_features_tensor);
22 |
23 |
24 | void three_nn_kernel_launcher_stack(int batch_size, int N, int M, const float *unknown,
25 | const int *unknown_batch_cnt, const float *known, const int *known_batch_cnt,
26 | float *dist2, int *idx);
27 |
28 |
29 | void three_interpolate_kernel_launcher_stack(int N, int channels,
30 | const float *features, const int *idx, const float *weight, float *out);
31 |
32 |
33 |
34 | void three_interpolate_grad_kernel_launcher_stack(int N, int channels, const float *grad_out,
35 | const int *idx, const float *weight, float *grad_features);
36 |
37 |
38 |
39 | #endif
--------------------------------------------------------------------------------
/pcdet/ops/pointnet2/pointnet2_stack/src/pointnet2_api.cpp:
--------------------------------------------------------------------------------
1 | #include
2 | #include
3 |
4 | #include "ball_query_gpu.h"
5 | #include "group_points_gpu.h"
6 | #include "sampling_gpu.h"
7 | #include "interpolate_gpu.h"
8 | #include "voxel_query_gpu.h"
9 | #include "vector_pool_gpu.h"
10 |
11 |
12 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
13 | m.def("ball_query_wrapper", &ball_query_wrapper_stack, "ball_query_wrapper_stack");
14 | m.def("voxel_query_wrapper", &voxel_query_wrapper_stack, "voxel_query_wrapper_stack");
15 |
16 | m.def("farthest_point_sampling_wrapper", &farthest_point_sampling_wrapper, "farthest_point_sampling_wrapper");
17 | m.def("stack_farthest_point_sampling_wrapper", &stack_farthest_point_sampling_wrapper, "stack_farthest_point_sampling_wrapper");
18 |
19 | m.def("group_points_wrapper", &group_points_wrapper_stack, "group_points_wrapper_stack");
20 | m.def("group_points_grad_wrapper", &group_points_grad_wrapper_stack, "group_points_grad_wrapper_stack");
21 |
22 | m.def("three_nn_wrapper", &three_nn_wrapper_stack, "three_nn_wrapper_stack");
23 | m.def("three_interpolate_wrapper", &three_interpolate_wrapper_stack, "three_interpolate_wrapper_stack");
24 | m.def("three_interpolate_grad_wrapper", &three_interpolate_grad_wrapper_stack, "three_interpolate_grad_wrapper_stack");
25 |
26 | m.def("query_stacked_local_neighbor_idxs_wrapper_stack", &query_stacked_local_neighbor_idxs_wrapper_stack, "query_stacked_local_neighbor_idxs_wrapper_stack");
27 | m.def("query_three_nn_by_stacked_local_idxs_wrapper_stack", &query_three_nn_by_stacked_local_idxs_wrapper_stack, "query_three_nn_by_stacked_local_idxs_wrapper_stack");
28 |
29 | m.def("vector_pool_wrapper", &vector_pool_wrapper_stack, "vector_pool_grad_wrapper_stack");
30 | m.def("vector_pool_grad_wrapper", &vector_pool_grad_wrapper_stack, "vector_pool_grad_wrapper_stack");
31 | }
32 |
--------------------------------------------------------------------------------
/pcdet/ops/pointnet2/pointnet2_stack/src/sampling.cpp:
--------------------------------------------------------------------------------
1 | #include
2 | #include
3 | #include
4 | #include
5 |
6 | #include "sampling_gpu.h"
7 |
8 | extern THCState *state;
9 | #define CHECK_CUDA(x) do { \
10 | if (!x.type().is_cuda()) { \
11 | fprintf(stderr, "%s must be CUDA tensor at %s:%d\n", #x, __FILE__, __LINE__); \
12 | exit(-1); \
13 | } \
14 | } while (0)
15 | #define CHECK_CONTIGUOUS(x) do { \
16 | if (!x.is_contiguous()) { \
17 | fprintf(stderr, "%s must be contiguous tensor at %s:%d\n", #x, __FILE__, __LINE__); \
18 | exit(-1); \
19 | } \
20 | } while (0)
21 | #define CHECK_INPUT(x) CHECK_CUDA(x);CHECK_CONTIGUOUS(x)
22 |
23 |
24 | int farthest_point_sampling_wrapper(int b, int n, int m,
25 | at::Tensor points_tensor, at::Tensor temp_tensor, at::Tensor idx_tensor) {
26 |
27 | CHECK_INPUT(points_tensor);
28 | CHECK_INPUT(temp_tensor);
29 | CHECK_INPUT(idx_tensor);
30 |
31 | const float *points = points_tensor.data();
32 | float *temp = temp_tensor.data();
33 | int *idx = idx_tensor.data();
34 |
35 | farthest_point_sampling_kernel_launcher(b, n, m, points, temp, idx);
36 | return 1;
37 | }
38 |
39 |
40 | int stack_farthest_point_sampling_wrapper(at::Tensor points_tensor,
41 | at::Tensor temp_tensor, at::Tensor xyz_batch_cnt_tensor, at::Tensor idx_tensor,
42 | at::Tensor num_sampled_points_tensor) {
43 |
44 | CHECK_INPUT(points_tensor);
45 | CHECK_INPUT(temp_tensor);
46 | CHECK_INPUT(idx_tensor);
47 | CHECK_INPUT(xyz_batch_cnt_tensor);
48 | CHECK_INPUT(num_sampled_points_tensor);
49 |
50 | int batch_size = xyz_batch_cnt_tensor.size(0);
51 | int N = points_tensor.size(0);
52 | const float *points = points_tensor.data();
53 | float *temp = temp_tensor.data();
54 | int *xyz_batch_cnt = xyz_batch_cnt_tensor.data();
55 | int *idx = idx_tensor.data();
56 | int *num_sampled_points = num_sampled_points_tensor.data();
57 |
58 | stack_farthest_point_sampling_kernel_launcher(N, batch_size, points, temp, xyz_batch_cnt, idx, num_sampled_points);
59 | return 1;
60 | }
--------------------------------------------------------------------------------
/pcdet/ops/pointnet2/pointnet2_stack/src/sampling_gpu.h:
--------------------------------------------------------------------------------
1 | #ifndef _SAMPLING_GPU_H
2 | #define _SAMPLING_GPU_H
3 |
4 | #include
5 | #include
6 | #include
7 |
8 |
9 | int farthest_point_sampling_wrapper(int b, int n, int m,
10 | at::Tensor points_tensor, at::Tensor temp_tensor, at::Tensor idx_tensor);
11 |
12 | void farthest_point_sampling_kernel_launcher(int b, int n, int m,
13 | const float *dataset, float *temp, int *idxs);
14 |
15 | int stack_farthest_point_sampling_wrapper(
16 | at::Tensor points_tensor, at::Tensor temp_tensor, at::Tensor xyz_batch_cnt_tensor,
17 | at::Tensor idx_tensor, at::Tensor num_sampled_points_tensor);
18 |
19 |
20 | void stack_farthest_point_sampling_kernel_launcher(int N, int batch_size,
21 | const float *dataset, float *temp, int *xyz_batch_cnt, int *idxs, int *num_sampled_points);
22 |
23 | #endif
24 |
--------------------------------------------------------------------------------
/pcdet/ops/pointnet2/pointnet2_stack/src/vector_pool_gpu.h:
--------------------------------------------------------------------------------
1 | /*
2 | Vector-pool aggregation based local feature aggregation for point cloud.
3 | PV-RCNN++: Point-Voxel Feature Set Abstraction With Local Vector Representation for 3D Object Detection
4 | https://arxiv.org/abs/2102.00463
5 |
6 | Written by Shaoshuai Shi
7 | All Rights Reserved 2020.
8 | */
9 |
10 |
11 | #ifndef _STACK_VECTOR_POOL_GPU_H
12 | #define _STACK_VECTOR_POOL_GPU_H
13 |
14 | #include
15 | #include
16 | #include
17 | #include
18 |
19 |
20 | int query_stacked_local_neighbor_idxs_kernel_launcher_stack(
21 | const float *support_xyz, const int *xyz_batch_cnt, const float *new_xyz, const int *new_xyz_batch_cnt,
22 | int *stack_neighbor_idxs, int *start_len, int *cumsum, int avg_length_of_neighbor_idxs,
23 | float max_neighbour_distance, int batch_size, int M, int nsample, int neighbor_type);
24 |
25 | int query_stacked_local_neighbor_idxs_wrapper_stack(at::Tensor support_xyz_tensor, at::Tensor xyz_batch_cnt_tensor,
26 | at::Tensor new_xyz_tensor, at::Tensor new_xyz_batch_cnt_tensor,
27 | at::Tensor stack_neighbor_idxs_tensor, at::Tensor start_len_tensor, at::Tensor cumsum_tensor,
28 | int avg_length_of_neighbor_idxs, float max_neighbour_distance, int nsample, int neighbor_type);
29 |
30 |
31 | int query_three_nn_by_stacked_local_idxs_kernel_launcher_stack(
32 | const float *support_xyz, const float *new_xyz, const float *new_xyz_grid_centers,
33 | int *new_xyz_grid_idxs, float *new_xyz_grid_dist2,
34 | const int *stack_neighbor_idxs, const int *start_len,
35 | int M, int num_total_grids);
36 |
37 | int query_three_nn_by_stacked_local_idxs_wrapper_stack(at::Tensor support_xyz_tensor,
38 | at::Tensor new_xyz_tensor, at::Tensor new_xyz_grid_centers_tensor,
39 | at::Tensor new_xyz_grid_idxs_tensor, at::Tensor new_xyz_grid_dist2_tensor,
40 | at::Tensor stack_neighbor_idxs_tensor, at::Tensor start_len_tensor,
41 | int M, int num_total_grids);
42 |
43 |
44 | int vector_pool_wrapper_stack(at::Tensor support_xyz_tensor, at::Tensor xyz_batch_cnt_tensor,
45 | at::Tensor support_features_tensor, at::Tensor new_xyz_tensor, at::Tensor new_xyz_batch_cnt_tensor,
46 | at::Tensor new_features_tensor, at::Tensor new_local_xyz,
47 | at::Tensor point_cnt_of_grid_tensor, at::Tensor grouped_idxs_tensor,
48 | int num_grid_x, int num_grid_y, int num_grid_z, float max_neighbour_distance, int use_xyz,
49 | int num_max_sum_points, int nsample, int neighbor_type, int pooling_type);
50 |
51 |
52 | int vector_pool_kernel_launcher_stack(
53 | const float *support_xyz, const float *support_features, const int *xyz_batch_cnt,
54 | const float *new_xyz, float *new_features, float * new_local_xyz, const int *new_xyz_batch_cnt,
55 | int *point_cnt_of_grid, int *grouped_idxs,
56 | int num_grid_x, int num_grid_y, int num_grid_z, float max_neighbour_distance,
57 | int batch_size, int N, int M, int num_c_in, int num_c_out, int num_total_grids, int use_xyz,
58 | int num_max_sum_points, int nsample, int neighbor_type, int pooling_type);
59 |
60 |
61 | int vector_pool_grad_wrapper_stack(at::Tensor grad_new_features_tensor,
62 | at::Tensor point_cnt_of_grid_tensor, at::Tensor grouped_idxs_tensor,
63 | at::Tensor grad_support_features_tensor);
64 |
65 |
66 | void vector_pool_grad_kernel_launcher_stack(
67 | const float *grad_new_features, const int *point_cnt_of_grid, const int *grouped_idxs,
68 | float *grad_support_features, int N, int M, int num_c_out, int num_c_in, int num_total_grids,
69 | int num_max_sum_points);
70 |
71 | #endif
72 |
--------------------------------------------------------------------------------
/pcdet/ops/pointnet2/pointnet2_stack/src/voxel_query.cpp:
--------------------------------------------------------------------------------
1 | #include
2 | #include
3 | #include
4 | #include
5 | #include
6 | #include
7 | #include
8 | #include
9 | #include "voxel_query_gpu.h"
10 |
11 | extern THCState *state;
12 |
13 | #define CHECK_CUDA(x) do { \
14 | if (!x.type().is_cuda()) { \
15 | fprintf(stderr, "%s must be CUDA tensor at %s:%d\n", #x, __FILE__, __LINE__); \
16 | exit(-1); \
17 | } \
18 | } while (0)
19 | #define CHECK_CONTIGUOUS(x) do { \
20 | if (!x.is_contiguous()) { \
21 | fprintf(stderr, "%s must be contiguous tensor at %s:%d\n", #x, __FILE__, __LINE__); \
22 | exit(-1); \
23 | } \
24 | } while (0)
25 | #define CHECK_INPUT(x) CHECK_CUDA(x);CHECK_CONTIGUOUS(x)
26 |
27 |
28 | int voxel_query_wrapper_stack(int M, int R1, int R2, int R3, int nsample, float radius,
29 | int z_range, int y_range, int x_range, at::Tensor new_xyz_tensor, at::Tensor xyz_tensor,
30 | at::Tensor new_coords_tensor, at::Tensor point_indices_tensor, at::Tensor idx_tensor) {
31 | CHECK_INPUT(new_coords_tensor);
32 | CHECK_INPUT(point_indices_tensor);
33 | CHECK_INPUT(new_xyz_tensor);
34 | CHECK_INPUT(xyz_tensor);
35 |
36 | const float *new_xyz = new_xyz_tensor.data();
37 | const float *xyz = xyz_tensor.data();
38 | const int *new_coords = new_coords_tensor.data();
39 | const int *point_indices = point_indices_tensor.data();
40 | int *idx = idx_tensor.data();
41 |
42 | voxel_query_kernel_launcher_stack(M, R1, R2, R3, nsample, radius, z_range, y_range, x_range, new_xyz, xyz, new_coords, point_indices, idx);
43 | return 1;
44 | }
45 |
--------------------------------------------------------------------------------
/pcdet/ops/pointnet2/pointnet2_stack/src/voxel_query_gpu.h:
--------------------------------------------------------------------------------
1 | #ifndef _STACK_VOXEL_QUERY_GPU_H
2 | #define _STACK_VOXEL_QUERY_GPU_H
3 |
4 | #include
5 | #include
6 | #include
7 | #include
8 |
9 | int voxel_query_wrapper_stack(int M, int R1, int R2, int R3, int nsample, float radius,
10 | int z_range, int y_range, int x_range, at::Tensor new_xyz_tensor, at::Tensor xyz_tensor,
11 | at::Tensor new_coords_tensor, at::Tensor point_indices_tensor, at::Tensor idx_tensor);
12 |
13 |
14 | void voxel_query_kernel_launcher_stack(int M, int R1, int R2, int R3, int nsample,
15 | float radius, int z_range, int y_range, int x_range, const float *new_xyz,
16 | const float *xyz, const int *new_coords, const int *point_indices, int *idx);
17 |
18 |
19 | #endif
20 |
--------------------------------------------------------------------------------
/pcdet/ops/roipoint_pool3d/roipoint_pool3d_utils.py:
--------------------------------------------------------------------------------
1 | import torch
2 | import torch.nn as nn
3 | from torch.autograd import Function
4 |
5 | from ...utils import box_utils
6 | from . import roipoint_pool3d_cuda
7 |
8 |
9 | class RoIPointPool3d(nn.Module):
10 | def __init__(self, num_sampled_points=512, pool_extra_width=1.0):
11 | super().__init__()
12 | self.num_sampled_points = num_sampled_points
13 | self.pool_extra_width = pool_extra_width
14 |
15 | def forward(self, points, point_features, boxes3d):
16 | """
17 | Args:
18 | points: (B, N, 3)
19 | point_features: (B, N, C)
20 | boxes3d: (B, M, 7), [x, y, z, dx, dy, dz, heading]
21 |
22 | Returns:
23 | pooled_features: (B, M, 512, 3 + C)
24 | pooled_empty_flag: (B, M)
25 | """
26 | return RoIPointPool3dFunction.apply(
27 | points, point_features, boxes3d, self.pool_extra_width, self.num_sampled_points
28 | )
29 |
30 |
31 | class RoIPointPool3dFunction(Function):
32 | @staticmethod
33 | def forward(ctx, points, point_features, boxes3d, pool_extra_width, num_sampled_points=512):
34 | """
35 | Args:
36 | ctx:
37 | points: (B, N, 3)
38 | point_features: (B, N, C)
39 | boxes3d: (B, num_boxes, 7), [x, y, z, dx, dy, dz, heading]
40 | pool_extra_width:
41 | num_sampled_points:
42 |
43 | Returns:
44 | pooled_features: (B, num_boxes, 512, 3 + C)
45 | pooled_empty_flag: (B, num_boxes)
46 | """
47 | assert points.shape.__len__() == 3 and points.shape[2] == 3
48 | batch_size, boxes_num, feature_len = points.shape[0], boxes3d.shape[1], point_features.shape[2]
49 | pooled_boxes3d = box_utils.enlarge_box3d(boxes3d.view(-1, 7), pool_extra_width).view(batch_size, -1, 7)
50 |
51 | pooled_features = point_features.new_zeros((batch_size, boxes_num, num_sampled_points, 3 + feature_len))
52 | pooled_empty_flag = point_features.new_zeros((batch_size, boxes_num)).int()
53 |
54 | roipoint_pool3d_cuda.forward(
55 | points.contiguous(), pooled_boxes3d.contiguous(),
56 | point_features.contiguous(), pooled_features, pooled_empty_flag
57 | )
58 |
59 | return pooled_features, pooled_empty_flag
60 |
61 | @staticmethod
62 | def backward(ctx, grad_out):
63 | raise NotImplementedError
64 |
65 |
66 | if __name__ == '__main__':
67 | pass
68 |
--------------------------------------------------------------------------------
/pcdet/ops/roipoint_pool3d/src/roipoint_pool3d.cpp:
--------------------------------------------------------------------------------
1 | #include
2 | #include
3 |
4 | #define CHECK_CUDA(x) do { \
5 | if (!x.type().is_cuda()) { \
6 | fprintf(stderr, "%s must be CUDA tensor at %s:%d\n", #x, __FILE__, __LINE__); \
7 | exit(-1); \
8 | } \
9 | } while (0)
10 | #define CHECK_CONTIGUOUS(x) do { \
11 | if (!x.is_contiguous()) { \
12 | fprintf(stderr, "%s must be contiguous tensor at %s:%d\n", #x, __FILE__, __LINE__); \
13 | exit(-1); \
14 | } \
15 | } while (0)
16 | #define CHECK_INPUT(x) CHECK_CUDA(x);CHECK_CONTIGUOUS(x)
17 |
18 |
19 | void roipool3dLauncher(int batch_size, int pts_num, int boxes_num, int feature_in_len, int sampled_pts_num,
20 | const float *xyz, const float *boxes3d, const float *pts_feature, float *pooled_features, int *pooled_empty_flag);
21 |
22 |
23 | int roipool3d_gpu(at::Tensor xyz, at::Tensor boxes3d, at::Tensor pts_feature, at::Tensor pooled_features, at::Tensor pooled_empty_flag){
24 | // params xyz: (B, N, 3)
25 | // params boxes3d: (B, M, 7)
26 | // params pts_feature: (B, N, C)
27 | // params pooled_features: (B, M, 512, 3+C)
28 | // params pooled_empty_flag: (B, M)
29 | CHECK_INPUT(xyz);
30 | CHECK_INPUT(boxes3d);
31 | CHECK_INPUT(pts_feature);
32 | CHECK_INPUT(pooled_features);
33 | CHECK_INPUT(pooled_empty_flag);
34 |
35 | int batch_size = xyz.size(0);
36 | int pts_num = xyz.size(1);
37 | int boxes_num = boxes3d.size(1);
38 | int feature_in_len = pts_feature.size(2);
39 | int sampled_pts_num = pooled_features.size(2);
40 |
41 |
42 | const float * xyz_data = xyz.data();
43 | const float * boxes3d_data = boxes3d.data();
44 | const float * pts_feature_data = pts_feature.data();
45 | float * pooled_features_data = pooled_features.data();
46 | int * pooled_empty_flag_data = pooled_empty_flag.data();
47 |
48 | roipool3dLauncher(batch_size, pts_num, boxes_num, feature_in_len, sampled_pts_num,
49 | xyz_data, boxes3d_data, pts_feature_data, pooled_features_data, pooled_empty_flag_data);
50 |
51 |
52 |
53 | return 1;
54 | }
55 |
56 |
57 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
58 | m.def("forward", &roipool3d_gpu, "roipool3d forward (CUDA)");
59 | }
60 |
61 |
--------------------------------------------------------------------------------
/pcdet/utils/__init__.py:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/CVMI-Lab/SparseKD/2e00edef6b94fb208ba64193940bc856e5effbcb/pcdet/utils/__init__.py
--------------------------------------------------------------------------------
/pcdet/utils/kd_utils/__init__.py:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/CVMI-Lab/SparseKD/2e00edef6b94fb208ba64193940bc856e5effbcb/pcdet/utils/kd_utils/__init__.py
--------------------------------------------------------------------------------
/pcdet/utils/kd_utils/kd_utils.py:
--------------------------------------------------------------------------------
1 | import torch
2 | from torch.nn.utils import clip_grad_norm_
3 | from pcdet.config import cfg
4 |
5 |
6 | def process_kd_config():
7 | """
8 | put kd related config to model
9 | Returns:
10 |
11 | """
12 | # Global dense head indicator.
13 | # Build KD head or not
14 | cfg.MODEL.KD = True
15 | cfg.MODEL_TEACHER.KD = True
16 |
17 | # Only student model have KD_LOSS config
18 | if cfg.KD_LOSS.ENABLED:
19 | cfg.MODEL.KD_LOSS = cfg.KD_LOSS
20 |
21 | parse_key_list = ['LOGIT_KD', 'FEATURE_KD', 'LABEL_ASSIGN_KD', 'VFE_KD', 'ROI_KD']
22 |
23 | for key in parse_key_list:
24 | if cfg.KD.get(key, None) and cfg.KD[key].ENABLED:
25 | cfg.MODEL.DENSE_HEAD[key] = cfg.KD[key]
26 | cfg.MODEL[key] = cfg.KD[key]
27 | cfg.MODEL_TEACHER[key] = cfg.KD[key]
28 |
29 | if cfg.KD.get('VFE_KD', None) and cfg.KD.VFE_KD.ENABLED:
30 | cfg.MODEL.VFE.VFE_KD = cfg.KD.VFE_KD
31 | cfg.MODEL_TEACHER.VFE.VFE_KD = cfg.KD.VFE_KD
32 |
33 |
34 | def prepare_kd_modules(teacher_model, model):
35 | # add a flag to indicate KD for each module in the detector
36 | if isinstance(model, torch.nn.parallel.DistributedDataParallel):
37 | for cur_module in model.module.module_list:
38 | cur_module.kd = True
39 | else:
40 | for cur_module in model.module_list:
41 | cur_module.kd = True
42 |
43 | put_teacher_prior_to_student(teacher_model, model, cfg.MODEL)
44 |
45 | # if student need some extra learnable moduleq
46 | if hasattr(model.dense_head.kd_head, 'register_extra_layers'):
47 | model.dense_head.kd_head.register_extra_layers(model.dense_head)
48 |
49 |
50 | def put_teacher_prior_to_student(teacher_model, student_model, model_cfg):
51 | student_model.kd_head.voxel_size_tea = teacher_model.kd_head.voxel_size
52 | student_model.kd_head.feature_map_stride_tea = teacher_model.kd_head.feature_map_stride
53 |
54 | if model_cfg.get('FEATURE_KD', None) and model_cfg.FEATURE_KD.get('ALIGN', None) and \
55 | model_cfg.FEATURE_KD.ALIGN.MODE == 'bn':
56 | student_model.dense_head.kd_head.get_prior_knowledge_from_teacher_model(
57 | teacher_model, model_cfg
58 | )
59 |
60 | if model_cfg.get('VFE_KD', None) and model_cfg.VFE_KD.get('CN_ALIGN', None) and \
61 | model_cfg.VFE_KD.CN_ALIGN.ENABLED and model_cfg.VFE_KD.CN_ALIGN.MODE == 'bn':
62 | student_model.kd_head.select_topk_channels_in_teacher_vfe(teacher_model)
63 |
64 |
65 | def pop_teacher_intermediate_features(batch):
66 | pop_list = ['pillar_features', 'spatial_features', 'spatial_features_2d']
67 |
68 | for key in pop_list:
69 | if key in batch:
70 | temp = batch.pop(key)
71 | del temp
72 |
73 |
74 | def cal_channel_attention_mask(feature):
75 | """
76 |
77 | Args:
78 | feature: [B, C, H, W]
79 |
80 | Returns:
81 | mask: [B, C]
82 | """
83 | bs, ch, height, width = feature.shape
84 | feat = feature.view(bs, ch, -1)
85 | mask = torch.abs(feat).mean(dim=-1)
86 |
87 | return mask
88 |
89 |
90 | def cal_spatial_attention_mask(feature):
91 | """
92 | Args:
93 | feature: [B, C, H, W]
94 |
95 | Returns:
96 | mask: [B, C]
97 | """
98 | mask = torch.abs(feature).mean(dim=1)
99 | return mask
100 |
101 |
--------------------------------------------------------------------------------
/pcdet/utils/kd_utils/kd_vis_utils.py:
--------------------------------------------------------------------------------
1 | import torch
2 | import matplotlib.pyplot as plt
3 | from . import kd_utils
4 |
5 |
6 | def vis_channel_sptial_magnitude(feature, k=16):
7 | ch_mag = kd_utils.cal_channel_attention_mask(feature)
8 | spatial_mag = kd_utils.cal_spatial_attention_mask(feature)
9 | ch_one_hot_mask = ch_mag > 0
10 |
11 | # magnitude range
12 | # print('student ch mag range: [%f, %f]' % (ch_mag.min(), ch_mag.max()))
13 | # print('student spatial mag range: [%f, %f]' % (spatial_mag.min(), spatial_mag.max()))
14 |
15 | fig, ax = plt.subplots(1, 3)
16 |
17 | ax[0].imshow(ch_mag.view(k, -1).cpu().numpy())
18 | ax[0].set_title('channel mask')
19 | ax[1].imshow(spatial_mag[0].cpu().numpy())
20 | ax[1].set_title('spatial mask')
21 | ax[2].imshow(ch_one_hot_mask.view(k, -1).detach().cpu().numpy(), vmin=0, vmax=1)
22 | ax[2].set_title('channel one-hot mask')
23 |
24 | plt.show()
25 |
26 | return
27 |
28 |
29 | def cal_feature_dist(feature, feature_tea, mode='cosine', topk=None):
30 | """_summary_
31 |
32 | Args:
33 | feature (_type_): _description_
34 | feature_tea (_type_): _description_
35 | mode (str, optional): [cosine, kl]. Defaults to 'cosine'.
36 | """
37 | bs = feature.shape[0]
38 |
39 | if topk is not None:
40 | feature = select_topk_feature_channel(feature, k=topk)
41 | feature_tea = select_topk_feature_channel(feature_tea, k=topk)
42 |
43 | if mode == 'cosine':
44 | dist = torch.nn.functional.cosine_similarity(feature.view(bs, -1), feature_tea.view(bs, -1))
45 | elif mode == 'kl':
46 | dist = torch.nn.functional.kl_div(feature, feature_tea)
47 | else:
48 | raise NotImplementedError
49 |
50 | return dist
51 |
52 |
53 | def select_topk_feature_channel(feature, k):
54 | ch_mag = kd_utils.cal_channel_attention_mask(feature)
55 | _, channel_idx = torch.topk(ch_mag, k=k)
56 |
57 | return feature[:, channel_idx, ...]
58 |
59 |
--------------------------------------------------------------------------------
/pcdet/utils/object3d_kitti.py:
--------------------------------------------------------------------------------
1 | import numpy as np
2 |
3 |
4 | def get_objects_from_label(label_file):
5 | with open(label_file, 'r') as f:
6 | lines = f.readlines()
7 | objects = [Object3d(line) for line in lines]
8 | return objects
9 |
10 |
11 | def cls_type_to_id(cls_type):
12 | type_to_id = {'Car': 1, 'Pedestrian': 2, 'Cyclist': 3, 'Van': 4}
13 | if cls_type not in type_to_id.keys():
14 | return -1
15 | return type_to_id[cls_type]
16 |
17 |
18 | class Object3d(object):
19 | def __init__(self, line):
20 | label = line.strip().split(' ')
21 | self.src = line
22 | self.cls_type = label[0]
23 | self.cls_id = cls_type_to_id(self.cls_type)
24 | self.truncation = float(label[1])
25 | self.occlusion = float(label[2]) # 0:fully visible 1:partly occluded 2:largely occluded 3:unknown
26 | self.alpha = float(label[3])
27 | self.box2d = np.array((float(label[4]), float(label[5]), float(label[6]), float(label[7])), dtype=np.float32)
28 | self.h = float(label[8])
29 | self.w = float(label[9])
30 | self.l = float(label[10])
31 | self.loc = np.array((float(label[11]), float(label[12]), float(label[13])), dtype=np.float32)
32 | self.dis_to_cam = np.linalg.norm(self.loc)
33 | self.ry = float(label[14])
34 | self.score = float(label[15]) if label.__len__() == 16 else -1.0
35 | self.level_str = None
36 | self.level = self.get_kitti_obj_level()
37 |
38 | def get_kitti_obj_level(self):
39 | height = float(self.box2d[3]) - float(self.box2d[1]) + 1
40 |
41 | if height >= 40 and self.truncation <= 0.15 and self.occlusion <= 0:
42 | self.level_str = 'Easy'
43 | return 0 # Easy
44 | elif height >= 25 and self.truncation <= 0.3 and self.occlusion <= 1:
45 | self.level_str = 'Moderate'
46 | return 1 # Moderate
47 | elif height >= 25 and self.truncation <= 0.5 and self.occlusion <= 2:
48 | self.level_str = 'Hard'
49 | return 2 # Hard
50 | else:
51 | self.level_str = 'UnKnown'
52 | return -1
53 |
54 | def generate_corners3d(self):
55 | """
56 | generate corners3d representation for this object
57 | :return corners_3d: (8, 3) corners of box3d in camera coord
58 | """
59 | l, h, w = self.l, self.h, self.w
60 | x_corners = [l / 2, l / 2, -l / 2, -l / 2, l / 2, l / 2, -l / 2, -l / 2]
61 | y_corners = [0, 0, 0, 0, -h, -h, -h, -h]
62 | z_corners = [w / 2, -w / 2, -w / 2, w / 2, w / 2, -w / 2, -w / 2, w / 2]
63 |
64 | R = np.array([[np.cos(self.ry), 0, np.sin(self.ry)],
65 | [0, 1, 0],
66 | [-np.sin(self.ry), 0, np.cos(self.ry)]])
67 | corners3d = np.vstack([x_corners, y_corners, z_corners]) # (3, 8)
68 | corners3d = np.dot(R, corners3d).T
69 | corners3d = corners3d + self.loc
70 | return corners3d
71 |
72 | def to_str(self):
73 | print_str = '%s %.3f %.3f %.3f box2d: %s hwl: [%.3f %.3f %.3f] pos: %s ry: %.3f' \
74 | % (self.cls_type, self.truncation, self.occlusion, self.alpha, self.box2d, self.h, self.w, self.l,
75 | self.loc, self.ry)
76 | return print_str
77 |
78 | def to_kitti_format(self):
79 | kitti_str = '%s %.2f %d %.2f %.2f %.2f %.2f %.2f %.2f %.2f %.2f %.2f %.2f %.2f %.2f' \
80 | % (self.cls_type, self.truncation, int(self.occlusion), self.alpha, self.box2d[0], self.box2d[1],
81 | self.box2d[2], self.box2d[3], self.h, self.w, self.l, self.loc[0], self.loc[1], self.loc[2],
82 | self.ry)
83 | return kitti_str
84 |
--------------------------------------------------------------------------------
/pcdet/utils/spconv_utils.py:
--------------------------------------------------------------------------------
1 | from typing import Set
2 |
3 | try:
4 | import spconv.pytorch as spconv
5 | except:
6 | import spconv as spconv
7 |
8 | import torch.nn as nn
9 |
10 |
11 | def find_all_spconv_keys(model: nn.Module, prefix="") -> Set[str]:
12 | """
13 | Finds all spconv keys that need to have weight's transposed
14 | """
15 | found_keys: Set[str] = set()
16 | for name, child in model.named_children():
17 | new_prefix = f"{prefix}.{name}" if prefix != "" else name
18 |
19 | if isinstance(child, spconv.conv.SparseConvolution):
20 | new_prefix = f"{new_prefix}.weight"
21 | found_keys.add(new_prefix)
22 |
23 | found_keys.update(find_all_spconv_keys(child, prefix=new_prefix))
24 |
25 | return found_keys
26 |
27 |
28 | def replace_feature(out, new_features):
29 | if "replace_feature" in out.__dir__():
30 | # spconv 2.x behaviour
31 | return out.replace_feature(new_features)
32 | else:
33 | out.features = new_features
34 | return out
35 |
36 |
37 | def clone_sp_tensor(out, batch_size):
38 | new_tensor = spconv.SparseConvTensor(
39 | features=out.features.clone().detach(),
40 | indices=out.indices.int(),
41 | spatial_shape=out.spatial_shape,
42 | batch_size=batch_size
43 | )
44 | return new_tensor
45 |
--------------------------------------------------------------------------------
/pcdet/utils/transform_utils.py:
--------------------------------------------------------------------------------
1 | import math
2 | import torch
3 |
4 | try:
5 | from kornia.geometry.conversions import (
6 | convert_points_to_homogeneous,
7 | convert_points_from_homogeneous,
8 | )
9 | except:
10 | pass
11 | # print('Warning: kornia is not installed. This package is only required by CaDDN')
12 |
13 |
14 | def project_to_image(project, points):
15 | """
16 | Project points to image
17 | Args:
18 | project [torch.tensor(..., 3, 4)]: Projection matrix
19 | points [torch.Tensor(..., 3)]: 3D points
20 | Returns:
21 | points_img [torch.Tensor(..., 2)]: Points in image
22 | points_depth [torch.Tensor(...)]: Depth of each point
23 | """
24 | # Reshape tensors to expected shape
25 | points = convert_points_to_homogeneous(points)
26 | points = points.unsqueeze(dim=-1)
27 | project = project.unsqueeze(dim=1)
28 |
29 | # Transform points to image and get depths
30 | points_t = project @ points
31 | points_t = points_t.squeeze(dim=-1)
32 | points_img = convert_points_from_homogeneous(points_t)
33 | points_depth = points_t[..., -1] - project[..., 2, 3]
34 |
35 | return points_img, points_depth
36 |
37 |
38 | def normalize_coords(coords, shape):
39 | """
40 | Normalize coordinates of a grid between [-1, 1]
41 | Args:
42 | coords: (..., 3), Coordinates in grid
43 | shape: (3), Grid shape
44 | Returns:
45 | norm_coords: (.., 3), Normalized coordinates in grid
46 | """
47 | min_n = -1
48 | max_n = 1
49 | shape = torch.flip(shape, dims=[0]) # Reverse ordering of shape
50 |
51 | # Subtract 1 since pixel indexing from [0, shape - 1]
52 | norm_coords = coords / (shape - 1) * (max_n - min_n) + min_n
53 | return norm_coords
54 |
55 |
56 | def bin_depths(depth_map, mode, depth_min, depth_max, num_bins, target=False):
57 | """
58 | Converts depth map into bin indices
59 | Args:
60 | depth_map: (H, W), Depth Map
61 | mode: string, Discretiziation mode (See https://arxiv.org/pdf/2005.13423.pdf for more details)
62 | UD: Uniform discretiziation
63 | LID: Linear increasing discretiziation
64 | SID: Spacing increasing discretiziation
65 | depth_min: float, Minimum depth value
66 | depth_max: float, Maximum depth value
67 | num_bins: int, Number of depth bins
68 | target: bool, Whether the depth bins indices will be used for a target tensor in loss comparison
69 | Returns:
70 | indices: (H, W), Depth bin indices
71 | """
72 | if mode == "UD":
73 | bin_size = (depth_max - depth_min) / num_bins
74 | indices = ((depth_map - depth_min) / bin_size)
75 | elif mode == "LID":
76 | bin_size = 2 * (depth_max - depth_min) / (num_bins * (1 + num_bins))
77 | indices = -0.5 + 0.5 * torch.sqrt(1 + 8 * (depth_map - depth_min) / bin_size)
78 | elif mode == "SID":
79 | indices = num_bins * (torch.log(1 + depth_map) - math.log(1 + depth_min)) / \
80 | (math.log(1 + depth_max) - math.log(1 + depth_min))
81 | else:
82 | raise NotImplementedError
83 |
84 | if target:
85 | # Remove indicies outside of bounds
86 | mask = (indices < 0) | (indices > num_bins) | (~torch.isfinite(indices))
87 | indices[mask] = num_bins
88 |
89 | # Convert to integer
90 | indices = indices.type(torch.int64)
91 | return indices
92 |
--------------------------------------------------------------------------------
/requirements.txt:
--------------------------------------------------------------------------------
1 | numpy<=1.20
2 | llvmlite
3 | numba
4 | torch>=1.1
5 | tensorboardX
6 | easydict
7 | pyyaml
8 | scikit-image
9 | tqdm
10 | torchvision
11 | SharedArray
12 | prettytable
13 |
--------------------------------------------------------------------------------
/tools/_init_path.py:
--------------------------------------------------------------------------------
1 | import sys
2 | sys.path.insert(0, '../')
--------------------------------------------------------------------------------
/tools/cfgs/dataset_configs/kitti_dataset.yaml:
--------------------------------------------------------------------------------
1 | DATASET: 'KittiDataset'
2 | DATA_PATH: '../data/kitti'
3 |
4 | POINT_CLOUD_RANGE: [0, -40, -3, 70.4, 40, 1]
5 |
6 | DATA_SPLIT: {
7 | 'train': train,
8 | 'test': val
9 | }
10 |
11 | INFO_PATH: {
12 | 'train': [kitti_infos_train.pkl],
13 | 'test': [kitti_infos_val.pkl],
14 | }
15 |
16 | GET_ITEM_LIST: ["points"]
17 | FOV_POINTS_ONLY: True
18 |
19 | DATA_AUGMENTOR:
20 | DISABLE_AUG_LIST: ['placeholder']
21 | AUG_CONFIG_LIST:
22 | - NAME: gt_sampling
23 | USE_ROAD_PLANE: True
24 | DB_INFO_PATH:
25 | - kitti_dbinfos_train_global.pkl
26 |
27 | USE_SHARED_MEMORY: True # set it to True to speed up (it costs about 15GB shared memory)
28 | DB_DATA_PATH:
29 | - kitti_dbinfos_train_global.npy
30 | PREPARE: {
31 | filter_by_min_points: ['Car:5', 'Pedestrian:5', 'Cyclist:5'],
32 | filter_by_difficulty: [-1],
33 | }
34 |
35 | SAMPLE_GROUPS: ['Car:20','Pedestrian:15', 'Cyclist:15']
36 | NUM_POINT_FEATURES: 4
37 | DATABASE_WITH_FAKELIDAR: False
38 | REMOVE_EXTRA_WIDTH: [0.0, 0.0, 0.0]
39 | LIMIT_WHOLE_SCENE: True
40 |
41 | - NAME: random_world_flip
42 | ALONG_AXIS_LIST: ['x']
43 |
44 | - NAME: random_world_rotation
45 | WORLD_ROT_ANGLE: [-0.78539816, 0.78539816]
46 |
47 | - NAME: random_world_scaling
48 | WORLD_SCALE_RANGE: [0.95, 1.05]
49 |
50 |
51 | POINT_FEATURE_ENCODING: {
52 | encoding_type: absolute_coordinates_encoding,
53 | used_feature_list: ['x', 'y', 'z', 'intensity'],
54 | src_feature_list: ['x', 'y', 'z', 'intensity'],
55 | }
56 |
57 |
58 | DATA_PROCESSOR:
59 | - NAME: mask_points_and_boxes_outside_range
60 | REMOVE_OUTSIDE_BOXES: True
61 |
62 | - NAME: shuffle_points
63 | SHUFFLE_ENABLED: {
64 | 'train': True,
65 | 'test': False
66 | }
67 |
68 | - NAME: transform_points_to_voxels
69 | VOXEL_SIZE: [0.05, 0.05, 0.1]
70 | MAX_POINTS_PER_VOXEL: 5
71 | MAX_NUMBER_OF_VOXELS: {
72 | 'train': 16000,
73 | 'test': 40000
74 | }
75 |
--------------------------------------------------------------------------------
/tools/cfgs/dataset_configs/nuscenes_dataset.yaml:
--------------------------------------------------------------------------------
1 | DATASET: 'NuScenesDataset'
2 | DATA_PATH: '../data/nuscenes'
3 |
4 | VERSION: 'v1.0-trainval'
5 | MAX_SWEEPS: 10
6 | PRED_VELOCITY: True
7 | SET_NAN_VELOCITY_TO_ZEROS: True
8 | FILTER_MIN_POINTS_IN_GT: 1
9 |
10 | DATA_SPLIT: {
11 | 'train': train,
12 | 'test': val
13 | }
14 |
15 | INFO_PATH: {
16 | 'train': [nuscenes_infos_10sweeps_train.pkl],
17 | 'test': [nuscenes_infos_10sweeps_val.pkl],
18 | }
19 |
20 | POINT_CLOUD_RANGE: [-51.2, -51.2, -5.0, 51.2, 51.2, 3.0]
21 |
22 | BALANCED_RESAMPLING: True
23 |
24 | DATA_AUGMENTOR:
25 | DISABLE_AUG_LIST: ['placeholder']
26 | AUG_CONFIG_LIST:
27 | - NAME: gt_sampling
28 | DB_INFO_PATH:
29 | - nuscenes_dbinfos_10sweeps_withvelo.pkl
30 | PREPARE: {
31 | filter_by_min_points: [
32 | 'car:5','truck:5', 'construction_vehicle:5', 'bus:5', 'trailer:5',
33 | 'barrier:5', 'motorcycle:5', 'bicycle:5', 'pedestrian:5', 'traffic_cone:5'
34 | ],
35 | }
36 |
37 | SAMPLE_GROUPS: [
38 | 'car:2','truck:3', 'construction_vehicle:7', 'bus:4', 'trailer:6',
39 | 'barrier:2', 'motorcycle:6', 'bicycle:6', 'pedestrian:2', 'traffic_cone:2'
40 | ]
41 |
42 | NUM_POINT_FEATURES: 5
43 | DATABASE_WITH_FAKELIDAR: False
44 | REMOVE_EXTRA_WIDTH: [0.0, 0.0, 0.0]
45 | LIMIT_WHOLE_SCENE: True
46 |
47 | - NAME: random_world_flip
48 | ALONG_AXIS_LIST: ['x', 'y']
49 |
50 | - NAME: random_world_rotation
51 | WORLD_ROT_ANGLE: [-0.3925, 0.3925]
52 |
53 | - NAME: random_world_scaling
54 | WORLD_SCALE_RANGE: [0.95, 1.05]
55 |
56 |
57 | POINT_FEATURE_ENCODING: {
58 | encoding_type: absolute_coordinates_encoding,
59 | used_feature_list: ['x', 'y', 'z', 'intensity', 'timestamp'],
60 | src_feature_list: ['x', 'y', 'z', 'intensity', 'timestamp'],
61 | }
62 |
63 |
64 | DATA_PROCESSOR:
65 | - NAME: mask_points_and_boxes_outside_range
66 | REMOVE_OUTSIDE_BOXES: True
67 |
68 | - NAME: shuffle_points
69 | SHUFFLE_ENABLED: {
70 | 'train': True,
71 | 'test': True
72 | }
73 |
74 | - NAME: transform_points_to_voxels
75 | VOXEL_SIZE: [0.1, 0.1, 0.2]
76 | MAX_POINTS_PER_VOXEL: 10
77 | MAX_NUMBER_OF_VOXELS: {
78 | 'train': 60000,
79 | 'test': 60000
80 | }
81 |
--------------------------------------------------------------------------------
/tools/cfgs/dataset_configs/waymo_dataset.yaml:
--------------------------------------------------------------------------------
1 | DATASET: 'WaymoDataset'
2 | DATA_PATH: '../data/waymo'
3 | PROCESSED_DATA_TAG: 'waymo_processed_data_v0_5_0'
4 |
5 | POINT_CLOUD_RANGE: [-75.2, -75.2, -2, 75.2, 75.2, 4]
6 |
7 | DATA_SPLIT: {
8 | 'train': train,
9 | 'test': val
10 | }
11 |
12 | SAMPLED_INTERVAL: {
13 | 'train': 5,
14 | 'test': 5
15 | }
16 |
17 | FILTER_EMPTY_BOXES_FOR_TRAIN: True
18 | DISABLE_NLZ_FLAG_ON_POINTS: True
19 |
20 | USE_SHARED_MEMORY: False # it will load the data to shared memory to speed up (DO NOT USE IT IF YOU DO NOT FULLY UNDERSTAND WHAT WILL HAPPEN)
21 | SHARED_MEMORY_FILE_LIMIT: 35000 # set it based on the size of your shared memory
22 |
23 | DATA_AUGMENTOR:
24 | DISABLE_AUG_LIST: ['placeholder']
25 | AUG_CONFIG_LIST:
26 | - NAME: gt_sampling
27 | USE_ROAD_PLANE: False
28 | DB_INFO_PATH:
29 | - waymo_processed_data_v0_5_0_waymo_dbinfos_train_sampled_1_global.pkl
30 |
31 | USE_SHARED_MEMORY: True # set it to True to speed up (it costs about 15GB shared memory)
32 | DB_DATA_PATH:
33 | - waymo_processed_data_v0_5_0_gt_database_train_sampled_1_global.npy
34 |
35 | PREPARE: {
36 | filter_by_min_points: ['Vehicle:5', 'Pedestrian:5', 'Cyclist:5'],
37 | filter_by_difficulty: [-1],
38 | }
39 |
40 | SAMPLE_GROUPS: ['Vehicle:15', 'Pedestrian:10', 'Cyclist:10']
41 | NUM_POINT_FEATURES: 5
42 | REMOVE_EXTRA_WIDTH: [0.0, 0.0, 0.0]
43 | LIMIT_WHOLE_SCENE: True
44 |
45 | - NAME: random_world_flip
46 | ALONG_AXIS_LIST: ['x', 'y']
47 |
48 | - NAME: random_world_rotation
49 | WORLD_ROT_ANGLE: [-0.78539816, 0.78539816]
50 |
51 | - NAME: random_world_scaling
52 | WORLD_SCALE_RANGE: [0.95, 1.05]
53 |
54 |
55 | POINT_FEATURE_ENCODING: {
56 | encoding_type: absolute_coordinates_encoding,
57 | used_feature_list: ['x', 'y', 'z', 'intensity', 'elongation'],
58 | src_feature_list: ['x', 'y', 'z', 'intensity', 'elongation'],
59 | }
60 |
61 |
62 | DATA_PROCESSOR:
63 | - NAME: mask_points_and_boxes_outside_range
64 | REMOVE_OUTSIDE_BOXES: True
65 |
66 | - NAME: shuffle_points
67 | SHUFFLE_ENABLED: {
68 | 'train': True,
69 | 'test': True
70 | }
71 |
72 | - NAME: transform_points_to_voxels
73 | VOXEL_SIZE: [0.1, 0.1, 0.15]
74 | MAX_POINTS_PER_VOXEL: 5
75 | MAX_NUMBER_OF_VOXELS: {
76 | 'train': 150000,
77 | 'test': 150000
78 | }
79 |
--------------------------------------------------------------------------------
/tools/cfgs/kitti_models/PartA2_free.yaml:
--------------------------------------------------------------------------------
1 | CLASS_NAMES: ['Car', 'Pedestrian', 'Cyclist']
2 |
3 | DATA_CONFIG:
4 | _BASE_CONFIG_: cfgs/dataset_configs/kitti_dataset.yaml
5 |
6 |
7 | MODEL:
8 | NAME: PointRCNN
9 |
10 | VFE:
11 | NAME: MeanVFE
12 |
13 | BACKBONE_3D:
14 | NAME: UNetV2
15 | RETURN_ENCODED_TENSOR: False
16 |
17 | POINT_HEAD:
18 | NAME: PointIntraPartOffsetHead
19 | CLS_FC: [128, 128]
20 | PART_FC: [128, 128]
21 | REG_FC: [128, 128]
22 | CLASS_AGNOSTIC: False
23 | USE_POINT_FEATURES_BEFORE_FUSION: False
24 | TARGET_CONFIG:
25 | GT_EXTRA_WIDTH: [0.2, 0.2, 0.2]
26 | BOX_CODER: PointResidualCoder
27 | BOX_CODER_CONFIG: {
28 | 'use_mean_size': True,
29 | 'mean_size': [
30 | [3.9, 1.6, 1.56],
31 | [0.8, 0.6, 1.73],
32 | [1.76, 0.6, 1.73]
33 | ]
34 | }
35 |
36 | LOSS_CONFIG:
37 | LOSS_REG: WeightedSmoothL1Loss
38 | LOSS_WEIGHTS: {
39 | 'point_cls_weight': 1.0,
40 | 'point_box_weight': 1.0,
41 | 'point_part_weight': 1.0,
42 | 'code_weights': [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0]
43 | }
44 |
45 | ROI_HEAD:
46 | NAME: PartA2FCHead
47 | CLASS_AGNOSTIC: True
48 |
49 | SHARED_FC: [256, 256, 256]
50 | CLS_FC: [256, 256]
51 | REG_FC: [256, 256]
52 | DP_RATIO: 0.3
53 | DISABLE_PART: True
54 | SEG_MASK_SCORE_THRESH: 0.0
55 |
56 | NMS_CONFIG:
57 | TRAIN:
58 | NMS_TYPE: nms_gpu
59 | MULTI_CLASSES_NMS: False
60 | NMS_PRE_MAXSIZE: 9000
61 | NMS_POST_MAXSIZE: 512
62 | NMS_THRESH: 0.8
63 | TEST:
64 | NMS_TYPE: nms_gpu
65 | MULTI_CLASSES_NMS: False
66 | NMS_PRE_MAXSIZE: 9000
67 | NMS_POST_MAXSIZE: 100
68 | NMS_THRESH: 0.85
69 |
70 | ROI_AWARE_POOL:
71 | POOL_SIZE: 12
72 | NUM_FEATURES: 128
73 | MAX_POINTS_PER_VOXEL: 128
74 |
75 | TARGET_CONFIG:
76 | BOX_CODER: ResidualCoder
77 | ROI_PER_IMAGE: 128
78 | FG_RATIO: 0.5
79 |
80 | SAMPLE_ROI_BY_EACH_CLASS: True
81 | CLS_SCORE_TYPE: roi_iou
82 |
83 | CLS_FG_THRESH: 0.75
84 | CLS_BG_THRESH: 0.25
85 | CLS_BG_THRESH_LO: 0.1
86 | HARD_BG_RATIO: 0.8
87 |
88 | REG_FG_THRESH: 0.65
89 |
90 | LOSS_CONFIG:
91 | CLS_LOSS: BinaryCrossEntropy
92 | REG_LOSS: smooth-l1
93 | CORNER_LOSS_REGULARIZATION: True
94 | LOSS_WEIGHTS: {
95 | 'rcnn_cls_weight': 1.0,
96 | 'rcnn_reg_weight': 1.0,
97 | 'rcnn_corner_weight': 1.0,
98 | 'code_weights': [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0]
99 | }
100 |
101 | POST_PROCESSING:
102 | RECALL_THRESH_LIST: [0.3, 0.5, 0.7]
103 | SCORE_THRESH: 0.1
104 | OUTPUT_RAW_SCORE: False
105 |
106 | EVAL_METRIC: kitti
107 |
108 | NMS_CONFIG:
109 | MULTI_CLASSES_NMS: False
110 | NMS_TYPE: nms_gpu
111 | NMS_THRESH: 0.1
112 | NMS_PRE_MAXSIZE: 4096
113 | NMS_POST_MAXSIZE: 500
114 |
115 |
116 | OPTIMIZATION:
117 | BATCH_SIZE_PER_GPU: 4
118 | NUM_EPOCHS: 80
119 |
120 | OPTIMIZER: adam_onecycle
121 | LR: 0.003
122 | WEIGHT_DECAY: 0.01
123 | MOMENTUM: 0.9
124 |
125 | MOMS: [0.95, 0.85]
126 | PCT_START: 0.4
127 | DIV_FACTOR: 10
128 | DECAY_STEP_LIST: [35, 45]
129 | LR_DECAY: 0.1
130 | LR_CLIP: 0.0000001
131 |
132 | LR_WARMUP: False
133 | WARMUP_EPOCH: 1
134 |
135 | GRAD_NORM_CLIP: 10
136 |
--------------------------------------------------------------------------------
/tools/cfgs/waymo_models/centerpoint.yaml:
--------------------------------------------------------------------------------
1 | CLASS_NAMES: ['Vehicle', 'Pedestrian', 'Cyclist']
2 |
3 | DATA_CONFIG:
4 | _BASE_CONFIG_: cfgs/dataset_configs/waymo_dataset.yaml
5 |
6 | MODEL:
7 | NAME: CenterPoint
8 | IGNORE_PRETRAIN_MODULES: ['placeholder']
9 |
10 | VFE:
11 | NAME: MeanVFE
12 |
13 | BACKBONE_3D:
14 | NAME: VoxelResBackBone8x
15 | ACT_FN: ReLU
16 | # [input, conv1, conv2, conv3, conv4, output]
17 | NUM_FILTERS: [16, 16, 32, 64, 128, 128]
18 | LAYER_NUMS: [1, 2, 3, 3, 3, 1]
19 | WIDTH: 1.0
20 |
21 | MAP_TO_BEV:
22 | NAME: HeightCompression
23 | NUM_BEV_FEATURES: 256
24 |
25 | BACKBONE_2D:
26 | NAME: BaseBEVBackbone
27 | ACT_FN: ReLU
28 | NORM_TYPE: BatchNorm2d
29 | WIDTH: 1.0
30 |
31 | LAYER_NUMS: [5, 5]
32 | LAYER_STRIDES: [1, 2]
33 | NUM_FILTERS: [128, 256]
34 | UPSAMPLE_STRIDES: [1, 2]
35 | NUM_UPSAMPLE_FILTERS: [256, 256]
36 |
37 | DENSE_HEAD:
38 | NAME: CenterHead
39 | CLASS_AGNOSTIC: False
40 | ACT_FN: ReLU
41 | NORM_TYPE: BatchNorm2d
42 |
43 | CLASS_NAMES_EACH_HEAD: [
44 | ['Vehicle', 'Pedestrian', 'Cyclist']
45 | ]
46 |
47 | SHARED_CONV_CHANNEL: 64
48 | USE_BIAS_BEFORE_NORM: True
49 | NUM_HM_CONV: 2
50 | SEPARATE_HEAD_CFG:
51 | HEAD_ORDER: ['center', 'center_z', 'dim', 'rot']
52 | HEAD_DICT: {
53 | 'center': {'out_channels': 2, 'num_conv': 2},
54 | 'center_z': {'out_channels': 1, 'num_conv': 2},
55 | 'dim': {'out_channels': 3, 'num_conv': 2},
56 | 'rot': {'out_channels': 2, 'num_conv': 2},
57 | }
58 |
59 | TARGET_ASSIGNER_CONFIG:
60 | FEATURE_MAP_STRIDE: 8
61 | NUM_MAX_OBJS: 500
62 | GAUSSIAN_OVERLAP: 0.1
63 | MIN_RADIUS: 2
64 |
65 | LOSS_CONFIG:
66 | LOSS_WEIGHTS: {
67 | 'cls_weight': 1.0,
68 | 'loc_weight': 2.0,
69 | 'code_weights': [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0]
70 | }
71 |
72 | POST_PROCESSING:
73 | SCORE_THRESH: 0.1
74 | POST_CENTER_LIMIT_RANGE: [-75.2, -75.2, -2, 75.2, 75.2, 4]
75 | MAX_OBJ_PER_SAMPLE: 500
76 | NMS_CONFIG:
77 | NMS_TYPE: nms_gpu
78 | NMS_THRESH: 0.7
79 | NMS_PRE_MAXSIZE: 4096
80 | NMS_POST_MAXSIZE: 500
81 |
82 | POST_PROCESSING:
83 | RECALL_THRESH_LIST: [0.3, 0.5, 0.7]
84 |
85 | EVAL_METRIC: waymo
86 | EVAL_CLASSES: {
87 | 'LEVEL_2/AP': [ 'OBJECT_TYPE_TYPE_VEHICLE_LEVEL_2/AP',
88 | 'OBJECT_TYPE_TYPE_PEDESTRIAN_LEVEL_2/AP',
89 | 'OBJECT_TYPE_TYPE_CYCLIST_LEVEL_2/AP'
90 | ],
91 | 'LEVEL_2/APH': [ 'OBJECT_TYPE_TYPE_VEHICLE_LEVEL_2/APH',
92 | 'OBJECT_TYPE_TYPE_PEDESTRIAN_LEVEL_2/APH',
93 | 'OBJECT_TYPE_TYPE_CYCLIST_LEVEL_2/APH'
94 | ]
95 | }
96 |
97 |
98 |
99 | OPTIMIZATION:
100 | BATCH_SIZE_PER_GPU: 4
101 | NUM_EPOCHS: 30
102 |
103 | OPTIMIZER: adam_onecycle
104 | LR: 0.003
105 | WEIGHT_DECAY: 0.01
106 | MOMENTUM: 0.9
107 |
108 | MOMS: [0.95, 0.85]
109 | PCT_START: 0.4
110 | DIV_FACTOR: 10
111 | DECAY_STEP_LIST: [35, 45]
112 | LR_DECAY: 0.1
113 | LR_CLIP: 0.0000001
114 |
115 | LR_WARMUP: False
116 | WARMUP_EPOCH: 1
117 |
118 | GRAD_NORM_CLIP: 10
119 |
120 | REMAP_PRETRAIN:
121 | ENABLED: False
122 | WAY: BN_SCALE
123 | BN_SCALE:
124 | ABS: True
125 | OFA:
126 | l1_norm: max
127 |
--------------------------------------------------------------------------------
/tools/cfgs/waymo_models/centerpoint_pillar_1x.yaml:
--------------------------------------------------------------------------------
1 | CLASS_NAMES: ['Vehicle', 'Pedestrian', 'Cyclist']
2 |
3 | DATA_CONFIG:
4 | _BASE_CONFIG_: cfgs/dataset_configs/waymo_dataset.yaml
5 |
6 | POINT_CLOUD_RANGE: [-74.88, -74.88, -2, 74.88, 74.88, 4.0]
7 | DATA_PROCESSOR:
8 | - NAME: mask_points_and_boxes_outside_range
9 | REMOVE_OUTSIDE_BOXES: True
10 |
11 | - NAME: shuffle_points
12 | SHUFFLE_ENABLED: {
13 | 'train': True,
14 | 'test': True
15 | }
16 |
17 | - NAME: transform_points_to_voxels
18 | VOXEL_SIZE: [ 0.32, 0.32, 6.0 ]
19 | MAX_POINTS_PER_VOXEL: 20
20 | MAX_NUMBER_OF_VOXELS: {
21 | 'train': 150000,
22 | 'test': 150000
23 | }
24 |
25 |
26 | MODEL:
27 | NAME: CenterPoint
28 |
29 | VFE:
30 | NAME: PillarVFE
31 | WITH_DISTANCE: False
32 | USE_ABSLOTE_XYZ: True
33 | USE_NORM: True
34 | NUM_FILTERS: [ 64, 64 ]
35 |
36 | MAP_TO_BEV:
37 | NAME: PointPillarScatter
38 | NUM_BEV_FEATURES: 64
39 |
40 | BACKBONE_2D:
41 | NAME: BaseBEVBackbone
42 | ACT_FN: ReLU
43 | LAYER_NUMS: [ 3, 5, 5 ]
44 | LAYER_STRIDES: [ 1, 2, 2 ]
45 | NUM_FILTERS: [ 64, 128, 256 ]
46 | UPSAMPLE_STRIDES: [ 1, 2, 4 ]
47 | NUM_UPSAMPLE_FILTERS: [ 128, 128, 128 ]
48 |
49 | DENSE_HEAD:
50 | NAME: CenterHead
51 | CLASS_AGNOSTIC: False
52 |
53 | CLASS_NAMES_EACH_HEAD: [
54 | ['Vehicle', 'Pedestrian', 'Cyclist']
55 | ]
56 |
57 | SHARED_CONV_CHANNEL: 64
58 | USE_BIAS_BEFORE_NORM: True
59 | NUM_HM_CONV: 2
60 | SEPARATE_HEAD_CFG:
61 | HEAD_ORDER: ['center', 'center_z', 'dim', 'rot']
62 | HEAD_DICT: {
63 | 'center': {'out_channels': 2, 'num_conv': 2},
64 | 'center_z': {'out_channels': 1, 'num_conv': 2},
65 | 'dim': {'out_channels': 3, 'num_conv': 2},
66 | 'rot': {'out_channels': 2, 'num_conv': 2},
67 | }
68 |
69 | TARGET_ASSIGNER_CONFIG:
70 | FEATURE_MAP_STRIDE: 1
71 | NUM_MAX_OBJS: 500
72 | GAUSSIAN_OVERLAP: 0.1
73 | MIN_RADIUS: 2
74 |
75 | LOSS_CONFIG:
76 | LOSS_WEIGHTS: {
77 | 'cls_weight': 1.0,
78 | 'loc_weight': 2.0,
79 | 'code_weights': [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0]
80 | }
81 |
82 | POST_PROCESSING:
83 | SCORE_THRESH: 0.1
84 | POST_CENTER_LIMIT_RANGE: [-80, -80, -10.0, 80, 80, 10.0]
85 | MAX_OBJ_PER_SAMPLE: 500
86 | NMS_CONFIG:
87 | NMS_TYPE: nms_gpu
88 | NMS_THRESH: 0.7
89 | NMS_PRE_MAXSIZE: 4096
90 | NMS_POST_MAXSIZE: 500
91 |
92 | POST_PROCESSING:
93 | RECALL_THRESH_LIST: [0.3, 0.5, 0.7]
94 |
95 | EVAL_METRIC: waymo
96 | EVAL_CLASSES: {
97 | 'LEVEL_2/AP': [ 'OBJECT_TYPE_TYPE_VEHICLE_LEVEL_2/AP',
98 | 'OBJECT_TYPE_TYPE_PEDESTRIAN_LEVEL_2/AP',
99 | 'OBJECT_TYPE_TYPE_CYCLIST_LEVEL_2/AP'
100 | ],
101 | 'LEVEL_2/APH': [ 'OBJECT_TYPE_TYPE_VEHICLE_LEVEL_2/APH',
102 | 'OBJECT_TYPE_TYPE_PEDESTRIAN_LEVEL_2/APH',
103 | 'OBJECT_TYPE_TYPE_CYCLIST_LEVEL_2/APH'
104 | ]
105 | }
106 |
107 |
108 | OPTIMIZATION:
109 | BATCH_SIZE_PER_GPU: 2
110 | NUM_EPOCHS: 30
111 |
112 | OPTIMIZER: adam_onecycle
113 | LR: 0.003
114 | WEIGHT_DECAY: 0.01
115 | MOMENTUM: 0.9
116 |
117 | MOMS: [0.95, 0.85]
118 | PCT_START: 0.4
119 | DIV_FACTOR: 10
120 | DECAY_STEP_LIST: [35, 45]
121 | LR_DECAY: 0.1
122 | LR_CLIP: 0.0000001
123 |
124 | LR_WARMUP: False
125 | WARMUP_EPOCH: 1
126 |
127 | GRAD_NORM_CLIP: 10
128 |
--------------------------------------------------------------------------------
/tools/cfgs/waymo_models/centerpoint_without_resnet.yaml:
--------------------------------------------------------------------------------
1 | CLASS_NAMES: ['Vehicle', 'Pedestrian', 'Cyclist']
2 |
3 | DATA_CONFIG:
4 | _BASE_CONFIG_: cfgs/dataset_configs/waymo_dataset.yaml
5 |
6 | MODEL:
7 | NAME: CenterPoint
8 |
9 | VFE:
10 | NAME: MeanVFE
11 |
12 | BACKBONE_3D:
13 | NAME: VoxelBackBone8x
14 | # [input, conv1, conv2, conv3, conv4, output]
15 | NUM_FILTERS: [ 16, 16, 32, 64, 64, 128 ]
16 | LAYER_NUMS: [ 1, 1, 3, 3, 3, 1 ]
17 |
18 | MAP_TO_BEV:
19 | NAME: HeightCompression
20 | NUM_BEV_FEATURES: 256
21 |
22 | BACKBONE_2D:
23 | NAME: BaseBEVBackbone
24 |
25 | LAYER_NUMS: [5, 5]
26 | LAYER_STRIDES: [1, 2]
27 | NUM_FILTERS: [128, 256]
28 | UPSAMPLE_STRIDES: [1, 2]
29 | NUM_UPSAMPLE_FILTERS: [256, 256]
30 |
31 | DENSE_HEAD:
32 | NAME: CenterHead
33 | CLASS_AGNOSTIC: False
34 |
35 | CLASS_NAMES_EACH_HEAD: [
36 | ['Vehicle', 'Pedestrian', 'Cyclist']
37 | ]
38 |
39 | SHARED_CONV_CHANNEL: 64
40 | USE_BIAS_BEFORE_NORM: True
41 | NUM_HM_CONV: 2
42 | SEPARATE_HEAD_CFG:
43 | HEAD_ORDER: ['center', 'center_z', 'dim', 'rot']
44 | HEAD_DICT: {
45 | 'center': {'out_channels': 2, 'num_conv': 2},
46 | 'center_z': {'out_channels': 1, 'num_conv': 2},
47 | 'dim': {'out_channels': 3, 'num_conv': 2},
48 | 'rot': {'out_channels': 2, 'num_conv': 2},
49 | }
50 |
51 | TARGET_ASSIGNER_CONFIG:
52 | FEATURE_MAP_STRIDE: 8
53 | NUM_MAX_OBJS: 500
54 | GAUSSIAN_OVERLAP: 0.1
55 | MIN_RADIUS: 2
56 |
57 | LOSS_CONFIG:
58 | LOSS_WEIGHTS: {
59 | 'cls_weight': 1.0,
60 | 'loc_weight': 2.0,
61 | 'code_weights': [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0]
62 | }
63 |
64 | POST_PROCESSING:
65 | SCORE_THRESH: 0.1
66 | POST_CENTER_LIMIT_RANGE: [-75.2, -75.2, -2, 75.2, 75.2, 4]
67 | MAX_OBJ_PER_SAMPLE: 500
68 | NMS_CONFIG:
69 | NMS_TYPE: nms_gpu
70 | NMS_THRESH: 0.7
71 | NMS_PRE_MAXSIZE: 4096
72 | NMS_POST_MAXSIZE: 500
73 |
74 | POST_PROCESSING:
75 | RECALL_THRESH_LIST: [0.3, 0.5, 0.7]
76 |
77 | EVAL_METRIC: waymo
78 | EVAL_CLASSES: {
79 | 'LEVEL_2/AP': [ 'OBJECT_TYPE_TYPE_VEHICLE_LEVEL_2/AP',
80 | 'OBJECT_TYPE_TYPE_PEDESTRIAN_LEVEL_2/AP',
81 | 'OBJECT_TYPE_TYPE_CYCLIST_LEVEL_2/AP'
82 | ],
83 | 'LEVEL_2/APH': [ 'OBJECT_TYPE_TYPE_VEHICLE_LEVEL_2/APH',
84 | 'OBJECT_TYPE_TYPE_PEDESTRIAN_LEVEL_2/APH',
85 | 'OBJECT_TYPE_TYPE_CYCLIST_LEVEL_2/APH'
86 | ]
87 | }
88 |
89 | OPTIMIZATION:
90 | BATCH_SIZE_PER_GPU: 4
91 | NUM_EPOCHS: 30
92 |
93 | OPTIMIZER: adam_onecycle
94 | LR: 0.003
95 | WEIGHT_DECAY: 0.01
96 | MOMENTUM: 0.9
97 |
98 | MOMS: [0.95, 0.85]
99 | PCT_START: 0.4
100 | DIV_FACTOR: 10
101 | DECAY_STEP_LIST: [35, 45]
102 | LR_DECAY: 0.1
103 | LR_CLIP: 0.0000001
104 |
105 | LR_WARMUP: False
106 | WARMUP_EPOCH: 1
107 |
108 | GRAD_NORM_CLIP: 10
109 |
--------------------------------------------------------------------------------
/tools/cfgs/waymo_models/cp-pillar/cp-pillar-v0.48.yaml:
--------------------------------------------------------------------------------
1 | CLASS_NAMES: ['Vehicle', 'Pedestrian', 'Cyclist']
2 |
3 | DATA_CONFIG:
4 | _BASE_CONFIG_: cfgs/dataset_configs/waymo_dataset.yaml
5 |
6 | POINT_CLOUD_RANGE: [-74.88, -74.88, -2, 74.88, 74.88, 4.0]
7 | DATA_PROCESSOR:
8 | - NAME: mask_points_and_boxes_outside_range
9 | REMOVE_OUTSIDE_BOXES: True
10 |
11 | - NAME: shuffle_points
12 | SHUFFLE_ENABLED: {
13 | 'train': True,
14 | 'test': True
15 | }
16 |
17 | - NAME: transform_points_to_voxels
18 | VOXEL_SIZE: [ 0.48, 0.48, 6.0 ]
19 | MAX_POINTS_PER_VOXEL: 36
20 | MAX_NUMBER_OF_VOXELS: {
21 | 'train': 150000,
22 | 'test': 150000
23 | }
24 |
25 | MODEL:
26 | NAME: CenterPoint
27 |
28 | VFE:
29 | NAME: PillarVFE
30 | WITH_DISTANCE: False
31 | USE_ABSLOTE_XYZ: True
32 | USE_NORM: True
33 | NUM_FILTERS: [ 64, 64 ]
34 |
35 | MAP_TO_BEV:
36 | NAME: PointPillarScatter
37 | NUM_BEV_FEATURES: 64
38 |
39 | BACKBONE_2D:
40 | NAME: BaseBEVBackbone
41 | LAYER_NUMS: [ 3, 5, 5 ]
42 | LAYER_STRIDES: [ 1, 2, 2 ]
43 | NUM_FILTERS: [ 64, 128, 256 ]
44 | UPSAMPLE_STRIDES: [ 1, 2, 4 ]
45 | NUM_UPSAMPLE_FILTERS: [ 128, 128, 128 ]
46 | FOCUS: False
47 | ACT_FN: ReLU
48 |
49 | DENSE_HEAD:
50 | NAME: CenterHead
51 | CLASS_AGNOSTIC: False
52 |
53 | CLASS_NAMES_EACH_HEAD: [
54 | ['Vehicle', 'Pedestrian', 'Cyclist']
55 | ]
56 |
57 | SHARED_CONV_CHANNEL: 64
58 | USE_BIAS_BEFORE_NORM: True
59 | NUM_HM_CONV: 2
60 | SEPARATE_HEAD_CFG:
61 | HEAD_ORDER: ['center', 'center_z', 'dim', 'rot']
62 | HEAD_DICT: {
63 | 'center': {'out_channels': 2, 'num_conv': 2},
64 | 'center_z': {'out_channels': 1, 'num_conv': 2},
65 | 'dim': {'out_channels': 3, 'num_conv': 2},
66 | 'rot': {'out_channels': 2, 'num_conv': 2},
67 | }
68 |
69 | TARGET_ASSIGNER_CONFIG:
70 | FEATURE_MAP_STRIDE: 1
71 | NUM_MAX_OBJS: 500
72 | GAUSSIAN_OVERLAP: 0.1
73 | MIN_RADIUS: 2
74 | SHARPER: False
75 |
76 | LOSS_CONFIG:
77 | LOSS_WEIGHTS: {
78 | 'cls_weight': 1.0,
79 | 'loc_weight': 2.0,
80 | 'code_weights': [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0]
81 | }
82 |
83 | POST_PROCESSING:
84 | SCORE_THRESH: 0.1
85 | POST_CENTER_LIMIT_RANGE: [-80, -80, -10.0, 80, 80, 10.0]
86 | MAX_OBJ_PER_SAMPLE: 500
87 | NMS_CONFIG:
88 | NMS_TYPE: nms_gpu
89 | NMS_THRESH: 0.7
90 | NMS_PRE_MAXSIZE: 4096
91 | NMS_POST_MAXSIZE: 500
92 |
93 | POST_PROCESSING:
94 | RECALL_THRESH_LIST: [0.3, 0.5, 0.7]
95 |
96 | EVAL_METRIC: waymo
97 | EVAL_CLASSES: {
98 | 'LEVEL_2/AP': [ 'OBJECT_TYPE_TYPE_VEHICLE_LEVEL_2/AP',
99 | 'OBJECT_TYPE_TYPE_PEDESTRIAN_LEVEL_2/AP',
100 | 'OBJECT_TYPE_TYPE_CYCLIST_LEVEL_2/AP'
101 | ],
102 | 'LEVEL_2/APH': [ 'OBJECT_TYPE_TYPE_VEHICLE_LEVEL_2/APH',
103 | 'OBJECT_TYPE_TYPE_PEDESTRIAN_LEVEL_2/APH',
104 | 'OBJECT_TYPE_TYPE_CYCLIST_LEVEL_2/APH'
105 | ]
106 | }
107 |
108 |
109 | OPTIMIZATION:
110 | BATCH_SIZE_PER_GPU: 2
111 | NUM_EPOCHS: 30
112 |
113 | OPTIMIZER: adam_onecycle
114 | LR: 0.003
115 | WEIGHT_DECAY: 0.01
116 | MOMENTUM: 0.9
117 |
118 | MOMS: [0.95, 0.85]
119 | PCT_START: 0.4
120 | DIV_FACTOR: 10
121 | DECAY_STEP_LIST: [35, 45]
122 | LR_DECAY: 0.1
123 | LR_CLIP: 0.0000001
124 |
125 | LR_WARMUP: False
126 | WARMUP_EPOCH: 1
127 |
128 | GRAD_NORM_CLIP: 10
129 |
--------------------------------------------------------------------------------
/tools/cfgs/waymo_models/cp-pillar/cp-pillar.yaml:
--------------------------------------------------------------------------------
1 | CLASS_NAMES: ['Vehicle', 'Pedestrian', 'Cyclist']
2 |
3 | DATA_CONFIG:
4 | _BASE_CONFIG_: cfgs/dataset_configs/waymo_dataset.yaml
5 |
6 | POINT_CLOUD_RANGE: [-74.24, -74.24, -2, 74.24, 74.24, 4.0]
7 | DATA_PROCESSOR:
8 | - NAME: mask_points_and_boxes_outside_range
9 | REMOVE_OUTSIDE_BOXES: True
10 |
11 | - NAME: shuffle_points
12 | SHUFFLE_ENABLED: {
13 | 'train': True,
14 | 'test': True
15 | }
16 |
17 | - NAME: transform_points_to_voxels
18 | VOXEL_SIZE: [ 0.32, 0.32, 6.0 ]
19 | MAX_POINTS_PER_VOXEL: 20
20 | MAX_NUMBER_OF_VOXELS: {
21 | 'train': 150000,
22 | 'test': 150000
23 | }
24 |
25 |
26 | MODEL:
27 | NAME: CenterPoint
28 |
29 | VFE:
30 | NAME: PillarVFE
31 | WITH_DISTANCE: False
32 | USE_ABSLOTE_XYZ: True
33 | USE_NORM: True
34 | NUM_FILTERS: [ 64, 64 ]
35 |
36 | MAP_TO_BEV:
37 | NAME: PointPillarScatter
38 | NUM_BEV_FEATURES: 64
39 |
40 | BACKBONE_2D:
41 | NAME: BaseBEVBackbone
42 | WIDTH: 1.0
43 |
44 | LAYER_NUMS: [ 3, 5, 5 ]
45 | LAYER_STRIDES: [ 1, 2, 2 ]
46 | NUM_FILTERS: [ 64, 128, 256 ]
47 | UPSAMPLE_STRIDES: [ 1, 2, 4 ]
48 | NUM_UPSAMPLE_FILTERS: [ 128, 128, 128 ]
49 |
50 | DENSE_HEAD:
51 | NAME: CenterHead
52 | CLASS_AGNOSTIC: False
53 |
54 | CLASS_NAMES_EACH_HEAD: [
55 | ['Vehicle', 'Pedestrian', 'Cyclist']
56 | ]
57 |
58 | SHARED_CONV_CHANNEL: 64
59 | USE_BIAS_BEFORE_NORM: True
60 | NUM_HM_CONV: 2
61 | SEPARATE_HEAD_CFG:
62 | HEAD_ORDER: ['center', 'center_z', 'dim', 'rot']
63 | HEAD_DICT: {
64 | 'center': {'out_channels': 2, 'num_conv': 2},
65 | 'center_z': {'out_channels': 1, 'num_conv': 2},
66 | 'dim': {'out_channels': 3, 'num_conv': 2},
67 | 'rot': {'out_channels': 2, 'num_conv': 2},
68 | }
69 |
70 | TARGET_ASSIGNER_CONFIG:
71 | FEATURE_MAP_STRIDE: 1
72 | NUM_MAX_OBJS: 500
73 | GAUSSIAN_OVERLAP: 0.1
74 | MIN_RADIUS: 2
75 |
76 | LOSS_CONFIG:
77 | LOSS_WEIGHTS: {
78 | 'cls_weight': 1.0,
79 | 'loc_weight': 2.0,
80 | 'code_weights': [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0]
81 | }
82 |
83 | POST_PROCESSING:
84 | SCORE_THRESH: 0.1
85 | POST_CENTER_LIMIT_RANGE: [-80, -80, -10.0, 80, 80, 10.0]
86 | MAX_OBJ_PER_SAMPLE: 500
87 | NMS_CONFIG:
88 | NMS_TYPE: nms_gpu
89 | NMS_THRESH: 0.7
90 | NMS_PRE_MAXSIZE: 4096
91 | NMS_POST_MAXSIZE: 500
92 |
93 | POST_PROCESSING:
94 | RECALL_THRESH_LIST: [0.3, 0.5, 0.7]
95 |
96 | EVAL_METRIC: waymo
97 | EVAL_CLASSES: {
98 | 'LEVEL_2/AP': [ 'OBJECT_TYPE_TYPE_VEHICLE_LEVEL_2/AP',
99 | 'OBJECT_TYPE_TYPE_PEDESTRIAN_LEVEL_2/AP',
100 | 'OBJECT_TYPE_TYPE_CYCLIST_LEVEL_2/AP'
101 | ],
102 | 'LEVEL_2/APH': [ 'OBJECT_TYPE_TYPE_VEHICLE_LEVEL_2/APH',
103 | 'OBJECT_TYPE_TYPE_PEDESTRIAN_LEVEL_2/APH',
104 | 'OBJECT_TYPE_TYPE_CYCLIST_LEVEL_2/APH'
105 | ]
106 | }
107 |
108 |
109 | OPTIMIZATION:
110 | BATCH_SIZE_PER_GPU: 2
111 | NUM_EPOCHS: 30
112 |
113 | OPTIMIZER: adam_onecycle
114 | LR: 0.003
115 | WEIGHT_DECAY: 0.01
116 | MOMENTUM: 0.9
117 |
118 | MOMS: [0.95, 0.85]
119 | PCT_START: 0.4
120 | DIV_FACTOR: 10
121 | DECAY_STEP_LIST: [35, 45]
122 | LR_DECAY: 0.1
123 | LR_CLIP: 0.0000001
124 |
125 | LR_WARMUP: False
126 | WARMUP_EPOCH: 1
127 |
128 | GRAD_NORM_CLIP: 10
129 |
--------------------------------------------------------------------------------
/tools/cfgs/waymo_models/cp-voxel/cp-voxel-s.yaml:
--------------------------------------------------------------------------------
1 | CLASS_NAMES: ['Vehicle', 'Pedestrian', 'Cyclist']
2 |
3 | DATA_CONFIG:
4 | _BASE_CONFIG_: cfgs/dataset_configs/waymo_dataset.yaml
5 |
6 | MODEL:
7 | NAME: CenterPoint
8 | IGNORE_PRETRAIN_MODULES: ['placeholder']
9 |
10 | VFE:
11 | NAME: MeanVFE
12 |
13 | BACKBONE_3D:
14 | NAME: VoxelResBackBone8x
15 | ACT_FN: ReLU
16 | # [input, conv1, conv2, conv3, conv4, output]
17 | NUM_FILTERS: [16, 16, 32, 64, 128, 128]
18 | LAYER_NUMS: [1, 2, 3, 3, 3, 1]
19 | WIDTH: 1.0
20 |
21 | MAP_TO_BEV:
22 | NAME: HeightCompression
23 | NUM_BEV_FEATURES: 256
24 |
25 | BACKBONE_2D:
26 | NAME: BaseBEVBackbone
27 | ACT_FN: ReLU
28 | NORM_TYPE: BatchNorm2d
29 | WIDTH: 0.5
30 |
31 | LAYER_NUMS: [5, 5]
32 | LAYER_STRIDES: [1, 2]
33 | NUM_FILTERS: [128, 256]
34 | UPSAMPLE_STRIDES: [1, 2]
35 | NUM_UPSAMPLE_FILTERS: [256, 256]
36 |
37 | DENSE_HEAD:
38 | NAME: CenterHead
39 | CLASS_AGNOSTIC: False
40 | ACT_FN: ReLU
41 | NORM_TYPE: BatchNorm2d
42 |
43 | CLASS_NAMES_EACH_HEAD: [
44 | ['Vehicle', 'Pedestrian', 'Cyclist']
45 | ]
46 |
47 | SHARED_CONV_CHANNEL: 32
48 | USE_BIAS_BEFORE_NORM: True
49 | NUM_HM_CONV: 2
50 | SEPARATE_HEAD_CFG:
51 | HEAD_ORDER: ['center', 'center_z', 'dim', 'rot']
52 | HEAD_DICT: {
53 | 'center': {'out_channels': 2, 'num_conv': 2},
54 | 'center_z': {'out_channels': 1, 'num_conv': 2},
55 | 'dim': {'out_channels': 3, 'num_conv': 2},
56 | 'rot': {'out_channels': 2, 'num_conv': 2},
57 | }
58 |
59 | TARGET_ASSIGNER_CONFIG:
60 | FEATURE_MAP_STRIDE: 8
61 | NUM_MAX_OBJS: 500
62 | GAUSSIAN_OVERLAP: 0.1
63 | MIN_RADIUS: 2
64 |
65 | LOSS_CONFIG:
66 | LOSS_WEIGHTS: {
67 | 'cls_weight': 1.0,
68 | 'loc_weight': 2.0,
69 | 'code_weights': [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0]
70 | }
71 |
72 | POST_PROCESSING:
73 | SCORE_THRESH: 0.1
74 | POST_CENTER_LIMIT_RANGE: [-75.2, -75.2, -2, 75.2, 75.2, 4]
75 | MAX_OBJ_PER_SAMPLE: 500
76 | NMS_CONFIG:
77 | NMS_TYPE: nms_gpu
78 | NMS_THRESH: 0.7
79 | NMS_PRE_MAXSIZE: 4096
80 | NMS_POST_MAXSIZE: 500
81 |
82 | POST_PROCESSING:
83 | RECALL_THRESH_LIST: [0.3, 0.5, 0.7]
84 |
85 | EVAL_METRIC: waymo
86 | EVAL_CLASSES: {
87 | 'LEVEL_2/AP': [ 'OBJECT_TYPE_TYPE_VEHICLE_LEVEL_2/AP',
88 | 'OBJECT_TYPE_TYPE_PEDESTRIAN_LEVEL_2/AP',
89 | 'OBJECT_TYPE_TYPE_CYCLIST_LEVEL_2/AP'
90 | ],
91 | 'LEVEL_2/APH': [ 'OBJECT_TYPE_TYPE_VEHICLE_LEVEL_2/APH',
92 | 'OBJECT_TYPE_TYPE_PEDESTRIAN_LEVEL_2/APH',
93 | 'OBJECT_TYPE_TYPE_CYCLIST_LEVEL_2/APH'
94 | ]
95 | }
96 |
97 |
98 |
99 | OPTIMIZATION:
100 | BATCH_SIZE_PER_GPU: 4
101 | NUM_EPOCHS: 30
102 |
103 | OPTIMIZER: adam_onecycle
104 | LR: 0.003
105 | WEIGHT_DECAY: 0.01
106 | MOMENTUM: 0.9
107 |
108 | MOMS: [0.95, 0.85]
109 | PCT_START: 0.4
110 | DIV_FACTOR: 10
111 | DECAY_STEP_LIST: [35, 45]
112 | LR_DECAY: 0.1
113 | LR_CLIP: 0.0000001
114 |
115 | LR_WARMUP: False
116 | WARMUP_EPOCH: 1
117 |
118 | GRAD_NORM_CLIP: 10
119 |
120 | REMAP_PRETRAIN:
121 | ENABLED: False
122 | WAY: BN_SCALE
123 | BN_SCALE:
124 | ABS: True
125 | OFA:
126 | l1_norm: max
127 |
--------------------------------------------------------------------------------
/tools/cfgs/waymo_models/cp-voxel/cp-voxel-xs.yaml:
--------------------------------------------------------------------------------
1 | CLASS_NAMES: ['Vehicle', 'Pedestrian', 'Cyclist']
2 |
3 | DATA_CONFIG:
4 | _BASE_CONFIG_: cfgs/dataset_configs/waymo_dataset.yaml
5 |
6 | MODEL:
7 | NAME: CenterPoint
8 | IGNORE_PRETRAIN_MODULES: ['placeholder']
9 |
10 | VFE:
11 | NAME: MeanVFE
12 |
13 | BACKBONE_3D:
14 | NAME: VoxelResBackBone8x
15 | ACT_FN: ReLU
16 | # [input, conv1, conv2, conv3, conv4, output]
17 | NUM_FILTERS: [16, 16, 32, 64, 128, 128]
18 | LAYER_NUMS: [1, 2, 3, 3, 3, 1]
19 | WIDTH: 0.75
20 |
21 | MAP_TO_BEV:
22 | NAME: HeightCompression
23 | NUM_BEV_FEATURES: 192
24 |
25 | BACKBONE_2D:
26 | NAME: BaseBEVBackbone
27 | ACT_FN: ReLU
28 | NORM_TYPE: BatchNorm2d
29 | WIDTH: 0.5
30 |
31 | LAYER_NUMS: [5, 5]
32 | LAYER_STRIDES: [1, 2]
33 | NUM_FILTERS: [128, 256]
34 | UPSAMPLE_STRIDES: [1, 2]
35 | NUM_UPSAMPLE_FILTERS: [256, 256]
36 |
37 | DENSE_HEAD:
38 | NAME: CenterHead
39 | CLASS_AGNOSTIC: False
40 | ACT_FN: ReLU
41 | NORM_TYPE: BatchNorm2d
42 |
43 | CLASS_NAMES_EACH_HEAD: [
44 | ['Vehicle', 'Pedestrian', 'Cyclist']
45 | ]
46 |
47 | SHARED_CONV_CHANNEL: 32
48 | USE_BIAS_BEFORE_NORM: True
49 | NUM_HM_CONV: 2
50 | SEPARATE_HEAD_CFG:
51 | HEAD_ORDER: ['center', 'center_z', 'dim', 'rot']
52 | HEAD_DICT: {
53 | 'center': {'out_channels': 2, 'num_conv': 2},
54 | 'center_z': {'out_channels': 1, 'num_conv': 2},
55 | 'dim': {'out_channels': 3, 'num_conv': 2},
56 | 'rot': {'out_channels': 2, 'num_conv': 2},
57 | }
58 |
59 | TARGET_ASSIGNER_CONFIG:
60 | FEATURE_MAP_STRIDE: 8
61 | NUM_MAX_OBJS: 500
62 | GAUSSIAN_OVERLAP: 0.1
63 | MIN_RADIUS: 2
64 |
65 | LOSS_CONFIG:
66 | LOSS_WEIGHTS: {
67 | 'cls_weight': 1.0,
68 | 'loc_weight': 2.0,
69 | 'code_weights': [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0]
70 | }
71 |
72 | POST_PROCESSING:
73 | SCORE_THRESH: 0.1
74 | POST_CENTER_LIMIT_RANGE: [-75.2, -75.2, -2, 75.2, 75.2, 4]
75 | MAX_OBJ_PER_SAMPLE: 500
76 | NMS_CONFIG:
77 | NMS_TYPE: nms_gpu
78 | NMS_THRESH: 0.7
79 | NMS_PRE_MAXSIZE: 4096
80 | NMS_POST_MAXSIZE: 500
81 |
82 | POST_PROCESSING:
83 | RECALL_THRESH_LIST: [0.3, 0.5, 0.7]
84 |
85 | EVAL_METRIC: waymo
86 | EVAL_CLASSES: {
87 | 'LEVEL_2/AP': [ 'OBJECT_TYPE_TYPE_VEHICLE_LEVEL_2/AP',
88 | 'OBJECT_TYPE_TYPE_PEDESTRIAN_LEVEL_2/AP',
89 | 'OBJECT_TYPE_TYPE_CYCLIST_LEVEL_2/AP'
90 | ],
91 | 'LEVEL_2/APH': [ 'OBJECT_TYPE_TYPE_VEHICLE_LEVEL_2/APH',
92 | 'OBJECT_TYPE_TYPE_PEDESTRIAN_LEVEL_2/APH',
93 | 'OBJECT_TYPE_TYPE_CYCLIST_LEVEL_2/APH'
94 | ]
95 | }
96 |
97 |
98 |
99 | OPTIMIZATION:
100 | BATCH_SIZE_PER_GPU: 4
101 | NUM_EPOCHS: 30
102 |
103 | OPTIMIZER: adam_onecycle
104 | LR: 0.003
105 | WEIGHT_DECAY: 0.01
106 | MOMENTUM: 0.9
107 |
108 | MOMS: [0.95, 0.85]
109 | PCT_START: 0.4
110 | DIV_FACTOR: 10
111 | DECAY_STEP_LIST: [35, 45]
112 | LR_DECAY: 0.1
113 | LR_CLIP: 0.0000001
114 |
115 | LR_WARMUP: False
116 | WARMUP_EPOCH: 1
117 |
118 | GRAD_NORM_CLIP: 10
119 |
120 | REMAP_PRETRAIN:
121 | ENABLED: False
122 | WAY: BN_SCALE
123 | BN_SCALE:
124 | ABS: True
125 | OFA:
126 | l1_norm: max
127 |
--------------------------------------------------------------------------------
/tools/cfgs/waymo_models/cp-voxel/cp-voxel-xxs.yaml:
--------------------------------------------------------------------------------
1 | CLASS_NAMES: ['Vehicle', 'Pedestrian', 'Cyclist']
2 |
3 | DATA_CONFIG:
4 | _BASE_CONFIG_: cfgs/dataset_configs/waymo_dataset.yaml
5 |
6 | MODEL:
7 | NAME: CenterPoint
8 | IGNORE_PRETRAIN_MODULES: ['placeholder']
9 |
10 | VFE:
11 | NAME: MeanVFE
12 |
13 | BACKBONE_3D:
14 | NAME: VoxelResBackBone8x
15 | ACT_FN: ReLU
16 | # [input, conv1, conv2, conv3, conv4, output]
17 | NUM_FILTERS: [16, 16, 32, 64, 128, 128]
18 | LAYER_NUMS: [1, 2, 3, 3, 3, 1]
19 | WIDTH: 0.5
20 |
21 | MAP_TO_BEV:
22 | NAME: HeightCompression
23 | NUM_BEV_FEATURES: 128
24 |
25 | BACKBONE_2D:
26 | NAME: BaseBEVBackbone
27 | ACT_FN: ReLU
28 | NORM_TYPE: BatchNorm2d
29 | WIDTH: 0.25
30 |
31 | LAYER_NUMS: [5, 5]
32 | LAYER_STRIDES: [1, 2]
33 | NUM_FILTERS: [128, 256]
34 | UPSAMPLE_STRIDES: [1, 2]
35 | NUM_UPSAMPLE_FILTERS: [256, 256]
36 |
37 | DENSE_HEAD:
38 | NAME: CenterHead
39 | CLASS_AGNOSTIC: False
40 | ACT_FN: ReLU
41 | NORM_TYPE: BatchNorm2d
42 |
43 | CLASS_NAMES_EACH_HEAD: [
44 | ['Vehicle', 'Pedestrian', 'Cyclist']
45 | ]
46 |
47 | SHARED_CONV_CHANNEL: 16
48 | USE_BIAS_BEFORE_NORM: True
49 | NUM_HM_CONV: 2
50 | SEPARATE_HEAD_CFG:
51 | HEAD_ORDER: ['center', 'center_z', 'dim', 'rot']
52 | HEAD_DICT: {
53 | 'center': {'out_channels': 2, 'num_conv': 2},
54 | 'center_z': {'out_channels': 1, 'num_conv': 2},
55 | 'dim': {'out_channels': 3, 'num_conv': 2},
56 | 'rot': {'out_channels': 2, 'num_conv': 2},
57 | }
58 |
59 | TARGET_ASSIGNER_CONFIG:
60 | FEATURE_MAP_STRIDE: 8
61 | NUM_MAX_OBJS: 500
62 | GAUSSIAN_OVERLAP: 0.1
63 | MIN_RADIUS: 2
64 |
65 | LOSS_CONFIG:
66 | LOSS_WEIGHTS: {
67 | 'cls_weight': 1.0,
68 | 'loc_weight': 2.0,
69 | 'code_weights': [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0]
70 | }
71 |
72 | POST_PROCESSING:
73 | SCORE_THRESH: 0.1
74 | POST_CENTER_LIMIT_RANGE: [-75.2, -75.2, -2, 75.2, 75.2, 4]
75 | MAX_OBJ_PER_SAMPLE: 500
76 | NMS_CONFIG:
77 | NMS_TYPE: nms_gpu
78 | NMS_THRESH: 0.7
79 | NMS_PRE_MAXSIZE: 4096
80 | NMS_POST_MAXSIZE: 500
81 |
82 | POST_PROCESSING:
83 | RECALL_THRESH_LIST: [0.3, 0.5, 0.7]
84 |
85 | EVAL_METRIC: waymo
86 | EVAL_CLASSES: {
87 | 'LEVEL_2/AP': [ 'OBJECT_TYPE_TYPE_VEHICLE_LEVEL_2/AP',
88 | 'OBJECT_TYPE_TYPE_PEDESTRIAN_LEVEL_2/AP',
89 | 'OBJECT_TYPE_TYPE_CYCLIST_LEVEL_2/AP'
90 | ],
91 | 'LEVEL_2/APH': [ 'OBJECT_TYPE_TYPE_VEHICLE_LEVEL_2/APH',
92 | 'OBJECT_TYPE_TYPE_PEDESTRIAN_LEVEL_2/APH',
93 | 'OBJECT_TYPE_TYPE_CYCLIST_LEVEL_2/APH'
94 | ]
95 | }
96 |
97 |
98 |
99 | OPTIMIZATION:
100 | BATCH_SIZE_PER_GPU: 4
101 | NUM_EPOCHS: 30
102 |
103 | OPTIMIZER: adam_onecycle
104 | LR: 0.003
105 | WEIGHT_DECAY: 0.01
106 | MOMENTUM: 0.9
107 |
108 | MOMS: [0.95, 0.85]
109 | PCT_START: 0.4
110 | DIV_FACTOR: 10
111 | DECAY_STEP_LIST: [35, 45]
112 | LR_DECAY: 0.1
113 | LR_CLIP: 0.0000001
114 |
115 | LR_WARMUP: False
116 | WARMUP_EPOCH: 1
117 |
118 | GRAD_NORM_CLIP: 10
119 |
120 | REMAP_PRETRAIN:
121 | ENABLED: False
122 | WAY: BN_SCALE
123 | BN_SCALE:
124 | ABS: True
125 | OFA:
126 | l1_norm: max
127 |
--------------------------------------------------------------------------------
/tools/cfgs/waymo_models/cp-voxel/cp-voxel.yaml:
--------------------------------------------------------------------------------
1 | CLASS_NAMES: ['Vehicle', 'Pedestrian', 'Cyclist']
2 |
3 | DATA_CONFIG:
4 | _BASE_CONFIG_: cfgs/dataset_configs/waymo_dataset.yaml
5 |
6 | MODEL:
7 | NAME: CenterPoint
8 | IGNORE_PRETRAIN_MODULES: ['placeholder']
9 |
10 | VFE:
11 | NAME: MeanVFE
12 |
13 | BACKBONE_3D:
14 | NAME: VoxelResBackBone8x
15 | ACT_FN: ReLU
16 | # [input, conv1, conv2, conv3, conv4, output]
17 | NUM_FILTERS: [16, 16, 32, 64, 128, 128]
18 | LAYER_NUMS: [1, 2, 3, 3, 3, 1]
19 | WIDTH: 1.0
20 |
21 | MAP_TO_BEV:
22 | NAME: HeightCompression
23 | NUM_BEV_FEATURES: 256
24 |
25 | BACKBONE_2D:
26 | NAME: BaseBEVBackbone
27 | ACT_FN: ReLU
28 | NORM_TYPE: BatchNorm2d
29 | WIDTH: 1.0
30 |
31 | LAYER_NUMS: [5, 5]
32 | LAYER_STRIDES: [1, 2]
33 | NUM_FILTERS: [128, 256]
34 | UPSAMPLE_STRIDES: [1, 2]
35 | NUM_UPSAMPLE_FILTERS: [256, 256]
36 |
37 | DENSE_HEAD:
38 | NAME: CenterHead
39 | CLASS_AGNOSTIC: False
40 | ACT_FN: ReLU
41 | NORM_TYPE: BatchNorm2d
42 |
43 | CLASS_NAMES_EACH_HEAD: [
44 | ['Vehicle', 'Pedestrian', 'Cyclist']
45 | ]
46 |
47 | SHARED_CONV_CHANNEL: 64
48 | USE_BIAS_BEFORE_NORM: True
49 | NUM_HM_CONV: 2
50 | SEPARATE_HEAD_CFG:
51 | HEAD_ORDER: ['center', 'center_z', 'dim', 'rot']
52 | HEAD_DICT: {
53 | 'center': {'out_channels': 2, 'num_conv': 2},
54 | 'center_z': {'out_channels': 1, 'num_conv': 2},
55 | 'dim': {'out_channels': 3, 'num_conv': 2},
56 | 'rot': {'out_channels': 2, 'num_conv': 2},
57 | }
58 |
59 | TARGET_ASSIGNER_CONFIG:
60 | FEATURE_MAP_STRIDE: 8
61 | NUM_MAX_OBJS: 500
62 | GAUSSIAN_OVERLAP: 0.1
63 | MIN_RADIUS: 2
64 |
65 | LOSS_CONFIG:
66 | LOSS_WEIGHTS: {
67 | 'cls_weight': 1.0,
68 | 'loc_weight': 2.0,
69 | 'code_weights': [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0]
70 | }
71 |
72 | POST_PROCESSING:
73 | SCORE_THRESH: 0.1
74 | POST_CENTER_LIMIT_RANGE: [-75.2, -75.2, -2, 75.2, 75.2, 4]
75 | MAX_OBJ_PER_SAMPLE: 500
76 | NMS_CONFIG:
77 | NMS_TYPE: nms_gpu
78 | NMS_THRESH: 0.7
79 | NMS_PRE_MAXSIZE: 4096
80 | NMS_POST_MAXSIZE: 500
81 |
82 | POST_PROCESSING:
83 | RECALL_THRESH_LIST: [0.3, 0.5, 0.7]
84 |
85 | EVAL_METRIC: waymo
86 | EVAL_CLASSES: {
87 | 'LEVEL_2/AP': [ 'OBJECT_TYPE_TYPE_VEHICLE_LEVEL_2/AP',
88 | 'OBJECT_TYPE_TYPE_PEDESTRIAN_LEVEL_2/AP',
89 | 'OBJECT_TYPE_TYPE_CYCLIST_LEVEL_2/AP'
90 | ],
91 | 'LEVEL_2/APH': [ 'OBJECT_TYPE_TYPE_VEHICLE_LEVEL_2/APH',
92 | 'OBJECT_TYPE_TYPE_PEDESTRIAN_LEVEL_2/APH',
93 | 'OBJECT_TYPE_TYPE_CYCLIST_LEVEL_2/APH'
94 | ]
95 | }
96 |
97 |
98 |
99 | OPTIMIZATION:
100 | BATCH_SIZE_PER_GPU: 4
101 | NUM_EPOCHS: 30
102 |
103 | OPTIMIZER: adam_onecycle
104 | LR: 0.003
105 | WEIGHT_DECAY: 0.01
106 | MOMENTUM: 0.9
107 |
108 | MOMS: [0.95, 0.85]
109 | PCT_START: 0.4
110 | DIV_FACTOR: 10
111 | DECAY_STEP_LIST: [35, 45]
112 | LR_DECAY: 0.1
113 | LR_CLIP: 0.0000001
114 |
115 | LR_WARMUP: False
116 | WARMUP_EPOCH: 1
117 |
118 | GRAD_NORM_CLIP: 10
119 |
120 | REMAP_PRETRAIN:
121 | ENABLED: False
122 | WAY: BN_SCALE
123 | BN_SCALE:
124 | ABS: True
125 | OFA:
126 | l1_norm: max
127 |
--------------------------------------------------------------------------------
/tools/cfgs/waymo_models/second.yaml:
--------------------------------------------------------------------------------
1 | CLASS_NAMES: ['Vehicle', 'Pedestrian', 'Cyclist']
2 |
3 | DATA_CONFIG:
4 | _BASE_CONFIG_: cfgs/dataset_configs/waymo_dataset.yaml
5 |
6 |
7 | MODEL:
8 | NAME: SECONDNet
9 |
10 | VFE:
11 | NAME: MeanVFE
12 |
13 | BACKBONE_3D:
14 | NAME: VoxelBackBone8x
15 |
16 | MAP_TO_BEV:
17 | NAME: HeightCompression
18 | NUM_BEV_FEATURES: 256
19 |
20 | BACKBONE_2D:
21 | NAME: BaseBEVBackbone
22 |
23 | LAYER_NUMS: [5, 5]
24 | LAYER_STRIDES: [1, 2]
25 | NUM_FILTERS: [128, 256]
26 | UPSAMPLE_STRIDES: [1, 2]
27 | NUM_UPSAMPLE_FILTERS: [256, 256]
28 |
29 | DENSE_HEAD:
30 | NAME: AnchorHeadSingle
31 | CLASS_AGNOSTIC: False
32 |
33 | USE_DIRECTION_CLASSIFIER: True
34 | DIR_OFFSET: 0.78539
35 | DIR_LIMIT_OFFSET: 0.0
36 | NUM_DIR_BINS: 2
37 |
38 | ANCHOR_GENERATOR_CONFIG: [
39 | {
40 | 'class_name': 'Vehicle',
41 | 'anchor_sizes': [[4.7, 2.1, 1.7]],
42 | 'anchor_rotations': [0, 1.57],
43 | 'anchor_bottom_heights': [0],
44 | 'align_center': False,
45 | 'feature_map_stride': 8,
46 | 'matched_threshold': 0.55,
47 | 'unmatched_threshold': 0.4
48 | },
49 | {
50 | 'class_name': 'Pedestrian',
51 | 'anchor_sizes': [[0.91, 0.86, 1.73]],
52 | 'anchor_rotations': [0, 1.57],
53 | 'anchor_bottom_heights': [0],
54 | 'align_center': False,
55 | 'feature_map_stride': 8,
56 | 'matched_threshold': 0.5,
57 | 'unmatched_threshold': 0.35
58 | },
59 | {
60 | 'class_name': 'Cyclist',
61 | 'anchor_sizes': [[1.78, 0.84, 1.78]],
62 | 'anchor_rotations': [0, 1.57],
63 | 'anchor_bottom_heights': [0],
64 | 'align_center': False,
65 | 'feature_map_stride': 8,
66 | 'matched_threshold': 0.5,
67 | 'unmatched_threshold': 0.35
68 | }
69 | ]
70 |
71 | TARGET_ASSIGNER_CONFIG:
72 | NAME: AxisAlignedTargetAssigner
73 | POS_FRACTION: -1.0
74 | SAMPLE_SIZE: 512
75 | NORM_BY_NUM_EXAMPLES: False
76 | MATCH_HEIGHT: False
77 | BOX_CODER: ResidualCoder
78 |
79 | LOSS_CONFIG:
80 | LOSS_WEIGHTS: {
81 | 'cls_weight': 1.0,
82 | 'loc_weight': 2.0,
83 | 'dir_weight': 0.2,
84 | 'code_weights': [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0]
85 | }
86 |
87 | POST_PROCESSING:
88 | RECALL_THRESH_LIST: [0.3, 0.5, 0.7]
89 | SCORE_THRESH: 0.1
90 | OUTPUT_RAW_SCORE: False
91 |
92 | EVAL_METRIC: waymo
93 |
94 | NMS_CONFIG:
95 | MULTI_CLASSES_NMS: False
96 | NMS_TYPE: nms_gpu
97 | NMS_THRESH: 0.7
98 | NMS_PRE_MAXSIZE: 4096
99 | NMS_POST_MAXSIZE: 500
100 |
101 |
102 | OPTIMIZATION:
103 | BATCH_SIZE_PER_GPU: 4
104 | NUM_EPOCHS: 30
105 |
106 | OPTIMIZER: adam_onecycle
107 | LR: 0.003
108 | WEIGHT_DECAY: 0.01
109 | MOMENTUM: 0.9
110 |
111 | MOMS: [0.95, 0.85]
112 | PCT_START: 0.4
113 | DIV_FACTOR: 10
114 | DECAY_STEP_LIST: [35, 45]
115 | LR_DECAY: 0.1
116 | LR_CLIP: 0.0000001
117 |
118 | LR_WARMUP: False
119 | WARMUP_EPOCH: 1
120 |
121 | GRAD_NORM_CLIP: 10
--------------------------------------------------------------------------------
/tools/scripts/dist_test.sh:
--------------------------------------------------------------------------------
1 | #!/usr/bin/env bash
2 |
3 | set -x
4 | NGPUS=$1
5 | PY_ARGS=${@:2}
6 |
7 | python -m torch.distributed.launch --nproc_per_node=${NGPUS} test.py --launcher pytorch ${PY_ARGS}
8 |
9 |
--------------------------------------------------------------------------------
/tools/scripts/dist_train.sh:
--------------------------------------------------------------------------------
1 | #!/usr/bin/env bash
2 |
3 | set -x
4 | NGPUS=$1
5 | PY_ARGS=${@:2}
6 |
7 | while true
8 | do
9 | PORT=$(( ((RANDOM<<15)|RANDOM) % 49152 + 10000 ))
10 | status="$(nc -z 127.0.0.1 $PORT < /dev/null &>/dev/null; echo $?)"
11 | if [ "${status}" != "0" ]; then
12 | break;
13 | fi
14 | done
15 | echo $PORT
16 |
17 | python -m torch.distributed.launch --nproc_per_node=${NGPUS} train.py --launcher pytorch ${PY_ARGS}
--------------------------------------------------------------------------------
/tools/scripts/slurm_test_mgpu.sh:
--------------------------------------------------------------------------------
1 | #!/usr/bin/env bash
2 |
3 | set -x
4 |
5 | PARTITION=$1
6 | GPUS=$2
7 | GPUS_PER_NODE=$GPUS
8 | PY_ARGS=${@:3}
9 | JOB_NAME=eval
10 | SRUN_ARGS=${SRUN_ARGS:-""}
11 |
12 | while true
13 | do
14 | PORT=$(( ((RANDOM<<15)|RANDOM) % 49152 + 10000 ))
15 | status="$(nc -z 127.0.0.1 $PORT < /dev/null &>/dev/null; echo $?)"
16 | if [ "${status}" != "0" ]; then
17 | break;
18 | fi
19 | done
20 | echo $PORT
21 |
22 | srun -p ${PARTITION} \
23 | --job-name=${JOB_NAME} \
24 | --gres=gpu:${GPUS_PER_NODE} \
25 | --ntasks=${GPUS} \
26 | --ntasks-per-node=${GPUS_PER_NODE} \
27 | --kill-on-bad-exit=1 \
28 | ${SRUN_ARGS} \
29 | python -u test.py --launcher slurm --tcp_port $PORT ${PY_ARGS}
30 |
31 |
--------------------------------------------------------------------------------
/tools/scripts/slurm_test_single.sh:
--------------------------------------------------------------------------------
1 | #!/usr/bin/env bash
2 |
3 | set -x
4 |
5 | PARTITION=$1
6 | GPUS=1
7 | GPUS_PER_NODE=1
8 | PY_ARGS=${@:2}
9 | JOB_NAME=eval
10 | SRUN_ARGS=${SRUN_ARGS:-""}
11 |
12 | srun -p ${PARTITION} \
13 | --job-name=${JOB_NAME} \
14 | --gres=gpu:${GPUS_PER_NODE} \
15 | --ntasks=${GPUS} \
16 | --ntasks-per-node=${GPUS_PER_NODE} \
17 | --kill-on-bad-exit=1 \
18 | ${SRUN_ARGS} \
19 | python -u test.py ${PY_ARGS}
20 |
--------------------------------------------------------------------------------
/tools/scripts/slurm_train.sh:
--------------------------------------------------------------------------------
1 | #!/usr/bin/env bash
2 |
3 | set -x
4 |
5 | PARTITION=$1
6 | JOB_NAME=$2
7 | GPUS=$3
8 | PY_ARGS=${@:4}
9 |
10 | GPUS_PER_NODE=${GPUS_PER_NODE:-8}
11 | CPUS_PER_TASK=${CPUS_PER_TASK:-8}
12 | SRUN_ARGS=${SRUN_ARGS:-""}
13 |
14 | while true
15 | do
16 | PORT=$(( ((RANDOM<<15)|RANDOM) % 49152 + 10000 ))
17 | status="$(nc -z 127.0.0.1 $PORT < /dev/null &>/dev/null; echo $?)"
18 | if [ "${status}" != "0" ]; then
19 | break;
20 | fi
21 | done
22 | echo $PORT
23 |
24 | srun -p ${PARTITION} \
25 | --job-name=${JOB_NAME} \
26 | --gres=gpu:${GPUS_PER_NODE} \
27 | --ntasks=${GPUS} \
28 | --ntasks-per-node=${GPUS_PER_NODE} \
29 | --cpus-per-task=${CPUS_PER_TASK} \
30 | --kill-on-bad-exit=1 \
31 | ${SRUN_ARGS} \
32 | python -u train.py --launcher slurm --tcp_port $PORT ${PY_ARGS}
33 |
--------------------------------------------------------------------------------
/tools/scripts/slurm_train_single.sh:
--------------------------------------------------------------------------------
1 | #!/usr/bin/env bash
2 |
3 | set -x
4 |
5 | PARTITION=$1
6 | JOB_NAME=$2
7 | GPUS=1
8 | PY_ARGS=${@:3}
9 |
10 | GPUS_PER_NODE=1
11 | SRUN_ARGS=${SRUN_ARGS:-""}
12 |
13 | srun -p ${PARTITION} \
14 | --job-name=${JOB_NAME} \
15 | --gres=gpu:${GPUS_PER_NODE} \
16 | --ntasks=${GPUS} \
17 | --ntasks-per-node=${GPUS_PER_NODE} \
18 | --kill-on-bad-exit=1 \
19 | ${SRUN_ARGS} \
20 | python -u train.py ${PY_ARGS}
21 |
--------------------------------------------------------------------------------
/tools/scripts/torch_train.sh:
--------------------------------------------------------------------------------
1 | #!/usr/bin/env bash
2 |
3 | set -x
4 | NGPUS=$1
5 | PY_ARGS=${@:2}
6 |
7 | while true
8 | do
9 | PORT=$(( ((RANDOM<<15)|RANDOM) % 49152 + 10000 ))
10 | status="$(nc -z 127.0.0.1 $PORT < /dev/null &>/dev/null; echo $?)"
11 | if [ "${status}" != "0" ]; then
12 | break;
13 | fi
14 | done
15 | echo $PORT
16 |
17 | torchrun --nproc_per_node=${NGPUS} --rdzv_endpoint=localhost:${PORT} train.py --launcher pytorch ${PY_ARGS}
18 |
--------------------------------------------------------------------------------
/tools/train_utils/optimization/__init__.py:
--------------------------------------------------------------------------------
1 | from functools import partial
2 |
3 | import torch.nn as nn
4 | import torch.optim as optim
5 | import torch.optim.lr_scheduler as lr_sched
6 |
7 | from .fastai_optim import OptimWrapper
8 | from .learning_schedules_fastai import CosineWarmupLR, OneCycle
9 |
10 |
11 | def build_optimizer(model, optim_cfg):
12 | if optim_cfg.OPTIMIZER == 'adam':
13 | optimizer = optim.Adam(model.parameters(), lr=optim_cfg.LR, weight_decay=optim_cfg.WEIGHT_DECAY)
14 | elif optim_cfg.OPTIMIZER == 'sgd':
15 | optimizer = optim.SGD(
16 | model.parameters(), lr=optim_cfg.LR, weight_decay=optim_cfg.WEIGHT_DECAY,
17 | momentum=optim_cfg.MOMENTUM
18 | )
19 | elif optim_cfg.OPTIMIZER == 'adam_onecycle':
20 | def children(m: nn.Module):
21 | return list(m.children())
22 |
23 | def num_children(m: nn.Module) -> int:
24 | return len(children(m))
25 |
26 | flatten_model = lambda m: sum(map(flatten_model, m.children()), []) if num_children(m) else [m]
27 | get_layer_groups = lambda m: [nn.Sequential(*flatten_model(m))]
28 |
29 | if optim_cfg.get('FIX_LAYERS', None) and optim_cfg.FIX_LAYERS.ENABLED:
30 | for key, params in model.named_parameters():
31 | module_name = key.split('.')[0]
32 | if module_name in optim_cfg.FIX_LAYERS.NAME:
33 | params.requires_grad = False
34 |
35 | if optim_cfg.get('EXCLUDE_LAYERS', None) and optim_cfg.EXCLUDE_LAYERS.ENABLED:
36 | for key, params in model.named_parameters():
37 | module_name = key.split('.')[0]
38 | if module_name in optim_cfg.EXCLUDE_LAYERS.NAME:
39 | params.exclude = True
40 |
41 | optimizer_func = partial(optim.Adam, betas=(0.9, 0.99))
42 | optimizer = OptimWrapper.create(
43 | optimizer_func, 3e-3, get_layer_groups(model), wd=optim_cfg.WEIGHT_DECAY, true_wd=True, bn_wd=True
44 | )
45 |
46 | # reset exclude to False
47 | if optim_cfg.get('EXCLUDE_LAYERS', None) and optim_cfg.EXCLUDE_LAYERS.ENABLED:
48 | for key, params in model.named_parameters():
49 | module_name = key.split('.')[0]
50 | if module_name in optim_cfg.EXCLUDE_LAYERS.NAME:
51 | params.exclude = False
52 |
53 | else:
54 | raise NotImplementedError
55 |
56 | return optimizer
57 |
58 |
59 | def build_scheduler(optimizer, total_iters_each_epoch, total_epochs, last_epoch, optim_cfg):
60 | decay_steps = [x * total_iters_each_epoch for x in optim_cfg.DECAY_STEP_LIST]
61 | def lr_lbmd(cur_epoch):
62 | cur_decay = 1
63 | for decay_step in decay_steps:
64 | if cur_epoch >= decay_step:
65 | cur_decay = cur_decay * optim_cfg.LR_DECAY
66 | return max(cur_decay, optim_cfg.LR_CLIP / optim_cfg.LR)
67 |
68 | lr_warmup_scheduler = None
69 | total_steps = total_iters_each_epoch * total_epochs
70 | if optim_cfg.OPTIMIZER == 'adam_onecycle':
71 | lr_scheduler = OneCycle(
72 | optimizer, total_steps, optim_cfg.LR, list(optim_cfg.MOMS), optim_cfg.DIV_FACTOR, optim_cfg.PCT_START
73 | )
74 | else:
75 | lr_scheduler = lr_sched.LambdaLR(optimizer, lr_lbmd, last_epoch=last_epoch)
76 |
77 | if optim_cfg.LR_WARMUP:
78 | lr_warmup_scheduler = CosineWarmupLR(
79 | optimizer, T_max=optim_cfg.WARMUP_EPOCH * len(total_iters_each_epoch),
80 | eta_min=optim_cfg.LR / optim_cfg.DIV_FACTOR
81 | )
82 |
83 | return lr_scheduler, lr_warmup_scheduler
84 |
--------------------------------------------------------------------------------
/tools/visual_utils/open3d_vis_utils.py:
--------------------------------------------------------------------------------
1 | """
2 | Open3d visualization tool box
3 | Written by Jihan YANG
4 | All rights preserved from 2021 - present.
5 | """
6 | import open3d
7 | import torch
8 | import matplotlib
9 | import numpy as np
10 |
11 | box_colormap = [
12 | [1, 1, 1],
13 | [0, 1, 0],
14 | [0, 1, 1],
15 | [1, 1, 0],
16 | ]
17 |
18 |
19 | def get_coor_colors(obj_labels):
20 | """
21 | Args:
22 | obj_labels: 1 is ground, labels > 1 indicates different instance cluster
23 |
24 | Returns:
25 | rgb: [N, 3]. color for each point.
26 | """
27 | colors = matplotlib.colors.XKCD_COLORS.values()
28 | max_color_num = obj_labels.max()
29 |
30 | color_list = list(colors)[:max_color_num+1]
31 | colors_rgba = [matplotlib.colors.to_rgba_array(color) for color in color_list]
32 | label_rgba = np.array(colors_rgba)[obj_labels]
33 | label_rgba = label_rgba.squeeze()[:, :3]
34 |
35 | return label_rgba
36 |
37 |
38 | def draw_scenes(points=None, gt_boxes=None, ref_boxes=None, ref_labels=None, ref_scores=None, point_colors=None, draw_origin=True):
39 | if isinstance(points, torch.Tensor):
40 | points = points.cpu().numpy()
41 | if isinstance(gt_boxes, torch.Tensor):
42 | gt_boxes = gt_boxes.cpu().numpy()
43 | if isinstance(ref_boxes, torch.Tensor):
44 | ref_boxes = ref_boxes.cpu().numpy()
45 |
46 | vis = open3d.visualization.Visualizer()
47 | vis.create_window()
48 |
49 | vis.get_render_option().point_size = 1.0
50 | vis.get_render_option().background_color = np.zeros(3)
51 |
52 | # draw origin
53 | if draw_origin:
54 | axis_pcd = open3d.geometry.TriangleMesh.create_coordinate_frame(size=1.0, origin=[0, 0, 0])
55 | vis.add_geometry(axis_pcd)
56 |
57 | if points is not None:
58 | pts = open3d.geometry.PointCloud()
59 | pts.points = open3d.utility.Vector3dVector(points[:, :3])
60 | vis.add_geometry(pts)
61 |
62 | if point_colors is None:
63 | pts.colors = open3d.utility.Vector3dVector(np.ones((points.shape[0], 3)))
64 | else:
65 | pts.colors = open3d.utility.Vector3dVector(point_colors)
66 |
67 | if gt_boxes is not None:
68 | vis = draw_box(vis, gt_boxes, (0, 0, 1))
69 |
70 | if ref_boxes is not None:
71 | vis = draw_box(vis, ref_boxes, (0, 1, 0), ref_labels, ref_scores)
72 |
73 | vis.run()
74 | vis.destroy_window()
75 |
76 |
77 | def translate_boxes_to_open3d_instance(gt_boxes):
78 | """
79 | 4-------- 6
80 | /| /|
81 | 5 -------- 3 .
82 | | | | |
83 | . 7 -------- 1
84 | |/ |/
85 | 2 -------- 0
86 | """
87 | center = gt_boxes[0:3]
88 | lwh = gt_boxes[3:6]
89 | axis_angles = np.array([0, 0, gt_boxes[6] + 1e-10])
90 | rot = open3d.geometry.get_rotation_matrix_from_axis_angle(axis_angles)
91 | box3d = open3d.geometry.OrientedBoundingBox(center, rot, lwh)
92 |
93 | line_set = open3d.geometry.LineSet.create_from_oriented_bounding_box(box3d)
94 |
95 | # import ipdb; ipdb.set_trace(context=20)
96 | lines = np.asarray(line_set.lines)
97 | lines = np.concatenate([lines, np.array([[1, 4], [7, 6]])], axis=0)
98 |
99 | line_set.lines = open3d.utility.Vector2iVector(lines)
100 |
101 | return line_set, box3d
102 |
103 |
104 | def draw_box(vis, gt_boxes, color=(0, 1, 0), ref_labels=None, score=None):
105 | for i in range(gt_boxes.shape[0]):
106 | line_set, box3d = translate_boxes_to_open3d_instance(gt_boxes[i])
107 | if ref_labels is None:
108 | line_set.paint_uniform_color(color)
109 | else:
110 | line_set.paint_uniform_color(box_colormap[ref_labels[i]])
111 |
112 | vis.add_geometry(line_set)
113 |
114 | # if score is not None:
115 | # corners = box3d.get_box_points()
116 | # vis.add_3d_label(corners[5], '%.2f' % score[i])
117 | return vis
118 |
--------------------------------------------------------------------------------