Customized op for WithinBox3d. #1020
Conversation
keras_cv/ops/point_cloud.py
Outdated
| from tensorflow.python.framework import load_library | ||
| from tensorflow.python.platform import resource_loader | ||
|
|
||
| custom_ops = load_library.load_op_library( |
There was a problem hiding this comment.
In order to support wheels without custom binaries, I think we'll need to try-catch this so that we don't break at import time if the .so files are missing.
For IoU3D I did this by making IoU3D an object, but IIRC we decided that's probably not the right solution, so maybe we just want to try/except this and raise an Exception when the op is called if we weren't able to load the custom ops at load time.
wdyt?
Still much slower than v1.
|
Can we add a little bit of c++ docstrings? |
| box_dimension = 20.0 | ||
|
|
||
|
|
||
| def get_points_boxes(): |
There was a problem hiding this comment.
Do you ensure no overlap between two bounding boxes? The v2 op assumes no overlaps, right?
There was a problem hiding this comment.
Good point. Added docstring
| for _ in range(5): | ||
| print(datetime.now()) | ||
| res = keras_cv.ops.within_box3d_index(points, boxes) | ||
| self.assertAllClose(res.shape, points.shape[:1]) |
There was a problem hiding this comment.
Is it possible to test assertAllClose(v1_points, v2_points), where v1_points and v2_points are generated by v1 and v2 OPs.
You can sort V1 and V2 points based on XYZ so that you can compare them
There was a problem hiding this comment.
will compare this once we have "is_within_any_box3d"
|
|
||
|
|
||
| # TODO(lengzhaoqi/tanzhenyu): compare the performance with v1 | ||
| def is_within_box3d_v2(points, boxes): |
There was a problem hiding this comment.
As far as I can tell, there's not a great way to efficiently verify that boxes don't overlap, so we should probably include info in the docstring warning users that we have a non-overlap precondition for boxes.
| // those points are within the box | ||
| auto within_fn = [&points_x_indices, &points_y_indices, &boxes_vec, &points_vec, &boxes_indices_t](int64_t begin, int64_t end) { | ||
| for (int64_t idx = begin; idx < end; ++idx) { | ||
| std::unordered_set<int>& set_a = points_x_indices[idx]; |
There was a problem hiding this comment.
A side note is https://github.com/abseil/abseil-cpp/blob/master/absl/container/flat_hash_map.h available here?
There was a problem hiding this comment.
can you elaborate? (I'm not sure if I understand the comment)
There was a problem hiding this comment.
flat_hash_set is more efficient than std:: unordered_set, but we can check the profiling first to see whether this is the bottleneck.
| std::unordered_set<int>& set_a = points_x_indices[idx]; | ||
| std::unordered_set<int>& set_b = points_y_indices[idx]; | ||
| std::unordered_set<int> p_set; | ||
| for (auto val : set_a) { |
There was a problem hiding this comment.
Nit: std::set_intersection?
There was a problem hiding this comment.
that only works for sorted range of set_a and set_b
|
/gcbrun |
|
the failing test is due to previous random_choice test, not related to this PR |
* Customized op for WithinBox3d. * add some parallelization. Still much slower than v1. * add more parallelization. * Replace with std binary search, add unit test. * add lazy loading. * fix format. * fix format. * address comment.
pip uninstall keras_cv
bazel build keras_cv/custom_ops:all
mv bazel-bin/keras_cv/custom_ops/*.so keras_cv/custom_ops/
bazel build build_pip_pkg
./bazel-bin/build_pip_pkg wheels
pip install wheels/keras_cv-0.0.0-cp39-cp39-linux_x86_64.whl