├── .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 | --------------------------------------------------------------------------------