1. add fp16 and multi-gpu support

2. fix some bugs of nuscenes dataset.

Author: traveller59

README.md

@@ -5,7 +5,7 @@ ONLY support python 3.6+, pytorch 1.0.0+. Tested in Ubuntu 16.04/18.04/Windows 1
 
 ## News
 
-2019-4-1: SECOND V1.6.0alpha released: New Data API, [NuScenes](https://www.nuscenes.org) support, [PointPillars](https://github.com/nutonomy/second.pytorch) support.
+2019-4-1: SECOND V1.6.0alpha released: New Data API, [NuScenes](https://www.nuscenes.org) support, [PointPillars](https://github.com/nutonomy/second.pytorch) support, fp16 and multi-gpu support.
 
 2019-3-21: SECOND V1.5.1 (minor improvement and bug fix) released! 
 
@@ -73,9 +73,11 @@ pip install numba scikit-image scipy pillow
 
 Follow instructions in [spconv](https://github.com/traveller59/spconv) to install spconv. 
 
+If you want to train with fp16 mixed precision (train faster in RTX series, Titan V/RTX and Tesla V100, but I only have 1080Ti), you need to install [apex](https://github.com/NVIDIA/apex).
+
 If you want to use NuScenes dataset, you need to install [nuscenes-devkit](https://github.com/nutonomy/nuscenes-devkit), I recommend to copy nuscenes in python-sdk to second/.. folder (equalivent to add it to PYTHONPATH) and manually install its dependencies, use pip to install devkit will install many fixed-version library.
 
-### 3. Setup cuda for numba
+### 3. Setup cuda for numba (will be removed in 1.6.0 release)
 
 you need to add following environment variable for numba.cuda, you can add them to ~/.bashrc:
 
@@ -128,12 +130,11 @@ Download NuScenes dataset:
        ├── maps          <-- unused
        └── v1.0-test     <-- metadata
 ```
-Since the dataset is really large, you can download parts of the dataset.
 
 Then run
 ```bash
-python create_data.py nuscenes_data_prep --data_path=NUSCENES_TRAINVAL_DATASET_ROOT --version="v1.0-trainval" --max_sweeps=9
-python create_data.py nuscenes_data_prep --data_path=NUSCENES_TEST_DATASET_ROOT --version="v1.0-test" --max_sweeps=9
+python create_data.py nuscenes_data_prep --data_path=NUSCENES_TRAINVAL_DATASET_ROOT --version="v1.0-trainval" --max_sweeps=10
+python create_data.py nuscenes_data_prep --data_path=NUSCENES_TEST_DATASET_ROOT --version="v1.0-test" --max_sweeps=10
 ```
 
 * Modify config file
@@ -168,10 +169,24 @@ eval_input_reader: {
 
 I recommend to use script.py to train and eval. see script.py for more details.
 
+#### train with single GPU
+
 ```bash
 python ./pytorch/train.py train --config_path=./configs/car.fhd.config --model_dir=/path/to/model_dir
 ```
 
+#### train with multiple GPU (need test, I only have one GPU)
+
+Assume you have 4 GPUs and want to train with 3 GPUs:
+
+```bash
+CUDA_VISIBLE_DEVICES=0,1,3 python ./pytorch/train.py train --config_path=./configs/car.fhd.config --model_dir=/path/to/model_dir --multi_gpu=True
+```
+
+#### train with fp16 (mixed precision)
+
+Modify config file, set enable_mixed_precision to true.
+
 * Make sure "/path/to/model_dir" doesn't exist if you want to train new model. A new directory will be created if the model_dir doesn't exist, otherwise will read checkpoints in it.
 
 * training process use batchsize=6 as default for 1080Ti, you need to reduce batchsize if your GPU has less memory.

RELEASE.md

@@ -12,6 +12,8 @@
 
 4. Full Tensorboard support.
 
+5. FP16 and multi-gpu (need test, I only have one gpu) support.
+
 ## Minor Improvements and Bug fixes
 
 1. Move all data-specific functions to their corresponding dataset file.

second/builder/dataset_builder.py

@@ -35,7 +35,8 @@ def build(input_reader_config,
           model_config,
           training,
           voxel_generator,
-          target_assigner):
+          target_assigner,
+          multi_gpu=False):
     """Builds a tensor dictionary based on the InputReader config.
 
     Args:
@@ -69,6 +70,7 @@ def build(input_reader_config,
     assert all([n != '' for n in target_assigner.classes]), "you must specify class_name in anchor_generators."
     dataset_cls = get_dataset_class(dataset_cfg.dataset_class_name)
     assert dataset_cls.NumPointFeatures >= 3, "you must set this to correct value"
+    assert dataset_cls.NumPointFeatures == num_point_features, "currently you need keep them same"
     prep_func = partial(
         prep_pointcloud,
         root_path=dataset_cfg.kitti_root_path,
@@ -95,7 +97,8 @@ def build(input_reader_config,
         remove_points_after_sample=prep_cfg.remove_points_after_sample,
         remove_environment=prep_cfg.remove_environment,
         use_group_id=prep_cfg.use_group_id,
-        out_size_factor=out_size_factor)
+        out_size_factor=out_size_factor,
+        multi_gpu=multi_gpu)
 
     ret = target_assigner.generate_anchors(feature_map_size)
     class_names = target_assigner.classes

second/configs/all.fhd.config

@@ -7,7 +7,7 @@ model: {
       max_number_of_points_per_voxel : 5
     }
     voxel_feature_extractor: {
-      module_class_name: "SimpleVoxel"
+      module_class_name: "SimpleVoxelRadius"
       num_filters: [16]
       with_distance: false
       num_input_features: 4
@@ -233,8 +233,8 @@ train_config: {
   steps_per_eval: 6190 # 1238 * 5
   save_checkpoints_secs : 1800 # half hour
   save_summary_steps : 10
-  enable_mixed_precision: false # for fp16 training, but sparseconvnet don't support fp16
-  loss_scale_factor : 512.0
+  enable_mixed_precision: false 
+  loss_scale_factor: 8.0
   clear_metrics_every_epoch: true
 }
 

second/configs/car.fhd.config

@@ -8,7 +8,7 @@ model: {
     }
 
     voxel_feature_extractor: {
-      module_class_name: "VoxelFeatureExtractorV3"
+      module_class_name: "SimpleVoxel"
       num_filters: [16]
       with_distance: false
       num_input_features: 4
@@ -186,8 +186,8 @@ train_config: {
 
   save_checkpoints_secs : 1800 # half hour
   save_summary_steps : 10
-  enable_mixed_precision: false # for fp16 training, but sparseconvnet don't support fp16
-  loss_scale_factor : 512.0
+  enable_mixed_precision: false
+  loss_scale_factor : 8.0
   clear_metrics_every_epoch: true
 }
 

second/configs/car.fhd.onestage.config

@@ -8,7 +8,7 @@ model: {
     }
 
     voxel_feature_extractor: {
-      module_class_name: "VoxelFeatureExtractorV3"
+      module_class_name: "SimpleVoxel"
       num_filters: [16]
       with_distance: false
       num_input_features: 4
@@ -170,8 +170,8 @@ train_config: {
   steps_per_eval: 3095 # 619 * 5
   save_checkpoints_secs : 1800 # half hour
   save_summary_steps : 10
-  enable_mixed_precision: false # for fp16 training, don't use this.
-  loss_scale_factor : 512.0
+  enable_mixed_precision: false 
+  loss_scale_factor : 8.0
   clear_metrics_every_epoch: true
 }
 

second/configs/car.lite.config

@@ -9,7 +9,7 @@ model: {
     }
 
     voxel_feature_extractor: {
-      module_class_name: "SimpleVoxel"
+      module_class_name: "SimpleVoxelRadius"
       num_filters: [16]
       with_distance: false
       num_input_features: 4
@@ -176,8 +176,8 @@ train_config: {
   steps_per_eval: 1550 # 310 * 5
   save_checkpoints_secs : 1800 # half hour
   save_summary_steps : 10
-  enable_mixed_precision: false # for fp16 training, but sparseconvnet don't support fp16
-  loss_scale_factor : 512.0
+  enable_mixed_precision: false 
+  loss_scale_factor: 8.0
   clear_metrics_every_epoch: true
 }
 

second/configs/car.lite.nu.config

@@ -2,30 +2,30 @@ model: {
   second: {
     voxel_generator {
       point_cloud_range : [-50, -50.0, -4, 50, 50, 2]
-      voxel_size : [0.1, 0.1, 0.3]
+      voxel_size : [0.05, 0.05, 0.15]
       max_number_of_points_per_voxel : 1
     }
 
     voxel_feature_extractor: {
-      module_class_name: "SimpleVoxel"
+      module_class_name: "SimpleVoxelRadius"
       num_filters: [16]
       with_distance: false
       num_input_features: 4
     }
     middle_feature_extractor: {
-      module_class_name: "SpMiddleHDLite"
+      module_class_name: "SpMiddleFHDLite"
       # num_filters_down1: [] # protobuf don't support empty list.
       # num_filters_down2: []
-      downsample_factor: 4
-      num_input_features: 3 # SimpleVoxel output 3 - 1 features
+      downsample_factor: 8
+      num_input_features: 3
     }
     rpn: {
       module_class_name: "RPNV2"
       layer_nums: [5]
       layer_strides: [1]
       num_filters: [128]
-      # upsample_strides: [1]
-      # num_upsample_filters: [128]
+      upsample_strides: [1]
+      num_upsample_filters: [128]
       use_groupnorm: false
       num_groups: 32
       num_input_features: 128
@@ -67,8 +67,8 @@ model: {
     use_multi_class_nms: false
     nms_pre_max_size: 1000
     nms_post_max_size: 100
-    nms_score_threshold: 0.3 # 0.4 in submit, but 0.3 can get better hard performance
-    nms_iou_threshold: 0.01
+    nms_score_threshold: 0.05 # 0.4 in submit, but 0.3 can get better hard performance
+    nms_iou_threshold: 0.5
 
     box_coder: {
       ground_box3d_coder: {
@@ -79,12 +79,12 @@ model: {
     target_assigner: {
       anchor_generators: {
         anchor_generator_range: {
-          sizes: [1.968534, 4.6260232, 1.7432361] # wlh
-          anchor_ranges: [-50, -50.0, -0.916, 50, 50, -0.916]
+          sizes: [1.95017717, 4.60718145, 1.72270761] # wlh
+          anchor_ranges: [-50, -50.0, -0.93897414, 50, 50, -0.93897414]
           rotations: [0, 1.57] # DON'T modify this unless you are very familiar with my code.
           matched_threshold : 0.6
           unmatched_threshold : 0.45
-          class_name: "vehicle.car"
+          class_name: "car"
         }
       }
       sample_positive_fraction : -1
@@ -102,17 +102,23 @@ train_input_reader: {
     dataset_class_name: "NuScenesDataset"
     kitti_info_path: "/media/yy/960evo/datasets/nuscene/v1.0-trainval/infos_train.pkl"
     kitti_root_path: "/media/yy/960evo/datasets/nuscene/v1.0-trainval"
+    # kitti_info_path: "/media/yy/960evo/datasets/nuscene/v1.0-mini/infos_train.pkl"
+    # kitti_root_path: "/media/yy/960evo/datasets/nuscene/v1.0-mini"
+
   }
 
-  batch_size: 8
+  batch_size: 6
   preprocess: {
-    max_number_of_voxels: 20000
-    shuffle_points: true
+    max_number_of_voxels: 63000
+    shuffle_points: false
     num_workers: 3
-    groundtruth_localization_noise_std: [1.0, 1.0, 0.5]
+    groundtruth_localization_noise_std: [0, 0, 0]
+    groundtruth_rotation_uniform_noise: [0, 0]
+
+    # groundtruth_localization_noise_std: [0.25, 0.25, 0.25]
     # groundtruth_rotation_uniform_noise: [-0.3141592654, 0.3141592654]
-    groundtruth_rotation_uniform_noise: [-0.78539816, 0.78539816]
-    global_rotation_uniform_noise: [-0.78539816, 0.78539816]
+    # groundtruth_rotation_uniform_noise: [-0.78539816, 0.78539816]
+    global_rotation_uniform_noise: [-1.57, 1.57]
     global_scaling_uniform_noise: [0.95, 1.05]
     global_random_rotation_range_per_object: [0, 0] # pi/4 ~ 3pi/4
     global_translate_noise_std: [0.2, 0.2, 0.2]
@@ -123,23 +129,9 @@ train_input_reader: {
     remove_unknown_examples: false
     remove_environment: false
     database_sampler {
-      database_info_path: "/media/yy/960evo/datasets/nuscene/v1.0-trainval/kitti_dbinfos_train.pkl"
-      sample_groups {
-        name_to_max_num {
-          key: "vehicle.car"
-          value: 5
-        }
-      }
-      database_prep_steps {
-        filter_by_min_num_points {
-          min_num_point_pairs {
-            key: "vehicle.car"
-            value: 5
-          }
-        }
-      }
-      global_random_rotation_range_per_object: [0, 0]
-      rate: 1.0
+      # leave this empty to disable database_sampler, nuscenes don't need sample
+      # because 1. the number of ground-truth is enough. 2. sweeps don't support 
+      # sample.
     }
   }
 }
@@ -160,24 +152,24 @@ train_config: {
     fixed_weight_decay: true
     use_moving_average: false
   }
-  steps: 70240 # 3512 * 20 (28096 / 8)
-  steps_per_eval: 7024 # 310 * 5
+  steps: 234450 # 4689 * 50 (28130 // 6 + 1)
+  steps_per_eval: 9378 # 4689 * 2
   save_checkpoints_secs : 1800 # half hour
   save_summary_steps : 10
-  enable_mixed_precision: false # for fp16 training, but sparseconvnet don't support fp16
-  loss_scale_factor : 512.0
+  enable_mixed_precision: false 
+  loss_scale_factor: 8.0
   clear_metrics_every_epoch: true
 }
 
 eval_input_reader: {
-  batch_size: 8
+  batch_size: 6
   dataset: {
     dataset_class_name: "NuScenesDataset"
     kitti_info_path: "/media/yy/960evo/datasets/nuscene/v1.0-trainval/infos_val.pkl"
     kitti_root_path: "/media/yy/960evo/datasets/nuscene/v1.0-trainval"
   }
   preprocess: {
-    max_number_of_voxels: 50000
+    max_number_of_voxels: 80000
     shuffle_points: false
     num_workers: 3
     anchor_area_threshold: -1

second/configs/people.fhd.config

@@ -8,7 +8,7 @@ model: {
     }
 
     voxel_feature_extractor: {
-      module_class_name: "VoxelFeatureExtractorV3"
+      module_class_name: "SimpleVoxel"
       num_filters: [16]
       with_distance: false
       num_input_features: 4
@@ -197,8 +197,8 @@ train_config: {
   steps_per_eval: 3095 # 619 * 5
   save_checkpoints_secs : 1800 # half hour
   save_summary_steps : 10
-  enable_mixed_precision: false # for fp16 training, but sparseconvnet don't support fp16
-  loss_scale_factor : 512.0
+  enable_mixed_precision: false 
+  loss_scale_factor: 8.0
   clear_metrics_every_epoch: true
 }
 

second/configs/pointpillars/car/xyres_16.config

@@ -99,8 +99,8 @@ model: {
 train_input_reader: {
   dataset: {
     dataset_class_name: "KittiDataset"
-    kitti_info_path: "/media/yy/960evo/datasets/kitti/kitti_infos_train.pkl"
-    kitti_root_path: "/media/yy/960evo/datasets/kitti"
+    kitti_info_path: "/media/yy/My Passport/datasets/kitti/kitti_infos_train.pkl"
+    kitti_root_path: "/media/yy/My Passport/datasets/kitti"
   }
 
   batch_size: 2
@@ -121,7 +121,7 @@ train_input_reader: {
     remove_unknown_examples: false
     remove_environment: false
     database_sampler {
-      database_info_path: "/media/yy/960evo/datasets/kitti/kitti_dbinfos_train.pkl"
+      database_info_path: "/media/yy/My Passport/datasets/kitti/kitti_dbinfos_train.pkl"
       sample_groups {
         name_to_max_num {
           key: "Car"
@@ -164,21 +164,21 @@ train_config: {
     use_moving_average: false
 
   }
-  steps: 92800 # 1856 steps per epoch * 50 epochs
-  steps_per_eval: 9280 # 1856 steps per epoch * 5 epochs
+  steps: 37120 # 92800 # 1856 steps per epoch * 50 epochs
+  steps_per_eval: 3712 # 9280 # 1856 steps per epoch * 5 epochs
   save_checkpoints_secs : 1800 # half hour
   save_summary_steps : 10
   enable_mixed_precision: false
-  loss_scale_factor : 512.0
+  loss_scale_factor: 8.0
   clear_metrics_every_epoch: true
 }
 
 eval_input_reader: {
   dataset: {
     dataset_class_name: "KittiDataset"
-    kitti_info_path: "/media/yy/960evo/datasets/kitti/kitti_infos_val.pkl"
-    # kitti_info_path: "/media/yy/960evo/datasets/kitti/kitti_infos_test.pkl"
-    kitti_root_path: "/media/yy/960evo/datasets/kitti"
+    kitti_info_path: "/media/yy/My Passport/datasets/kitti/kitti_infos_val.pkl"
+    # kitti_info_path: "/media/yy/My Passport/datasets/kitti/kitti_infos_test.pkl"
+    kitti_root_path: "/media/yy/My Passport/datasets/kitti"
   }
   batch_size: 2