diff --git a/.gitignore b/.gitignore index 7eed5e2..744215b 100644 --- a/.gitignore +++ b/.gitignore @@ -25,4 +25,7 @@ cancel.sh # cuda build files *.egg-info build* -dist* \ No newline at end of file +dist* + +data +*__pycache__ \ No newline at end of file diff --git a/README.md b/README.md index 41f9f6e..a1aa2ae 100644 --- a/README.md +++ b/README.md @@ -16,6 +16,11 @@ It is also an official implementation of the following papers (sorted by the tim Preprint; Under review; 2025 [ Strategy ] [ Self-Supervised ] - [ [arXiv](https://arxiv.org/abs/2503.00803) ] [ [Project](https://kin-zhang.github.io/HiMo/) ] +- **VoteFlow: Enforcing Local Rigidity in Self-Supervised Scene Flow** +*Yancong Lin\*, Shiming Wang\*, Liangliang Nan, Julian Kooij, Holger Caesar* +Conference on Computer Vision and Pattern Recognition (**CVPR**) 2025 +[ Backbone ] [ Self-Supervised ] - [ [arXiv](https://arxiv.org/abs/2503.22328) ] [ [Project](https://github.com/tudelft-iv/VoteFlow/)] → [here](#VoteFLow) + - **Flow4D: Leveraging 4D Voxel Network for LiDAR Scene Flow Estimation** *Jaeyeul Kim, Jungwan Woo, Ukcheol Shin, Jean Oh, Sunghoon Im* IEEE Robotics and Automation Letters (**RA-L**) 2025 @@ -46,6 +51,7 @@ Additionally, *OpenSceneFlow* integrates following excellent works: [ICLR'24 Zer - [x] [NSFP](https://arxiv.org/abs/2111.01253): NeurIPS 2021, faster 3x than original version because of [our CUDA speed up](assets/cuda/README.md), same (slightly better) performance. - [x] [FastNSF](https://arxiv.org/abs/2304.09121): ICCV 2023. SSL optimization-based. - [ ] [ICP-Flow](https://arxiv.org/abs/2402.17351): CVPR 2024. SSL optimization-based. Done coding, public after review. +- [ ] [EulerFlow](https://arxiv.org/abs/2410.02031): ICLR 2025. SSL optimization-based. In my plan, haven't coding yet. @@ -128,6 +134,27 @@ And free yourself from trainning, you can download the pretrained weight from [H mamba activate opensf ``` +### VoteFLow +Extra pakcges needed for VoteFlow, [pytorch3d](https://pytorch3d.org/) (prefer 0.7.7) and [torch-scatter](https://github.com/rusty1s/pytorch_scatter?tab=readme-ov-file) (prefer 2.1.2): + +```bash +# Install Pytorch3d +conda install pytorch3d -c pytorch3d + +# Install torch-scatter +pip install torch-scatter -f https://data.pyg.org/whl/torch-2.0.0+cu117.html +``` + +Train VoteFlow with the leaderboard submit config. [Runtime: Around 32 hours in 4 x V100 GPUs.] +```bash +python train.py model=voteflow lr=2e-4 lr_scheduler=step epochs=12 batch_size=4 model.target.m=8 model.target.n=128 loss_fn=seflowLoss "add_seloss={chamfer_dis: 1.0, static_flow_loss: 1.0, dynamic_chamfer_dis: 1.0, cluster_based_pc0pc1: 1.0}" +``` + +Pretrained weight can be downloaded through: +```bash +wget https://huggingface.co/kin-zhang/OpenSceneFlow/resolve/main/voteflow_best.ckpt +``` + ### Flow4D Train Flow4D with the leaderboard submit config. [Runtime: Around 18 hours in 4x RTX 3090 GPUs.] @@ -327,6 +354,13 @@ And our excellent collaborators works contributed to this codebase also: journal={arXiv preprint arXiv:2501.17821}, year={2025} } + +@inproceedings{lin2025voteflow, + title={VoteFlow: Enforcing Local Rigidity in Self-Supervised Scene Flow}, + author={Lin, Yancong and Wang, Shiming and Nan, Liangliang and Kooij, Julian and Caesar, Holger}, + booktitle={CVPR}, + year={2025}, +} ``` Thank you for your support! ❤️ diff --git a/conf/config.yaml b/conf/config.yaml index 1e34f76..4a80895 100644 --- a/conf/config.yaml +++ b/conf/config.yaml @@ -27,6 +27,7 @@ gradient_clip_val: 5.0 # optimizer ==> Adam lr: 2e-6 +lr_scheduler: None # choices: [cosine, step, linear] loss_fn: deflowLoss # choices: [ff3dLoss, zeroflowLoss, deflowLoss, seflowLoss] add_seloss: # {chamfer_dis: 1.0, static_flow_loss: 1.0, dynamic_chamfer_dis: 1.0, cluster_based_pc0pc1: 1.0} diff --git a/conf/model/voteflow.yaml b/conf/model/voteflow.yaml new file mode 100644 index 0000000..3362d8e --- /dev/null +++ b/conf/model/voteflow.yaml @@ -0,0 +1,18 @@ +name: voteflow + +target: + _target_: src.models.VoteFlow + using_voting: True + nframes: 1 + m: 8 + n: 128 + input_channels: 32 + output_channels: 64 + point_cloud_range: ${point_cloud_range} + voxel_size: ${voxel_size} + grid_feature_size: [512, 512] + decoder_layers: 4 + use_ball_query: False + vol_conv_hidden_dim: 16 + +val_monitor: val/Dynamic/Mean \ No newline at end of file diff --git a/src/models/__init__.py b/src/models/__init__.py index 4735f1c..410b773 100644 --- a/src/models/__init__.py +++ b/src/models/__init__.py @@ -14,6 +14,13 @@ from .fastflow3d import FastFlow3D from .nsfp import NSFP +# check README for package: +try: + from .voteflow import VoteFlow +except ImportError as e: + print("\033[93m--- WARNING [model]: VoteFlow is not imported, as it requires some lib which is not installed.") + print(f"Detail error message\033[0m: {e}. Just ignore this warning if code runs without these models.") + # following need install extra package: # * pip install spconv-cu117 try: diff --git a/src/models/basic/voteflow_plugin/__init__.py b/src/models/basic/voteflow_plugin/__init__.py new file mode 100644 index 0000000..55a35ae --- /dev/null +++ b/src/models/basic/voteflow_plugin/__init__.py @@ -0,0 +1,3 @@ +from .hough_transformation import HT_CUDA +from .utils import calculate_unq_voxels, batched_masked_gather, pad_to_batch +from .voteflow_module import VolConvBN, VoteFlowLinearDecoder \ No newline at end of file diff --git a/src/models/basic/voteflow_plugin/hough_transformation/__init__.py b/src/models/basic/voteflow_plugin/hough_transformation/__init__.py new file mode 100755 index 0000000..be0a4d0 --- /dev/null +++ b/src/models/basic/voteflow_plugin/hough_transformation/__init__.py @@ -0,0 +1,4 @@ +from .ht_cuda import HT_CUDA +__all__ = [ + 'HT_CUDA', +] diff --git a/src/models/basic/voteflow_plugin/hough_transformation/cpp_im2ht/cuda_utils.h b/src/models/basic/voteflow_plugin/hough_transformation/cpp_im2ht/cuda_utils.h new file mode 100755 index 0000000..2a231e1 --- /dev/null +++ b/src/models/basic/voteflow_plugin/hough_transformation/cpp_im2ht/cuda_utils.h @@ -0,0 +1,77 @@ +// Copyright (c) Facebook, Inc. and its affiliates. +// +// This source code is licensed under the MIT license found in the +// LICENSE file in the root directory of this source tree. + +// #include +// #include +// #include +#include +#include + + +#define CUDA_KERNEL_LOOP(i, n) \ + for (int i = blockIdx.x * blockDim.x + threadIdx.x; \ + i < (n); \ + i += blockDim.x * gridDim.x) + +// CUDA: thread number configuration. +// Use 1024 threads per block, which requires cuda sm_2x or above, +// or fall back to attempt compatibility (best of luck to you). +//#if __CUDA_ARCH__ >= 200 +// const int CUDA_NUM_THREADS = 1024; +//#else +// const int CUDA_NUM_THREADS = 512; +//#endif +const int CUDA_NUM_THREADS = 1024; + +inline int GET_BLOCKS(const int N) +{ + return (N + CUDA_NUM_THREADS - 1) / CUDA_NUM_THREADS; +} + + +// #ifndef CUDA_NUM_THREADS +// #define CUDA_NUM_THREADS 1024 +// #endif + +// #define CUDA_KERNEL_LOOP(i, n) \ +// for (int i = blockIdx.x * blockDim.x + threadIdx.x; \ +// i < (n); \ +// i += blockDim.x * gridDim.x) + +// const int CUDA_NUM_THREADS = 1024; + +// inline int GET_BLOCKS(const int N) +// { +// return (N + CUDA_NUM_THREADS - 1) / CUDA_NUM_THREADS; +// } + + +// // C++ interface +// // NOTE: AT_ASSERT has become AT_CHECK on master after 0.4. +// // #define CHECK_CUDA(x) AT_CHECK(x.type().is_cuda(), #x " must be a CUDA tensor") +// #define CHECK_CONTIGUOUS(x) AT_CHECK(x.is_contiguous(), #x " must be contiguous") +// #define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x) + +// #define CHECK_CUDA(x) \ +// do { \ +// AT_CHECK(x.type().is_cuda(), #x " must be a CUDA tensor"); \ +// } while (0) + +// #define CHECK_CONTIGUOUS(x) \ +// do { \ +// AT_CHECK(x.is_contiguous(), #x " must be a contiguous tensor"); \ +// } while (0) + +// #define CHECK_IS_INT(x) \ +// do { \ +// AT_CHECK(x.scalar_type() == at::ScalarType::Int, \ +// #x " must be an int tensor"); \ +// } while (0) + +// #define CHECK_IS_FLOAT(x) \ +// do { \ +// AT_CHECK(x.scalar_type() == at::ScalarType::Float, \ +// #x " must be a float tensor"); \ +// } while (0) diff --git a/src/models/basic/voteflow_plugin/hough_transformation/cpp_im2ht/ht_cuda.cu b/src/models/basic/voteflow_plugin/hough_transformation/cpp_im2ht/ht_cuda.cu new file mode 100755 index 0000000..095be69 --- /dev/null +++ b/src/models/basic/voteflow_plugin/hough_transformation/cpp_im2ht/ht_cuda.cu @@ -0,0 +1,116 @@ +#include +#include +#include +#include +#include +#include +#include "im2ht_cuda.cuh" + + +at::Tensor +ht_cuda_forward( + const at::Tensor &feats_src_dst, + const at::Tensor &voxels_src, + const at::Tensor &voxels_dst, + const at::Tensor &idxs_src, + const at::Tensor &idxs_dst, + const int h, + const int w, + const int d + ) +{ + AT_ASSERTM(feats_src_dst.is_contiguous(), "feats tensor has to be contiguous"); + AT_ASSERTM(voxels_src.is_contiguous(), "voxels tensor has to be contiguous"); + AT_ASSERTM(voxels_dst.is_contiguous(), "voxels tensor has to be contiguous"); + AT_ASSERTM(idxs_src.is_contiguous(), "idxs tensor has to be contiguous"); + AT_ASSERTM(idxs_dst.is_contiguous(), "idxs tensor has to be contiguous"); + + // AT_ASSERTM(input.type().is_cuda(), "input must be a CUDA tensor"); + + const int b = feats_src_dst.size(0); + const int l = feats_src_dst.size(1); // num of points + const int c = feats_src_dst.size(-1); // num of channels + const int m = idxs_src.size(2); // m + const int n = idxs_dst.size(2); // n + + // AT_ASSERTM(m_ == m && n_ <= n && k_ == k, + // "input shape and predefined shape do not match: (%04d x %04d x %04d vs %04d x %04d x %04d).", k_, m_, n_, k, m, n); + + // printf("dimensions b=%06d, l=%06d, m=%06d, n=%06d, h=%06d, w= %06d, d= %06d \n", + // b, l, m, n, h, w, d); + auto vol_ht = at::zeros({b, l, c, h, w}, feats_src_dst.options()); + + AT_DISPATCH_FLOATING_TYPES(vol_ht.type(), "im2ht_cuda_forward", ([&] { + im2ht_cuda_forward(at::cuda::getCurrentCUDAStream(), + feats_src_dst.data() , + vol_ht.data(), + voxels_src.data() , + voxels_dst.data(), + idxs_src.data(), + idxs_dst.data(), + b, l, c, + m, n, + h, w, d + ); + + })); + // std::cout <<"output" < +at::Tensor +ht_cuda_backward( + const at::Tensor &grad_vol, + const at::Tensor &voxels_src, + const at::Tensor &voxels_dst, + const at::Tensor &idxs_src, + const at::Tensor &idxs_dst, + const int h, + const int w, + const int d + ) +{ + + AT_ASSERTM(grad_vol.is_contiguous(), "grad_output tensor has to be contiguous"); + AT_ASSERTM(voxels_src.is_contiguous(), "voxels tensor has to be contiguous"); + AT_ASSERTM(voxels_dst.is_contiguous(), "voxels tensor has to be contiguous"); + AT_ASSERTM(idxs_src.is_contiguous(), "idxs tensor has to be contiguous"); + AT_ASSERTM(idxs_dst.is_contiguous(), "idxs tensor has to be contiguous"); + + // AT_ASSERTM(grad_output.type().is_cuda(), "grad_output must be a CUDA tensor"); + + // grad_output: [b, l, c, h, w] + const int b = grad_vol.size(0); + const int l = grad_vol.size(1); + const int c = grad_vol.size(2); + + const int m = idxs_src.size(2); // m + const int n = idxs_dst.size(2); // n + + // printf("h_ = %04d, w_ = %04d, d_ = %04d vs h = %04d, w = %04d, d = %04d\n", h_, w_, d_, h, w, d); + // AT_ASSERTM(h_ == h && w_ == w && d_== d, + // "grad_out shape and predefined shape do not match: ", h_, ' ', w_, ' ', d_, 'vs', h, ' ', w, ' ', d); + + auto grad_feats_src_dst = at::zeros({b, l, n, c}, grad_vol.options()); + + AT_DISPATCH_FLOATING_TYPES(grad_vol.type(), "im2ht_cuda_backward", ([&] { + im2ht_cuda_backward(at::cuda::getCurrentCUDAStream(), + grad_feats_src_dst.data(), + grad_vol.data(), + voxels_src.data(), + voxels_dst.data(), + idxs_src.data(), + idxs_dst.data(), + b, l, c, + m, n, + h, w, d + ); + + })); + + return grad_feats_src_dst; +} diff --git a/src/models/basic/voteflow_plugin/hough_transformation/cpp_im2ht/ht_cuda.h b/src/models/basic/voteflow_plugin/hough_transformation/cpp_im2ht/ht_cuda.h new file mode 100755 index 0000000..f8e0300 --- /dev/null +++ b/src/models/basic/voteflow_plugin/hough_transformation/cpp_im2ht/ht_cuda.h @@ -0,0 +1,29 @@ +#pragma once +#include + + +at::Tensor +ht_cuda_forward( + const at::Tensor &feats_src_dst, + const at::Tensor &voxels_src, + const at::Tensor &voxels_dst, + const at::Tensor &idxs_src, + const at::Tensor &idxs_dst, + const int h, + const int w, + const int d + ); + +// std::vector +at::Tensor +ht_cuda_backward( + const at::Tensor &grad_output, + const at::Tensor &voxels_src, + const at::Tensor &voxels_dst, + const at::Tensor &idxs_src, + const at::Tensor &idxs_dst, + const int h, + const int w, + const int d + ); + diff --git a/src/models/basic/voteflow_plugin/hough_transformation/cpp_im2ht/im2ht.cpp b/src/models/basic/voteflow_plugin/hough_transformation/cpp_im2ht/im2ht.cpp new file mode 100755 index 0000000..7ef830c --- /dev/null +++ b/src/models/basic/voteflow_plugin/hough_transformation/cpp_im2ht/im2ht.cpp @@ -0,0 +1,64 @@ +#include "im2ht.h" +#include "ht_cuda.h" + +at::Tensor +// std::vector +im2ht_forward( + const at::Tensor &feats, + const at::Tensor &voxels_src, + const at::Tensor &voxels_dst, + const at::Tensor &idxs_src, + const at::Tensor &idxs_dst, + const int h, + const int w, + const int d + ) +{ + if ( + feats.type().is_cuda() && \ + voxels_src.type().is_cuda() && voxels_dst.type().is_cuda() && \ + idxs_src.type().is_cuda() && idxs_dst.type().is_cuda() + ) + { + return ht_cuda_forward(feats, + voxels_src, voxels_dst, + idxs_src, idxs_dst, + h, w, d + ); + } + AT_ERROR("Not implemented on the CPU"); +} + + +at::Tensor +// std::vector +im2ht_backward(const at::Tensor &grad_vol, + const at::Tensor &voxels_src, + const at::Tensor &voxels_dst, + const at::Tensor &idxs_src, + const at::Tensor &idxs_dst, + const int h, + const int w, + const int d + ) +{ + if ( + grad_vol.type().is_cuda() && \ + voxels_src.type().is_cuda() && voxels_dst.type().is_cuda() && \ + idxs_src.type().is_cuda() && idxs_dst.type().is_cuda() + ) + { + return ht_cuda_backward(grad_vol, + voxels_src, voxels_dst, + idxs_src, idxs_dst, + h, w, d + ); + } + AT_ERROR("Not implemented on the CPU"); +} + + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + m.def("im2ht_forward", &im2ht_forward, "Forward pass of ht"); + m.def("im2ht_backward", &im2ht_backward, "Backward pass of ht"); +} diff --git a/src/models/basic/voteflow_plugin/hough_transformation/cpp_im2ht/im2ht.h b/src/models/basic/voteflow_plugin/hough_transformation/cpp_im2ht/im2ht.h new file mode 100755 index 0000000..ccb263d --- /dev/null +++ b/src/models/basic/voteflow_plugin/hough_transformation/cpp_im2ht/im2ht.h @@ -0,0 +1,28 @@ +#pragma once +#include + +// std::vector +at::Tensor +im2ht_forward( + const at::Tensor &feats_src_dst, + const at::Tensor &voxels_src, + const at::Tensor &voxels_dst, + const at::Tensor &idxs_src, + const at::Tensor &idxs_dst, + const int h, + const int w, + const int d + ); + +at::Tensor +// std::vector +im2ht_backward( + const at::Tensor &grad_output, + const at::Tensor &voxels_src, + const at::Tensor &voxels_dst, + const at::Tensor &idxs_src, + const at::Tensor &idxs_dst, + const int h, + const int w, + const int d + ); diff --git a/src/models/basic/voteflow_plugin/hough_transformation/cpp_im2ht/im2ht_cuda.cuh b/src/models/basic/voteflow_plugin/hough_transformation/cpp_im2ht/im2ht_cuda.cuh new file mode 100755 index 0000000..72d9b4e --- /dev/null +++ b/src/models/basic/voteflow_plugin/hough_transformation/cpp_im2ht/im2ht_cuda.cuh @@ -0,0 +1,206 @@ +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +// #include +#include +// #include + +#include "cuda_utils.h" + +template +__global__ void im2ht_cuda_forward_kernel(const int64_t n_threads, + const scalar_t *feats, + scalar_t *vol_ht, + const scalar_t *voxels_src, + const scalar_t *voxels_dst, + const scalar_t *idxs_src, + const scalar_t *idxs_dst, + const int b, const int l, const int c, + const int m, const int n, + const int h, const int w, const int d + ) +{ + CUDA_KERNEL_LOOP(index, n_threads) + { + // [b, l, c, m, n] + int b_tmp = index / l / c / m / n % b; // b + int l_tmp = index / c / m / n % l; // l + int c_tmp = index / m / n % c; // c + int m_tmp = index / n % m; // m + int n_tmp = index % n; // n + + int src_tmp = idxs_src[b_tmp * (l * m) + l_tmp * m + m_tmp]; // point idx in [b, l, m] + if (src_tmp>=0 ) + { + int dst_tmp = idxs_dst[b_tmp * (l * n) + src_tmp * n + n_tmp]; // point idx in [b, l, n] + if (dst_tmp>=0) + { + // Z-Y-X after voxelization + int y_src_tmp = voxels_src[b_tmp * (l * 2) + src_tmp * 2 + 0]; // in [b, l, 2] + int x_src_tmp = voxels_src[b_tmp * (l * 2) + src_tmp * 2 + 1]; // in [b, l, 2] + + int y_dst_tmp = voxels_dst[b_tmp * (l * 2) + dst_tmp * 2 + 0]; // in [b, l, 2] + int x_dst_tmp = voxels_dst[b_tmp * (l * 2) + dst_tmp * 2 + 1]; // in [b, l, 2] + + scalar_t feat_tmp = feats[b_tmp * (l * n * c) + src_tmp * (n * c) + n_tmp * (c) + c_tmp]; // in [b, l, n, c] + + int bin_x = x_dst_tmp - x_src_tmp; + int bin_y = y_dst_tmp - y_src_tmp; + bin_x += w/2; + bin_y += h/2; + + if (y_src_tmp >=0 && x_src_tmp >=0 && y_dst_tmp >= 0 && x_dst_tmp >=0) + { + if (bin_x >=0 && bin_x < w && bin_y >= 0 && bin_y < h) + { + // (0, 0) is at position (h//2, w//2) + // vol_ht: [b, l, h, w]; + int64_t offset_vol_tmp = b_tmp * (l * c * h * w) + l_tmp * (c * h * w) + c_tmp * (h * w) + bin_y * w + bin_x; + atomicAdd(vol_ht+offset_vol_tmp, feat_tmp); + } + } + } + } + } +} + +template +__global__ void im2ht_cuda_backward_kernel(const int n_threads, + scalar_t* grad_feats, + const scalar_t* grad_vol_ht, + const scalar_t *voxels_src, + const scalar_t *voxels_dst, + const scalar_t *idxs_src, + const scalar_t *idxs_dst, + const int b, const int l, const int c, + const int m, const int n, + const int h, const int w, const int d + ) +{ + // todo: coalesce + CUDA_KERNEL_LOOP(index, n_threads) + { + // [b, l, c, m, n] + int b_tmp = index / l / c / m / n % b; // b + int l_tmp = index / c / m / n % l; // l + int c_tmp = index / m / n % c; // c + int m_tmp = index / n % m; // m + int n_tmp = index % n; // n + + int src_tmp = idxs_src[b_tmp * (l * m) + l_tmp * m + m_tmp]; // point idx in [b, l, m] + if (src_tmp>=0 ) + { + int dst_tmp = idxs_dst[b_tmp * (l * n) + src_tmp * n + n_tmp]; // point idx in [b, l, n] + if (dst_tmp>=0) + { + // Z-Y-X after voxelization + int y_src_tmp = voxels_src[b_tmp * (l * 2) + src_tmp * 2 + 0]; // in [b, l, 2] + int x_src_tmp = voxels_src[b_tmp * (l * 2) + src_tmp * 2 + 1]; // in [b, l, 2] + + int y_dst_tmp = voxels_dst[b_tmp * (l * 2) + dst_tmp * 2 + 0]; // in [b, l, 2] + int x_dst_tmp = voxels_dst[b_tmp * (l * 2) + dst_tmp * 2 + 1]; // in [b, l, 2] + + int bin_x = x_dst_tmp - x_src_tmp; + int bin_y = y_dst_tmp - y_src_tmp; + bin_x += w/2; + bin_y += h/2; + + if (y_src_tmp >=0 && x_src_tmp >=0 && y_dst_tmp >= 0 && x_dst_tmp >=0) + { + if (bin_x >=0 && bin_x < w && bin_y >= 0 && bin_y < h) + { + // (0, 0) is at position (h//2, w//2) + // vol_ht: [b, l, c, h, w]; + // output gradients + int offset_vol_tmp = b_tmp * (l * c * h * w) + l_tmp * (c * h * w) + c_tmp * (h * w) + bin_y * w + bin_x; // [b, l, c, h, w] + scalar_t grad_vol_tmp = grad_vol_ht[offset_vol_tmp]; // + // input grads + int offset_feat_tmp = b_tmp * (l * n * c) + src_tmp * (n * c) + n_tmp * (c) + c_tmp; // in [b, l, n, c] + atomicAdd(grad_feats+offset_feat_tmp, grad_vol_tmp); + } + } + } + } + } +} + +// %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% // +template +void im2ht_cuda_forward(cudaStream_t stream, + const scalar_t* feats, + scalar_t* vol_ht, + const scalar_t* voxels_src, + const scalar_t* voxels_dst, + const scalar_t* idxs_src, + const scalar_t* idxs_dst, + const int b, const int l, const int c, + const int m, const int n, + const int h, const int w, const int d + ) +{ + const int num_kernels = b * l * c * m * n; + // printf("dimensions num_kernels=%16d, b=%06d, l=%06d, m=%06d, n=%06d, h=%06d, w= %06d, d= %06d \n", + // num_kernels, b, l, m, n, h, w, d); + im2ht_cuda_forward_kernel + <<>>(num_kernels, + feats, + vol_ht, + voxels_src, + voxels_dst, + idxs_src, + idxs_dst, + b, l, c, + m, n, + h, w, d + ); + + cudaError_t err = cudaGetLastError(); + if (err != cudaSuccess) + { + printf("error in im2ht_gpu_kernel: %s\n", cudaGetErrorString(err)); + } +} + +template +void im2ht_cuda_backward(cudaStream_t stream, + scalar_t* grad_feats, + const scalar_t* grad_vol, + const scalar_t* voxels_src, + const scalar_t* voxels_dst, + const scalar_t* idxs_src, + const scalar_t* idxs_dst, + const int b, const int l, const int c, + const int m, const int n, + const int h, const int w, const int d + ) +{ + + const int num_kernels = b * l * c * m * n; + // ***********************************************************************// + im2ht_cuda_backward_kernel + <<>>(num_kernels, + grad_feats, + grad_vol, + voxels_src, + voxels_dst, + idxs_src, + idxs_dst, + b, l, c, + m, n, + h, w, d + ); + cudaError_t err = cudaGetLastError(); + if (err != cudaSuccess) + { + printf("error in ht2im_gpu_kernel: %s\n", cudaGetErrorString(err)); + } + +} + diff --git a/src/models/basic/voteflow_plugin/hough_transformation/cpp_im2ht/setup.py b/src/models/basic/voteflow_plugin/hough_transformation/cpp_im2ht/setup.py new file mode 100755 index 0000000..6c2521f --- /dev/null +++ b/src/models/basic/voteflow_plugin/hough_transformation/cpp_im2ht/setup.py @@ -0,0 +1,37 @@ +from setuptools import setup +from torch.utils.cpp_extension import BuildExtension, CUDAExtension + +setup( + name='im2ht', + # ext_modules=[ + # CUDAExtension(name='im2ht', sources=['im2ht.cpp', 'ht_cuda.cu'], + # extra_compile_args={'cxx': ['-g'], 'nvcc': ['-arch=sm_60']}), + # ], + ext_modules=[ + CUDAExtension(name='im2ht', sources=['im2ht.cpp', 'ht_cuda.cu']), + ], + cmdclass={ + 'build_ext': BuildExtension + }) + + + +# def load_cpp_ext(ext_name): +# root_dir = os.path.join(os.path.split(__file__)[0]) +# src_dir = os.path.join(root_dir, "cpp") +# tar_dir = os.path.join(src_dir, "build", ext_name) +# os.makedirs(tar_dir, exist_ok=True) +# srcs = glob(f"{src_dir}/*.cu") + glob(f"{src_dir}/*.cpp") + +# with warnings.catch_warnings(): +# warnings.simplefilter("ignore") +# from torch.utils.cpp_extension import load + +# ext = load( +# name=ext_name, +# sources=srcs, +# extra_cflags=["-O3"], +# extra_cuda_cflags=[], +# build_directory=tar_dir, +# ) +# return ext \ No newline at end of file diff --git a/src/models/basic/voteflow_plugin/hough_transformation/ht_cuda.py b/src/models/basic/voteflow_plugin/hough_transformation/ht_cuda.py new file mode 100755 index 0000000..0acc055 --- /dev/null +++ b/src/models/basic/voteflow_plugin/hough_transformation/ht_cuda.py @@ -0,0 +1,52 @@ +import torch.nn as nn +from .im2ht import IM2HT + +class HT_CUDA(nn.Module): + """ + for each voxel in src: + find its m knn voxels in src: + for each voxel_src: + find its n knn voxels in dst; + calculate a transaltion which determines the bin in voting (bin_y, bin_x); + vote incrementally by the amount of correponding src and dst features; + """ + """ mapping from points to histgrams + input: feats_src_dst [b, l, n, c] + input: voxels_src [b, l, 2] + input: voxels_dst [b, l, 2] + input: idxs_src [b, l, m] + input: idxs_dst [b, l, n] + input: bins_x int + input: bins_y int + input: bins_z int + output: HT feaures [b, l, h, w], c*2 stands for the concatenation of features from src and dst. + """ + def __init__(self, h, w, d): + super(HT_CUDA, self).__init__() + self.h, self.w, self.d = h, w, d + assert h%2==0 + assert w%2==0 + self.ht = IM2HT() + + self.__repr__() + + def __repr__(self): + return self.__class__.__name__ + f'( h: {self.h:04d}, w: {self.w:04d}, d: {self.d:04d} )' + + def forward(self, feats_src_dst, voxels_src, voxels_dst, idxs_src, idxs_dst): + # print('ht_cuda forward: ', feats.shape, idxs_fps.shape, idxs_src.shape, idxs_dst.shape, bins_x.shape, bins_y.shape, bins_z.shape) + # print(f'( k:{self.k:04d}, m: {self.m:04d}, n: {self.n:04d}, h: {self.h:04d}, w: {self.w:04d}, d: {self.d:04d} )') + assert feats_src_dst.dim()==4 + assert voxels_src.shape == voxels_dst.shape + + b, l, n = idxs_dst.shape + _, _, m = idxs_src.shape + + datatype = feats_src_dst.dtype + voxels_src = voxels_src.to(datatype) + voxels_dst = voxels_dst.to(datatype) + idxs_src = idxs_src.to(datatype) + idxs_dst = idxs_dst.to(datatype) + vol = self.ht(feats_src_dst, voxels_src, voxels_dst, idxs_src, idxs_dst, self.h, self.w, self.d) + # print('vol: ', vol.shape, vol.min(), vol.max()) + return vol diff --git a/src/models/basic/voteflow_plugin/hough_transformation/im2ht.py b/src/models/basic/voteflow_plugin/hough_transformation/im2ht.py new file mode 100755 index 0000000..e4d12e7 --- /dev/null +++ b/src/models/basic/voteflow_plugin/hough_transformation/im2ht.py @@ -0,0 +1,103 @@ +import os +import warnings +from glob import glob + +from torch import nn +from torch.autograd import Function +from torch.autograd.function import once_differentiable + +def load_cpp_ext(ext_name): + root_dir = os.path.join(os.path.split(__file__)[0]) + src_dir = os.path.join(root_dir, "cpp_im2ht") + tar_dir = os.path.join(src_dir, "build", ext_name) + os.makedirs(tar_dir, exist_ok=True) + srcs = glob(f"{src_dir}/*.cu") + glob(f"{src_dir}/*.cpp") + + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + from torch.utils.cpp_extension import load + ext = load( + name=ext_name, + sources=srcs, + extra_cflags=["-O3"], + extra_cuda_cflags=[], + build_directory=tar_dir, + ) + return ext + +# defer calling load_cpp_ext to make CUDA_VISIBLE_DEVICES happy +im2ht = None + + +class IM2HTFunction(Function): + @staticmethod + def forward(ctx, + feats_src_dst, + voxels_src, + voxels_dst, + idxs_src, + idxs_dst, + h, w, d + ): + + ctx.h = h + ctx.w = w + ctx.d = d + ctx.save_for_backward(voxels_src, voxels_dst, idxs_src, idxs_dst) + vol = im2ht.im2ht_forward( + feats_src_dst, + voxels_src, + voxels_dst, + idxs_src, + idxs_dst, + h, + w, + d + ) + return vol + + @staticmethod + @once_differentiable + def backward(ctx, grad_vol): + voxels_src, voxels_dst, idxs_src, idxs_dst = ctx.saved_tensors # it is a list of saved tensors! + + grad_input_src_dst = im2ht.im2ht_backward( + grad_vol, + voxels_src, + voxels_dst, + idxs_src, + idxs_dst, + ctx.h, + ctx.w, + ctx.d + ) + + return ( + grad_input_src_dst, + None, # voxels_src, + None, # voxels_dst, + None, # idxs_src, + None, # idxs_dst, + None, # h + None, # w + None # d + ) + + +class IM2HT(nn.Module): + def __init__(self, ): + super(IM2HT, self).__init__() + + global im2ht + im2ht = load_cpp_ext("im2ht") + + def forward(self, feats_src_dst, voxels_src, voxels_dst, idxs_src, idxs_dst, h, w, d): + # print('IM2HT forward', feats.shape, idxs_fps.shape, idxs_src.shape, idxs_dst.shape, bins_x.shape, bins_y.shape, bins_z.shape) + return IM2HTFunction.apply( + feats_src_dst.contiguous(), + voxels_src.contiguous(), + voxels_dst.contiguous(), + idxs_src.contiguous(), + idxs_dst.contiguous(), + h, w, d + ) diff --git a/src/models/basic/voteflow_plugin/hough_transformation/test_backward.py b/src/models/basic/voteflow_plugin/hough_transformation/test_backward.py new file mode 100755 index 0000000..a018c3c --- /dev/null +++ b/src/models/basic/voteflow_plugin/hough_transformation/test_backward.py @@ -0,0 +1,73 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch.autograd import gradcheck + +import random +import numpy as np +import math +import scipy.io as sio +import glob + +from models.ht.ht_cuda import HT_CUDA +from models.ht.im2ht import IM2HTFunction + +random.seed(0) +np.random.seed(0) +torch.manual_seed(0) + + +if torch.cuda.is_available(): + # device_name = "cuda" + # torch.backends.cudnn.benchmark = True + # torch.backends.cudnn.deterministic = True + # torch.cuda.manual_seed(0) + + device_name = "cuda" + torch.cuda.manual_seed(0) + # torch.backends.cudnn.enabled = False #greatly slow down the speed + torch.backends.cudnn.deterministic = True + torch.cuda.manual_seed(0) + torch.cuda.empty_cache() + print("Let's use", torch.cuda.device_count(), "GPU(s)!") +else: + device_name = "cpu" + print("CUDA is not available") +device = torch.device(device_name) +print('device: ', device) + + +b = 2 +l = 64 +c=3 +m = 4 +n = 8 +max_x=0.5 +max_y=0.5 +max_z=0.2 +voxel=0.1 +nx = math.ceil(max_x / voxel) * 2 # assume 120km/h, along both directions +ny = math.ceil(max_y / voxel) * 2 # assume 120km/h, along both directions +nz = math.ceil(max_z / voxel) * 2 + +vote = HT_CUDA(nx, ny, nz) +vote = vote.double().to(device) + +print('grad check***********') + +feats = torch.rand(size=(b, l, n, 3), requires_grad=True).double().to(device) +voxels_src = torch.randint(size=(b, l, 2), low=0, high=nx).double().to(device) +voxels_dst = torch.randint(size=(b, l, 2), low=0, high=nx).double().to(device) +idxs_src = torch.randint(0, l, size=(b, l, m)).double().to(device) +idxs_dst = torch.randint(0, l, size=(b, l, n)).double().to(device) +print('input shapes: ', feats.shape, voxels_src.shape, voxels_dst.shape, idxs_src.shape, idxs_dst.shape) +# grad_output = torch.randn(vol.shape).double().to(device) +# vol.backward(gradient=grad_output) +# grad_input = fs_window.grad +# print(grad_input.shape) + +# only able to test inputs with tiny sizes +res = gradcheck(vote, (feats, voxels_src, voxels_dst, idxs_src, idxs_dst), raise_exception=True) +# res = gradcheck(IM2HTFunction.apply, (feats, bins_x, bins_y, bins_z, flags, idxs, m, n, h, w, d), raise_exception=True) +# # res=gradcheck(myconv, input, eps=1e-3, atol=1e-3, rtol=1e-2, raise_exception=True) +print('grad check', res) diff --git a/src/models/basic/voteflow_plugin/hough_transformation/test_forward.py b/src/models/basic/voteflow_plugin/hough_transformation/test_forward.py new file mode 100644 index 0000000..3d5af04 --- /dev/null +++ b/src/models/basic/voteflow_plugin/hough_transformation/test_forward.py @@ -0,0 +1,247 @@ +import numpy as np +import random +import glob +import warnings +import parmap +import os +import itertools +import shutil +import argparse +import copy +import sys +import yaml +import json +import gc +import datetime, time +import open3d as o3d +import seaborn as sns +import multiprocessing +import matplotlib.pyplot as plt +import plotly +from tqdm import tqdm +import torch +import torch.nn as nn +import pytorch3d.ops as pytorch3d_ops +import plotly.graph_objs as go +from util_model import Decoder +from util_loss import warped_pc_loss +from utils_visualization import visualize_pcd +warnings.filterwarnings('ignore') +np.set_printoptions(suppress=True) +# add cuda implementation +from ht.ht_cuda import HT_CUDA + +# this loads a pair of frames and gt flow. Note 1). the source has been compensated by ego motion; and 2) ground has been removed. +def dataloader_minimal(data_path): + data = np.load(data_path) + pcl_0 = data['pc1'] + pcl_1 = data['pc2'] + valid_0 = data['pc1_flows_valid_idx'] + valid_1 = data['pc2_flows_valid_idx'] + flow_0_1 = data['gt_flow_0_1'] + class_0 = data['pc1_classes'] + class_1 = data['pc2_classes'] + + pcl_0 = pcl_0[valid_0] + pcl_1 = pcl_1[valid_1] + flow_0_1 = flow_0_1[valid_0] + class_0 = class_0[valid_0] + class_1 = class_1[valid_1] + # print('pc range: ', + # pcl_0[:, 0].min(), pcl_0[:, 0].max(), + # pcl_0[:, 1].min(), pcl_0[:, 1].max(), + # pcl_1[:, 0].min(), pcl_1[:, 0].max(), + # pcl_1[:, 1].min(), pcl_1[:, 1].max(), + # ) + # print('class 0: ', np.unique(class_0)) + # visualize_pcd( + # np.concatenate([pcl_0, pcl_1, pcl_0+flow_0_1], axis=0), + # np.concatenate([np.zeros(len(pcl_0))+1, np.zeros(len(pcl_1))+2, np.zeros(len(pcl_0))+0], axis=0), + # num_colors=3, + # title=f'visualize input: src-g, dst-b, src+flow-r: {data_path}' + # ) + data_dict = { + 'point_src': pcl_0, + 'point_dst': pcl_1, + 'scene_flow': flow_0_1, + 'data_path': data_path + } + return data_dict + +if __name__ == "__main__": + + # Initialization + random.seed(0) + np.random.seed(0) + # fix open3d seed. question mark there: sometimes o3d outputs different results. + o3d.utility.random.seed(0) # only in open3d>=0.16 + multiprocessing.set_start_method('forkserver') # or 'spawn' + + torch.manual_seed(0) + if torch.cuda.is_available(): + assert torch.cuda.device_count()==1 + device_name = f'cuda:{torch.cuda.current_device()}' + torch.backends.cudnn.deterministic = True + torch.cuda.manual_seed(0) + torch.cuda.empty_cache() + gc.collect() + torch.set_default_tensor_type("torch.cuda.FloatTensor") + print(f"Let's use {torch.cuda.device_count()}, {torch.cuda.get_device_name()} GPU(s)!") + else: + print("CUDA is not available") + device_name = 'cpu' + device = torch.device(device_name) + print(f'device: {device}') + + model = Decoder().to(device) + b = 4 + l = 40000 + m = 9 # neighbors in src + n =16 # neighbors in dst + vote = HT_CUDA(m, n, 3.3, 3.3, 0.1, 0.1).to(device) + + files = glob.glob(os.path.join('demo.npz')) + print('total files: ', len(files)) + for file in files: + data = dataloader_minimal(file) + points_src = data['point_src'] + points_dst = data['point_dst'] + + # random downsampling + points_src = points_src[np.random.choice(np.arange(0, len(points_src)), l, replace=False)] + points_dst = points_dst[np.random.choice(np.arange(0, len(points_src)), l, replace=False)] + + # coordinates + points_src = torch.from_numpy(points_src).float().to(device) + points_dst = torch.from_numpy(points_dst).float().to(device) + print('points: ', points_src.shape, points_dst.shape) + + # feats. this is a minimal exmaple. in pratice, they are learned from some backbone model; + feats_src = torch.randn((len(points_src), 64)).float().to(device) + feats_dst = torch.randn((len(points_dst), 64)).float().to(device) + feats_src /= torch.linalg.norm(feats_src, dim=-1, keepdims=True) + feats_dst /= torch.linalg.norm(feats_dst, dim=-1, keepdims=True) + print('feats: ', feats_src.shape, feats_dst.shape) + + # batch, think about how to batch point clouds of various length + points_src, points_dst = points_src[None].expand(b, l, 3).contiguous(), points_dst[None].expand(b, l, 3).contiguous() + feats_src, feats_dst = feats_src[None].expand(b, l, 64).contiguous(), feats_dst[None].expand(b, l, 64).contiguous() + + # only need to calculate once per batch + with torch.no_grad(): + + # process coordinates + # knn of src in dst + knn_dst = pytorch3d_ops.knn_points(points_src, points_dst, lengths1=None, lengths2=None, K=16, return_nn=True, return_sorted=False) + dists_dst, idxs_dst, pts_dst = knn_dst[0], knn_dst[1], knn_dst[2] + print('dst: ', dists_dst.shape, idxs_dst.shape, pts_dst.shape) + + # translations, between src and its knn in dst + ts_dst_src = pts_dst - points_src[:, :, None, :] + print('translations: ', ts_dst_src.shape) + + # define a local window for each point in src, using knn for simplicity + knn_src = pytorch3d_ops.knn_points(points_src, points_src, lengths1=None, lengths2=None, K=9, return_nn=True, return_sorted=False) + dists_src, idxs_src, pts_src = knn_src[0], knn_src[1], knn_src[2] + print('src: ', dists_src.shape, idxs_src.shape, pts_src.shape) + + # all translations within in a window, between src and its knn in dst + ts_window = pytorch3d_ops.knn_gather(ts_dst_src.view(b, l, -1), idxs_src, lengths=None) + ts_window = ts_window.view(idxs_src.shape[0], idxs_src.shape[1], idxs_src.shape[2], 16, 3) + print('ts_window: ', ts_window.shape) + + # quantization + n_intervals = 3.33 // 0.1 # both directions + intervals = torch.arange(-0.1*n_intervals, 0.1*n_intervals + 1e-8, 0.1) # [-3.3, -3.2, ..., 3.2, 3.3] + intervals_z = torch.arange(-0.1, 0.1 + 1e-8, 0.1) # [-0.1, 0, 0.1] + print('intervals xy: ', n_intervals, 0.1*n_intervals, intervals.shape, intervals, intervals_z) + + # tricky there, check: https://pytorch.org/docs/stable/generated/torch.bucketize.html + x_idxs = torch.bucketize(ts_window[:, :, :, :, 0], intervals, right=True) + y_idxs = torch.bucketize(ts_window[:, :, :, :, 1], intervals, right=True) + z_idxs = torch.bucketize(ts_window[:, :, :, :, 2], intervals_z, right=True) + x_idxs -= 1 + y_idxs -= 1 + z_idxs -= 1 + print('ts_idxs: ', x_idxs.shape, y_idxs.shape, z_idxs.shape) + + # # # sanity check: find out shared translations + # for i in range(0, x_idxs.shape[0]): + # for j in range(0, x_idxs.shape[1]): + # x_idxs_temp = x_idxs[i, j] + # y_idxs_temp = y_idxs[i, j] + # z_idxs_temp = z_idxs[i, j] + # for p in range(0, ts_window.shape[2]): + # for q in range(0, ts_window.shape[3]): + # print('t, idx: ', ts_window[i, j, p, q], x_idxs[i, j, p, q], y_idxs[i, j, p, q], z_idxs[i, j, p, q]) + # idxs_temp = torch.stack([x_idxs_temp, y_idxs_temp, z_idxs_temp], dim=2) + # print('idxs_temp: ', idxs_temp.shape) + # bin_unq, count_unq = torch.unique(idxs_temp.view(-1, 3), dim=0, return_counts=True) + # print('unique bins: ', bin_unq, count_unq) + + # process feats: gradient backpropagation during training + # feats: knn of src in dst + feats_dst_nn = pytorch3d_ops.knn_gather(feats_dst, idxs_dst, lengths=None) + print('feat dst nn: ', feats_dst_nn.shape) + + # feats similarity, between src and its knn in dst + feats = (feats_src[:, :, None, :] * feats_dst_nn).sum(-1) + print('feats: ', feats.shape) + + # all feats within in a window, between src and its knn in dst + feats_window = pytorch3d_ops.knn_gather(feats, idxs_src, lengths=None) + print('feats_window: ', feats_window.shape) + + # # simple test: no feature learning + feats_window.fill_(1.0) + + # # # Hough voting: pytorch implemenation vs custom cuda implementation + # 1. test pytorch implementation + start_time_pytorch = time.time() + b_idxs = torch.arange(0, b)[:, None, None, None].expand(b, l, 9, 16).flatten() + p_idxs = torch.arange(0, l)[None, :, None, None].expand(b, l, 9, 16).flatten() + x_idxs = x_idxs.flatten() + y_idxs = y_idxs.flatten() + z_idxs = z_idxs.flatten() + + # delete out-of-range values later + valid = torch.all( + torch.stack([x_idxs=0, y_idxs=0, z_idxs=0], dim=1), + dim=1).flatten() + print('valid: ', valid.shape, valid.sum()) + # build a 3d volume (x/y: -3.3 +3.3, z:-0.1 +0.1) + vol = torch.zeros((b, l, len(intervals)-1, len(intervals)-1, len(intervals_z)-1)).float() + + # incremental voting: a 3d representation per point + # https://discuss.pytorch.org/t/indexing-with-repeating-indices-numpy-add-at/10223/13 + # option 1: + # idxs = b_idxs * vol.shape[1:].numel() + p_idxs * vol.shape[2:].numel() + x_idxs * vol.shape[3:].numel() + y_idxs * vol.shape[4:].numel() + z_idxs # flatten to 1d indices + # print('flatten to 1d indices: ', vol.shape, vol.shape[1:].numel(), vol.shape[2:].numel(), vol.shape[3:].numel(), vol.shape[4:].numel()) + # feats_window = feats_window.flatten() + # print('idxs, feats_window: ', idxs.shape, feats_window.shape) + # vol.flatten().put_(idxs[valid], feats_window[valid], accumulate=True) + # vol.flatten().put_(idxs[valid], torch.ones(sum(valid)).float(), accumulate=True) + # vol = vol.view(b, n_pts, len(intervals)-1, len(intervals)-1, len(intervals_z)-1) + # option 2: + vol.index_put_([b_idxs[valid], p_idxs[valid], x_idxs[valid], y_idxs[valid], z_idxs[valid]], feats_window.flatten()[valid], accumulate=True) + print('vote: ', vol.shape, vol.max(), vol.min(), time.time()-start_time_pytorch) + + # 2. test custom cuda implementation + start_time_cuda = time.time() + vol2, _, _, _ = vote(feats_window, ts_window) + vol2 = vol2.permute(0, 1, 3, 4, 2) + print('vote2: ', vol2.shape, vol2.max(), vol2.min(), time.time()-start_time_cuda) + + # check result + print('vol==vol2: ', torch.equal(vol, vol2), (vol-vol2).abs().max()) + + # add conv + fc layers for final prediction + flow = model.forward(vol.view(b*l, len(intervals)-1, len(intervals)-1, len(intervals_z)-1).permute(0, 3, 1, 2).contiguous()) + print('flow: ', flow.shape) + + # calculate loss during training + loss = warped_pc_loss(points_src.view(-1, 3)+flow, points_dst.view(-1, 3)) + print('loss: ', loss.shape, loss) + + + diff --git a/src/models/basic/voteflow_plugin/utils.py b/src/models/basic/voteflow_plugin/utils.py new file mode 100644 index 0000000..e823fd1 --- /dev/null +++ b/src/models/basic/voteflow_plugin/utils.py @@ -0,0 +1,45 @@ +import torch +import pytorch3d.ops as pytorch3d_ops + +def calculate_unq_voxels(coords, image_dims): + unqs, idxs = torch.unique(coords[:, 1]*image_dims[0]+coords[:, 2], return_inverse=True, sorted=True) + unqs_voxel = torch.stack([unqs//image_dims[0], unqs%image_dims[0]], dim=1) + return unqs_voxel, idxs + +# https://pytorch3d.readthedocs.io/en/latest/_modules/pytorch3d/ops/utils.html +def batched_masked_gather(x: torch.Tensor, idxs: torch.Tensor, mask: torch.Tensor,fill_value=-1.0) -> torch.Tensor: + assert x.dim() == 3 # [b, m, c] + assert idxs.dim() == 3 # [b, n, k] + assert idxs.shape == mask.shape + b, m, c = x.shape + + idxs_masked = idxs.clone() + idxs_masked[~mask] = 0 + l, n, k = idxs.shape + y = pytorch3d_ops.knn_gather(x, idxs_masked) # [b, n, k, c] + y[~mask, :] = fill_value + + return y + +def pad_to_batch(idxs, l): + if idxs is None: + return idxs + if idxs.dim()==2: + assert idxs.shape[0]<=l + if idxs.shape[0] torch.Tensor: + x = self.conv(x) + x = F.relu(x) + return x + +class ConvBNBlock(nn.Module): + def __init__(self, in_num_channels: int, out_num_channels: int, + kernel_size=3, stride=1, padding=1): + super().__init__() + self.conv = nn.Conv2d(in_num_channels, out_num_channels, kernel_size, stride, padding, bias=True) + self.bn = nn.BatchNorm2d(out_num_channels) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = self.conv(x) + x = self.bn(x) + x = F.relu(x) + return x + +class VolConvBN(nn.Module): + def __init__(self, h, w, hidden_dim=16, dim_output=64): + super().__init__() + assert h%2==0 + assert w%2==0 + self.conv1 = ConvBNBlock(in_num_channels=1, out_num_channels=hidden_dim, stride=2) + self.conv2 = ConvBNBlock(in_num_channels=hidden_dim, out_num_channels=hidden_dim, stride=2) + self.linear = nn.Linear(math.ceil(h/4) * math.ceil(w/4) * hidden_dim, dim_output) + self.relu = nn.ReLU() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + b, l, _, h, w = x.shape + x = self.conv1(x.view(b*l, 1, h, w)) + x = self.conv2(x) + # print(x.shape) + x = self.linear(x.view(b*l, -1)) + # print('after linear:', x.shape) + x = self.relu(x) + return x.view(b, l, -1) + + +class VoteFlowLinearDecoder(nn.Module): + def __init__(self, dim_input=16, layer_size=1, filter_size=128): + super().__init__() + # self.linear = nn.Linear(m*dim_input, dim_output) + offset_encoder_channels = 128 + self.offset_encoder = nn.Linear(3, offset_encoder_channels) + filter_size=filter_size + + decoder = nn.ModuleList() + decoder.append(torch.nn.Linear(dim_input + offset_encoder_channels, filter_size)) + decoder.append(torch.nn.GELU()) + for _ in range(layer_size-1): + decoder.append(torch.nn.Linear(filter_size, filter_size)) + decoder.append(torch.nn.ReLU()) + decoder.append(torch.nn.Linear(filter_size, 3)) + + self.decoder = nn.Sequential(*decoder) + print(self.decoder) + + def forward(self, x: torch.Tensor, pts_offsets: torch.Tensor) -> torch.Tensor: + b, l, _ = x.shape + pts_offsets_feats = self.offset_encoder(pts_offsets) + x = torch.cat([x, pts_offsets_feats], dim=-1) + # print('decoder conv: ', x.shape) + x = self.decoder(x) + + return x diff --git a/src/models/voteflow.py b/src/models/voteflow.py new file mode 100644 index 0000000..0b708d3 --- /dev/null +++ b/src/models/voteflow.py @@ -0,0 +1,375 @@ +""" +This file is directly copied from: +https://github.com/tudelft-iv/VoteFlow + +with slightly modification to have unified format with all benchmark. +""" + + +import math +import numpy as np + +import dztimer +import torch +import torch.nn as nn +import pytorch3d.ops as pytorch3d_ops + +from .basic.encoder import DynamicEmbedder +from .basic.unet import FastFlow3DUNet +from .basic import cal_pose0to1 + +from .basic.voteflow_plugin.hough_transformation import HT_CUDA +from .basic.voteflow_plugin.voteflow_module import VolConvBN, VoteFlowLinearDecoder +from .basic.voteflow_plugin.utils import calculate_unq_voxels, batched_masked_gather, pad_to_batch + +import warnings +warnings.filterwarnings('ignore') +np.set_printoptions(suppress=True) + +class VoteFlow(nn.Module): + def __init__(self, + nframes=1, + m=8, + n=128, + using_voting=True, + input_channels=32, + output_channels=64, + point_cloud_range = [-51.2, -51.2, -3, 51.2, 51.2, 3], + voxel_size=(0.2, 0.2, 6), + grid_feature_size = [512, 512], + decoder_layers=4, + use_ball_query=False, + vol_conv_hidden_dim=16, + **kwargs): + super().__init__() + + assert len(point_cloud_range)==6 + assert len(voxel_size)==3 + + # a hack there, may and may not have any impact depends on the setting. + assert voxel_size[0]==voxel_size[1] + self.using_voting = using_voting + print('Using voting:', self.using_voting) + + # General architecture: + pseudo_image_dims = grid_feature_size[:2] #int((point_cloud_range[3]-point_cloud_range[0])/voxel_size[0]), int((point_cloud_range[4]-point_cloud_range[1])/voxel_size[1])) # ignore z dimension + + self.point_cloud_range = point_cloud_range + self.voxel_size = voxel_size + self.pseudo_image_dims = pseudo_image_dims + self.backbone = FastFlow3DUNet() ## output_channel 64 + + if self.using_voting: + self.decoder = VoteFlowLinearDecoder(dim_input= output_channels * 2 + input_channels * 2, layer_size=decoder_layers) + print('decoder:', self.decoder) + else: + self.decoder = VoteFlowLinearDecoder(dim_input= output_channels + input_channels * 2, layer_size=decoder_layers) + print('decoder:', self.decoder) + + + self.embedder = DynamicEmbedder(voxel_size=voxel_size, + pseudo_image_dims=pseudo_image_dims, + point_cloud_range=point_cloud_range, + feat_channels=input_channels) + + # Voting Module settings + if self.using_voting: + ## how many bins inside a voxel after quantization + ## assume 72km/h (20m/s), along x/y; 0.1m along z + e = 1e-8 + radius = 2 + nx = math.ceil((radius+e)*nframes / voxel_size[0]) + ny = math.ceil((radius+e)*nframes / voxel_size[1]) + nz = math.ceil((0.1+e) / voxel_size[2]) # +/-0.1 + + self.nx = nx*2 # +/-x + self.ny = ny*2 # +/-y + self.nz = nz*2 # +/-z + print('n x/y/z: ', self.nx, self.ny, self.nz) + + self.nframes = nframes + self.m = m # m knn within src, for each src point + self.using_ball_query = use_ball_query + self.n = n # n knn between src and dst, for each src voxel + self.radius_dst = max(nx, ny) # define a search window for a src voxel in dst voxels for calculating translations + + if self.using_ball_query: + self.radius_src = math.ceil(max((radius+e)/voxel_size[0], (radius+e)/voxel_size[1])) # define a search window (in meters) within src voxels, aka the rigid motion window + print(f'using ball query to search window radius in source pc: {self.radius_src}, m={self.m};') + else: + print(f'using knn to search window radius in source pc, m={self.m};') + print(f'using ball query to search in target pc, n={self.n}, search window radius: {self.radius_dst}.') + + self.vote = HT_CUDA(self.ny, self.nx, self.nz) + + self.volconv = VolConvBN(self.ny, self.nx, hidden_dim=vol_conv_hidden_dim, dim_output=output_channels) + + self.timer = dztimer.Timing() + self.timer.start("Total") + + def process_points_per_pair(self, voxel_info_src, voxel_info_dst): + valid_point_idxs_src = voxel_info_src['point_idxes'] # [N_valid_pts] + valid_point_idxs_dst = voxel_info_dst['point_idxes'] + valid_voxel_coords_src = voxel_info_src['voxel_coords'] #[N_valid_pts, 3] + valid_voxel_coords_dst = voxel_info_dst['voxel_coords'] + + unq_voxel_coords_src, point_voxel_idxs_src = calculate_unq_voxels(valid_voxel_coords_src, self.pseudo_image_dims) + unq_voxel_coords_dst, point_voxel_idxs_dst = calculate_unq_voxels(valid_voxel_coords_dst, self.pseudo_image_dims) + + # unq_voxel_coords_src # [N_valid_voxels, 2] + # point_voxel_idxs_src # [N_valid_pts], the index of the unique voxel for each point + if self.using_voting: + dists_dst, knn_idxs_dst, _ = pytorch3d_ops.ball_query(unq_voxel_coords_src[None].float(), unq_voxel_coords_dst[None].float(), lengths1=None, lengths2=None, K=self.n, return_nn=False, radius=self.radius_dst) + if self.using_ball_query: + dists_src, knn_idxs_src, _ = pytorch3d_ops.ball_query(unq_voxel_coords_src[None].float(), unq_voxel_coords_src[None].float(), lengths1=None, lengths2=None, K=self.m, return_nn=False, radius=self.radius_src) + else: + dists_src, knn_idxs_src, _ = pytorch3d_ops.knn_points(unq_voxel_coords_src[None].float(), unq_voxel_coords_src[None].float(), lengths1=None, lengths2=None, K=self.m, return_nn=False, return_sorted=False) + else: + knn_idxs_dst=[None] + knn_idxs_src=[None] + + return valid_point_idxs_src, valid_point_idxs_dst, point_voxel_idxs_src, point_voxel_idxs_dst, unq_voxel_coords_src, unq_voxel_coords_dst, knn_idxs_src[0], knn_idxs_dst[0] + + def preprocessing(self, points_src, points_dst, voxel_info_list_src, voxel_info_list_dst): + point_masks_src = [] + point_masks_dst = [] + point_voxel_idxs_src = [] + point_voxel_idxs_dst = [] + point_offsets_src = [] + + unq_voxels_src = [] + unq_voxels_dst = [] + knn_idxs_src = [] + knn_idxs_dst = [] + + l_points = 0 # for padding + l_voxels = 0 # for padding + for points_src_, points_dst_, voxel_info_src, voxel_info_dst in zip(points_src, points_dst, voxel_info_list_src, voxel_info_list_dst): + point_idxs_src_, point_idxs_dst_, \ + point_voxel_idxs_src_, point_voxel_idxs_dst_, \ + unq_voxels_src_, unq_voxels_dst_, \ + knn_idxs_src_, knn_idxs_dst_ = self.process_points_per_pair(voxel_info_src, voxel_info_dst) + # print('process_points_per_pair: ', point_idxs_src_.shape, point_voxel_idxs_src_.shape, unq_voxels_src_.shape, knn_idxs_dst_.shape) + + # generate a mask for easy batching and sequential processing + point_masks_src_ = torch.zeros((len(points_src_)), device=points_src_.device, dtype=points_src_.dtype) + point_masks_src_[point_idxs_src_] = 1 + point_masks_dst_ = torch.zeros((len(points_dst_)), device=points_dst_.device, dtype=points_dst_.dtype) + point_masks_dst_[point_idxs_dst_] = 1 + point_masks_src.append(point_masks_src_) + point_masks_dst.append(point_masks_dst_) + + point_voxel_idxs_src.append(point_voxel_idxs_src_) + point_voxel_idxs_dst.append(point_voxel_idxs_dst_) + point_offsets_src.append(voxel_info_src['point_offsets']) + unq_voxels_src.append(unq_voxels_src_) + unq_voxels_dst.append(unq_voxels_dst_) + knn_idxs_src.append(knn_idxs_src_) + knn_idxs_dst.append(knn_idxs_dst_) + + l_voxels = max(l_voxels, max(unq_voxels_src_.shape[0], unq_voxels_dst_.shape[0])) + l_points = max(l_points, max(point_voxel_idxs_src_.shape[0], point_voxel_idxs_dst_.shape[0])) + # print('l', l_points, l_voxels) + + # padding + for i, (point_voxel_idxs_src_, point_voxel_idxs_dst_, point_offsets_src_, unq_voxels_src_, unq_voxels_dst_, knn_idxs_src_, knn_idxs_dst_) \ + in enumerate( zip(point_voxel_idxs_src, point_voxel_idxs_dst, point_offsets_src, unq_voxels_src, unq_voxels_dst, knn_idxs_src, knn_idxs_dst) ): + unq_voxels_src_= pad_to_batch(unq_voxels_src_, l_voxels) + unq_voxels_dst_ = pad_to_batch(unq_voxels_dst_, l_voxels) + point_voxel_idxs_src_= pad_to_batch(point_voxel_idxs_src_, l_points) + point_voxel_idxs_dst_= pad_to_batch(point_voxel_idxs_dst_, l_points) + point_offsets_src_ = pad_to_batch(point_offsets_src_, l_points) + + knn_idxs_src_ = pad_to_batch(knn_idxs_src_, l_voxels) + knn_idxs_dst_= pad_to_batch(knn_idxs_dst_, l_voxels) + + unq_voxels_src[i] = unq_voxels_src_ + unq_voxels_dst[i] = unq_voxels_dst_ + point_voxel_idxs_src[i] = point_voxel_idxs_src_ + point_voxel_idxs_dst[i] = point_voxel_idxs_dst_ + point_offsets_src[i] = point_offsets_src_ + + knn_idxs_src[i] = knn_idxs_src_ + knn_idxs_dst[i] = knn_idxs_dst_ + + unq_voxels_src = torch.stack(unq_voxels_src, dim=0) + unq_voxels_dst = torch.stack(unq_voxels_dst, dim=0) + point_voxel_idxs_src = torch.stack(point_voxel_idxs_src, dim=0) + point_voxel_idxs_dst = torch.stack(point_voxel_idxs_dst, dim=0) + point_offsets_src = torch.stack(point_offsets_src, dim=0) + point_masks_src = torch.stack(point_masks_src, dim=0) + point_masks_dst = torch.stack(point_masks_dst, dim=0) + + if self.using_voting: + knn_idxs_src = torch.stack(knn_idxs_src, dim=0) + knn_idxs_dst = torch.stack(knn_idxs_dst, dim=0) + else: + knn_idxs_src = None + knn_idxs_dst = None + + return point_masks_src, point_masks_dst, point_voxel_idxs_src, point_voxel_idxs_dst, point_offsets_src, unq_voxels_src, unq_voxels_dst, knn_idxs_src, knn_idxs_dst + + def extract_voxel_from_image(self, image, voxels): + # image: [b, c, h, w]; voxels: [b, num, 2] + idxs = voxels[:, :, 0] * self.pseudo_image_dims[0] + voxels[:, :, 1] + + mask = voxels.min(-1)[0]>=0 + + b, c, h, w = image.shape + + feats_per_voxel = batched_masked_gather(image.view(b, c, h*w).permute(0, 2, 1), idxs[:, :, None].long(), mask[:, :, None], fill_value=0) + + # print('point per voxel: ', feats_per_voxel.shape) + return feats_per_voxel[:, :, 0, :] # [b, num , c] + + def extract_point_from_voxel(self, voxels, idxs): + # voxels: [b, l, c]; idxs: [b, num] + feats_per_point = batched_masked_gather(voxels, idxs[:,:,None].long(), idxs[:, :, None]>=0, fill_value=0) + # print('point per point: ', feats_per_point.shape) + return feats_per_point[:, :, 0, :]# [b, num, c] + + def extract_point_from_image(self, image, voxels, point_idxs): + feats_per_voxel = self.extract_voxel_from_image(image, voxels) + feats_per_point = self.extract_point_from_voxel(feats_per_voxel, point_idxs) + # print('point per point merged: ', feats_per_point.shape) + return feats_per_point + + # adapted from zeroflow + def concat_feats(self, feats_vote, feats_before, feats_after): + feats = torch.cat([feats_vote, feats_before, feats_after], dim=-1) + return feats + + def _model_forward(self, points_src, points_dst): + # assert points_src.shape==points_dst.shape + + self.timer[1][0].start("Voxelization") + pseudoimages_src, voxel_infos_lst_src = self.embedder(points_src) + pseudoimages_dst, voxel_infos_lst_dst = self.embedder(points_dst) + self.timer[1][0].stop() + + self.timer[1][1].start("Preprocessing") + with torch.no_grad(): + point_masks_src, point_masks_dst, \ + point_voxel_idxs_src, point_voxel_idxs_dst, \ + point_offsets_src, \ + voxels_src, voxels_dst, \ + knn_idxs_src, knn_idxs_dst = self.preprocessing(points_src, points_dst, voxel_infos_lst_src, voxel_infos_lst_dst) + self.timer[1][1].stop() + + self.timer[1][2].start("Feature extraction") + pseudoimages_grid = self.backbone(pseudoimages_src, pseudoimages_dst) + + feats_point_src_init = self.extract_point_from_image(torch.cat([pseudoimages_src, pseudoimages_dst], dim=1), voxels_src, point_voxel_idxs_src) + feats_point_src_grid = self.extract_point_from_image(pseudoimages_grid, voxels_src, point_voxel_idxs_src) + + self.timer[1][2].stop() + + if self.using_voting: + self.timer[1][3].start("Gathering") + + feats_voxel_src = self.extract_voxel_from_image(pseudoimages_src, voxels_src) # [B, N_valid_voxels, C] + feats_voxel_dst = self.extract_voxel_from_image(pseudoimages_dst, voxels_dst) # [B, N_valid_voxels, C] + + feats_voxel_dst_inflate = batched_masked_gather(feats_voxel_dst, knn_idxs_dst.long(), knn_idxs_dst>=0, fill_value=0) + + corr_src_dst = torch.nn.functional.cosine_similarity(feats_voxel_src[:, :, None, :], feats_voxel_dst_inflate, dim=-1) + self.timer[1][3].stop() + + self.timer[1][4].start("Voting") + + voting_vols= self.vote(corr_src_dst[:, :, :, None], voxels_src, voxels_dst, knn_idxs_src, knn_idxs_dst) # [b, l, c, ny, nx] + + vols = self.volconv(voting_vols) + feats_point_vol = self.extract_point_from_voxel(vols, point_voxel_idxs_src) + + self.timer[1][4].stop() + + self.timer[1][5].start("Decoding") + + feats_cat = torch.cat([feats_point_vol, feats_point_src_init, feats_point_src_grid], -1) + flows = self.decoder(feats_cat, point_offsets_src) + self.timer[1][5].stop() + else: + feats_voxel_src = None + feats_voxel_dst = None + corr_src_dst = None + voting_vols = None + + self.timer[1][5].start("Decoding") + feats_cat = torch.cat([feats_point_src_init, feats_point_src_grid], -1) + flows = self.decoder(feats_cat, point_offsets_src) + self.timer[1][5].stop() + + pc0_points_lst = [e["points"] for e in voxel_infos_lst_src] + pc1_points_lst = [e["points"] for e in voxel_infos_lst_dst] + + pc0_valid_point_idxes = [e["point_idxes"] for e in voxel_infos_lst_src] + pc1_valid_point_idxes = [e["point_idxes"] for e in voxel_infos_lst_dst] + + flows_reshape = [] + for (flow, valid_pts) in zip(flows, pc0_valid_point_idxes): + flow = flow[:valid_pts.shape[0], :] + flows_reshape.append(flow) + + model_res = { + "pseudoimages_src": pseudoimages_src, + "pseudoimages_dst": pseudoimages_dst, + "pseudoimages_grid": pseudoimages_grid, + "feats_voxel_src": feats_voxel_src, + "feats_voxel_dst": feats_voxel_dst, + "corr": corr_src_dst, + "voxels_src": voxels_src, + "voting_vol": voting_vols, + "points_src_offset": point_offsets_src, + "points_src_voxel_idx": point_voxel_idxs_src, + "flow": flows_reshape, + "pc0_points_lst": pc0_points_lst, + "pc1_points_lst": pc1_points_lst, + "pc0_valid_point_idxes": pc0_valid_point_idxes, + "pc1_valid_point_idxes": pc1_valid_point_idxes, + } + + return model_res + + + def forward(self, batch): + """ + input: using the batch from dataloader, which is a dict + Detail: [pc0, pc1, pose0, pose1] + output: the predicted flow, pose_flow, and the valid point index of pc0 + """ + self.timer[0].start("Data Preprocess") + batch_sizes = len(batch["pose0"]) + + pose_flows = [] + transform_pc0s = [] + for batch_id in range(batch_sizes): + selected_pc0 = batch["pc0"][batch_id] + self.timer[0][0].start("pose") + with torch.no_grad(): + if 'ego_motion' in batch: + pose_0to1 = batch['ego_motion'][batch_id] + else: + pose_0to1 = cal_pose0to1(batch["pose0"][batch_id], batch["pose1"][batch_id]) + self.timer[0][0].stop() + + self.timer[0][1].start("transform") + # transform selected_pc0 to pc1 + transform_pc0 = selected_pc0 @ pose_0to1[:3, :3].T + pose_0to1[:3, 3] + self.timer[0][1].stop() + pose_flows.append(transform_pc0 - selected_pc0) + transform_pc0s.append(transform_pc0) + + pc0s = torch.stack(transform_pc0s, dim=0) + pc1s = batch["pc1"] + self.timer[0].stop() + + self.timer[1].start("Model Forward") + model_res = self._model_forward(pc0s, pc1s) + self.timer[1].stop() + + ret_dict = model_res + ret_dict["pose_flow"] = pose_flows + + return ret_dict + \ No newline at end of file diff --git a/src/trainer.py b/src/trainer.py index fe2cb16..53954fa 100644 --- a/src/trainer.py +++ b/src/trainer.py @@ -63,6 +63,7 @@ def __init__(self, cfg, eval=False): self.batch_size = int(cfg.batch_size) if 'batch_size' in cfg else 1 self.lr = cfg.lr if 'lr' in cfg else None + self.lr_scheduler = cfg.lr_scheduler if 'lr_scheduler' in cfg else None self.epochs = cfg.epochs if 'epochs' in cfg else None self.metrics = OfficialMetrics() @@ -172,8 +173,14 @@ def train_validation_step_(self, batch, res_dict): pass def configure_optimizers(self): - optimizer = optim.Adam(self.model.parameters(), lr=self.lr) - return optimizer + if self.lr is None: + optimizer = optim.Adam(self.model.parameters(), lr=self.lr) + return optimizer + elif self.lr_scheduler == 'step': + optimizer = optim.Adam(self.model.parameters(), lr=self.lr) + lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=self.trainer.max_epochs//2, gamma=0.1) + return {'optimizer': optimizer, + 'lr_scheduler': lr_scheduler} def on_train_epoch_start(self): self.time_start_train_epoch = time.time()