Mixed Precision Support

[clee-ai]

I implemented basic half precision support compatible with torch.cuda.amp.autocast. I also annotated the c++ convolution code a bit.

I experimented a lot with resizing tensors to have dimensions of multiples of 8, but it seems like it won't change execution time significantly, so I left that out. With size=400000, batch_size=4, on my 2080ti, I get the following results (I attached the nvprof benchmarks as well):

Default precision (nvprof):

Done in: 15.723482s
max_memory_allocated 3492.76513671875 MB
max_memory_reserved 5024.0 MB

Mixed precision, no optimization (nvprof):

Done in: 14.207781s
max_memory_allocated 1876.51904296875 MB
max_memory_reserved 3284.0 MB

Mixed precision, all mm ops in multiples of 8 (nvprof):

Done in: 14.748600s
max_memory_allocated 1924.66552734375 MB
max_memory_reserved 3274.0 MB

Looking at the nvprof results, it looks like barely any computation time is spent on mm ops anyways:

            Type  Time(%)      Time     Calls       Avg       Min       Max  Name
 GPU activities:   43.98%  3.13228s        40  78.307ms  64.375ms  110.07ms  void at::native::batch_norm_backward_kernel
                   20.28%  1.44425s        40  36.106ms  31.259ms  45.500ms  void at::native::batch_norm_collect_statistics_kernel
                    6.36%  453.17ms        40  11.329ms  7.7208ms  17.821ms  void at::native::batch_norm_transform_input_kernel
                    6.30%  449.02ms        20  22.451ms  7.9090ms  39.716ms  cuckooLookupKernel_Multi
                    3.83%  272.89ms        10  27.289ms  26.730ms  29.211ms  void cunn_ClassNLLCriterion_updateOutput_kernel
                    3.22%  229.55ms      2080  110.36us  18.337us  1.8771ms  void gather_kernel

But mixed precision will slightly speed up models and reduce their memory footprint to some degree. Let me know if you have any questions/suggestions. (should address #17 )

[zhijian-liu]

Thank you so much for your great efforts! We were currently running experiments using this new version to verify whether it has any negative influence on the performance. Btw, I'm wondering if it is possible to also support mixed precision for spvoxelize and spdevoxelize. Thanks!

[clee-ai]

Great! Please let me know your results for performance.

It should be no problem adding those functions. If possible, it would be great if you could provide me with a minimum test script that uses these functions, similar to examples/test.py, so that I can make sure it works correctly, if you have one. If not I will try to make my own.

[zhijian-liu]

The inference of SPVNAS should be a pretty good example to test these functions: https://github.com/mit-han-lab/spvnas. Thanks!

[zhijian-liu]

The large-scale experiments of MinkowskiNet on NuScenes have just finished:

• The performance of the mixed-precision training is almost the same as that of the full-precision training: 76.43 v.s. 76.78.
• The speedup is fairly limited: the total training time is reduced from 7.3 hours to 6.4 hours (around 10% reduction).
• The memory reduction is very significant: the memory usage is reduced from 48.8G to 28.8G (40% reduction).

[clee-ai]

Great, I'm glad it works! 10% and 40% was about what I saw as well in my tests.

I added support for insertion and devoxelization in half/double precision with my latest commits. It worked well on my spvnas inference test but I didn't test the backwards functions. Please let me know how it looks in your tests!

[zhijian-liu]

Thanks for the efforts! I will launch some large-scale experiments to test these functions as well.

UPDATE: The results of SPVNAS are similar to those of MinkowskiNet.

[zhijian-liu]

The implementation looks great! I think it's ready to be merged.

[kentang-mit]

Thanks @CCInc for the great effort on mix precision support! I've also been through the changes and believe that this pull request is ready for merging.

[clee-ai]

Great, glad to hear it! I will also be happy to help implement SPVNAS architecture search or any other tasks you need, feel free to let me know in an email.