To produce and deal with non-matrix & packed data, and efficiently perform per-pack or pack-wise operations with Pytorch-CUDA exentions.
Pytorch deals with matrix tensors, i.e. batched data with the same data size on each batch.
However, there are situations where we have to deal with non-matrix data with different data size on each batch (or, as named by kaolin "nuggets"). This is especially the case for efficient 3D neural rendering in which each batch or "nuggets" represents buffer data of each ray.
Hence, we develop a toolbox for such needs, trying our best to deal with pack production, pack reduction, per-pack operations and so on.
We acknowledge that the current repository does not cover all required use cases, and our implementations are still far from complete. Therefore, any collaboration or pull-requests are warmly welcomed ! 🤗
| batched data | packed data ("nuggets" in kaolin) |
|---|---|
 |
 |
| batch infos: batch_size=4 batch_data_size=5 |
pack_infos:torch.tensor([[0,4],[4,1],[5,5],[10,3]])which means: first_inds = pack_infos[:,0] = [0,4,5,10]lengths = pack_infos[:,1] = [4,1,5,3] |
interleave_arange_simple
interleave_arange
interleave_linstep
implementations are borrowed from kaolin
packed_sum
packed_mean
implementations are borrowed from kaolin
packed_cumsum
packed_cumprod
packed_diff
forward difference:
packed_diff(feats, appends=...)
forward difference with appends:
packed_backward_diff(feats, prepends=...)
backward difference with prepends:
packed_searchsorted
packed_searchsorted_packed_vals
You achieve multiple packs merging via packed_sort.
However, If you are only merging two packs with one of them already sorted, it's faster to use APIs here since they are based on searchsorted which is faster than general sorting.
merge_two_packs_sorted_aligned
Merge two packs with aligned nuggets
merge_two_packs_sorted_a_includes_b
Merge two packs with not aligned nuggets, but a's nuggets includes b's nuggets.
merge_two_packs_sorted
Merge two packs with intersecting nuggets.
NOTE: The unique nuggets of b must be sorted in advance since their orders will be untouched.
