PV-RAFT: Point-Voxel Correlation Fields for Scene Flow Estimation of Point Clouds

Correlation Lookup: The correlation volume $\mathbf{C}$ is built only once and is kept as a lookup table for flow estimations in different steps. Given a source point $p_{1}=(x_{1},y_{1},z_{1})\in P_{1}$, a target point $p_{2}=(x_{2},y_{2},z_{2})\in P_{2}$ and an estimated scene flow $f=(f_{1},f_{2},f_{3})\in\textbf{f}$, the source point is expected to move to $q=(x_{1}+f_{1},x_{2}+f_{2},x_{3}+f_{3})\in Q$, where $Q$ is the translated point cloud. We can easily get the correlation fields between $Q$ and $P_{2}$ by searching the neighbors of $Q$ in $P_{2}$ and looking up the corresponding correlation values in $\mathbf{C}$. Such looking-up procedure avoids extracting features of $Q$ and calculating matrix dot product repeatedly while keeping the all-pairs correlations available at the same time. Since 3D points data is not structured in the dense voxel, grid sampling is no longer useful and we cannot directly convert 2D method [37] into 3D version. Thus, the main challenge is how to locate neighbors and look up correlation values efficiently in the 3D space.

Truncated Correlation: According to our experimental results, not all correlation entries are useful in the subsequent correlation lookup process. The pairs with higher similarity often guide the correct direction of flow estimation, while dissimilar pairs tend to make little contribution. To save memory and increase calculation efficiency in correlation lookup, for each point in $P_{1}$, we select its top-$M$ highest correlations. Specifically, we will get truncated correlation fields $\mathbf{C}_{M}\in\mathbb{R}^{N_{1}\times M}$, where $M<N_{2}$ is the pre-defined truncation number. The point branch and voxel branch are built upon truncated correlation fields.

Point Branch: A common practice to locate neighbors in 3D point clouds is to use K-Nearest Neighbors (KNN) algorithm. Suppose the top-k nearest neighbors of $Q$ in $P_{2}$ is $\mathcal{N}_{k}=\mathcal{N}(Q)_{k}$ and their corresponding correlation values are $\textbf{C}_{M}(\mathcal{N}_{k})$, the correlation feature between $Q$ and $P_{2}$ can be defined as:

where concat stands for concatenation and $\max$ indicates a max pooling operation on $k$ dimension. We briefly note $\mathcal{N}(Q)$ as $\mathcal{N}$ in the following statements as all neighbors are based on $Q$ in this paper. The point branch extracts fine-grained correlation features of the estimated flow since the nearest neighbors are often close to the query point, illustrated in the upper branch of Figure 1. While the point branch is able to capture local correlations, long-range relations are often not taken into account in KNN scenario. Existing methods try to solve this problem by implementing the coarse-to-fine strategy, but error often accumulates if estimates in the coarse stage are not accurate.

Voxel Branch: To tackle the problem mentioned above, we propose a voxel branch to capture long-range correlation features. Instead of voxelizing $Q$ directly, we build voxel neighbor cubes centered around $Q$ and check which points in $P_{2}$ lie in these cubes. Moreover, we also need to know each point’s relative direction to $Q$. Therefore, if we denote sub-cube side length by $r$ and cube resolution by $a$, then the neighbor cube of $Q$ would be a $a\times a\times a$ Rubik’s cube:

where $\mathbf{i}=[i,j,k]^{T},\lceil{-\frac{a}{2}}\rceil\leq i,j,k\leq\lceil{\frac{a}{2}}\rceil\in\mathbb{Z}$ and each $r\times r\times r$ sub-cube $\mathcal{N}_{r}^{(\mathbf{i})}$ indicates a specific direction of neighbor points (\eg, $[0,0,0]^{T}$ indicates the central sub-cube). Then we identify all neighbor points in the sub-cube $\mathcal{N}_{r}^{(\mathbf{i})}$ and average their correlation values to get sub-cube features. The correlation feature between $Q$ and $P_{2}$ can be defined as:

where $n_{\mathbf{i}}$ is the number of points in $P_{2}$ that lie in the $\mathbf{i}^{th}$ sub-cube of $Q$ and $\textbf{C}_{v}(Q,P_{2})\in\mathbb{R}^{N_{1}\times a^{3}}$. Please refer to the lower branch of Figure 1 for illustration.

Feature Extraction: The feature extractor $E_{\theta}$ encodes point clouds with mere coordinates information into higher dimensional feature space, as $E_{\theta}:\mathbb{R}^{n\times 3}\mapsto\mathbb{R}^{n\times D}$. Our backbone framework is based on PointNet++ [29]. For consecutive point clouds input $P_{1},P_{2}$, the feature extractor outputs $E_{\theta}(P_{1}),E_{\theta}(P_{2})$ as backbone features. Besides, we design a content feature extractor $E_{\gamma}$ to encode context feature of $P_{1}$. Its structure is exactly the same as feature extractor $E_{\theta}$, without weight sharing. The output context feature $E_{\gamma}(P_{1})$ is used as auxiliary context information in GRU iteration.

Correlation Fields Construction: As is introduced in Section 3.1, we build all-pair correlation fields $\mathbf{C}$ based on backbone features $E_{\theta}(P_{1}),E_{\theta}(P_{2})$. Then we truncate it according to correlation value sorting and keep it as a lookup table for later iterative updates.

Iterative Flow Estimation: The iterative flow estimation begins with the initialize state $\mathbf{f}_{0}=0$. With each iteration, the scene flow estimation is updated upon the current state: $\mathbf{f}_{t+1}=\mathbf{f}_{t}+\Delta\mathbf{f}$. Eventually, the sequence converges to the final prediction $\mathbf{f_{T}}\to\textbf{f}^{*}$. Each iteration takes the following variables as input: (1) correlation features, (2) current flow estimate, (3) hidden states from the previous iteration, (4) context features. First, the correlation features are the combination of both fine-grained point-based ones and long-range pyramid-voxel-based ones:

Second, the current flow estimation is simply the direction vector between $Q_{t}$ and $P_{1}$:

Third, the hidden state $h_{t}$ is calculated by GRU cell[37]:

where $x_{t}$ is a concatenation of correlation $\mathbf{C}_{t}$, current flow $\mathbf{f}_{t}$ and context features $E_{\gamma}(P_{1})$. Finally, the hidden state $h_{t}$ is fed into a small convolutional network to get the final scene flow estimate $\textbf{f}^{*}$. The detailed iterative update process is illustrated in Figure 3.

Flow Refinement: The purpose of designing this flow refinement module is to make scene flow prediction $\textbf{f}^{*}$ smoother in the 3D space. Specifically, the estimated scene flow from previous stages is fed into three convolutional layers and one fully connected layer. To update flow for more iterations without out of memory, the refinement module is not trained end-to-end with other modules. We first train the backbone and iterative update module, then we freeze the weights and train the refinement module alone.

Flow Supervision: We follow the common practice of supervised scene flow learning to design our loss function. In detail, we use $l_{1}$-norm between the ground truth flow $\mathbf{f}_{gt}$ and estimated flow $\mathbf{f}_{est}$ for each iteration:

where $T$ is the total amount of iterative updates, $\mathbf{f}_{est}^{(t)}$ is the flow estimate at $t^{th}$ iteration, and $w^{(t)}$ is the weight of $t^{th}$ iteration:

where $\gamma$ is a hyper-parameter and we set $\gamma=0.8$ in our experiments.

Refinement Supervision: When we freeze the weights of previous stages and only train the refinement module, we design a similar refinement loss:

where $\mathbf{f}_{ref}$ is the flow prediction from refinement module.

Datasets: Same with [6, 18, 50, 26], we trained our model on the FlyingThings3D [20] dataset and evaluated it on both FlyingThings3D [20] and KITTI [21, 22] datasets. We followed [6] to preprocess data. As a large-scale synthetic dataset, FlyingThings3D is the first benchmark for scene flow estimation. With the objects from ShapeNet [3], FlyingThings3D consists of rendered stereo and RGB-D images. Totally, there are 19,640 pairs of samples in the training set and 3,824 pairs in the test set. Besides, we kept aside 2000 samples from the training set for validation. We lifted depth images to point clouds and optical flow to scene flow instead of operating on RGB images. As another benchmark, KITTI Scene Flow 2015 is a dataset for scene flow estimation in real scans [21, 22]. It is built from KITTI raw data by annotating dynamic motions. Following previous works [6, 18, 50, 26], we evaluated on 142 samples in the training set since point clouds were not available in the test set. Ground points were removed by height (0.3m). Further, we deleted points whose depths are larger than 35m.

Implementation Details: We randomly sampled 8192 points in each point cloud to train PV-RAFT. For the point branch, we searched 32 nearest neighbors. For the voxel branch, we set cube resolution $a=3$ and built 3-level pyramid with $r=0.25,0.5,1$. To save memory, we set truncation number $M$ as 512. We updated scene flow for 8 iterations during training and evaluated the model with 32 flow updates. The backbone and iterative module were trained for 20 epochs. Then, we fixed their weights with 32 iterations and trained the refinement module for another 10 epochs. PV-RAFT was implemented in PyTorch [24]. We utilized Adam optimizer [13] with initial learning rate as 0.001 .

Evaluation Metrics: We adopted four evaluation metrics used in [6, 18, 50, 26], including EPE, Acc Strict, Acc Relax and Outliers. We denote estimated scene flow and ground-truth scene flow as $f_{est}$ and $f_{gt}$ respectively. The evaluation metrics are defined as follows:

$\bullet$ EPE: $||f_{est}-f_{gt}||_{2}$. The end point error averaged on each point in meters.

$\bullet$ Acc Strict: the percentage of points whose EPE $<0.05m$ or relative error $<5\%$.

$\bullet$ Acc Relax: the percentage of points whose EPE $<0.1m$ or relative error $<10\%$.

$\bullet$ Outliers: the percentage of points whose EPE $>0.3m$ or relative error $>10\%$.

Effects of Truncation Operation: We introduce the truncation operation to reduce running memory while maintain the performance. To prove this viewpoint, we conducted experiments with different truncation numbers $M$, which are shown in Table 3. On the one hand, when $M$ is too small, the accuracy will degrade due to the lack of correlation information. On the other hand, achieving the comparable performance with $M=512$, the model adopting $M=1024$ needs about 14G running memory, which is not available on many GPU services (\egRTX 2080 Ti). This result indicates that top 512 correlations are enough to accurately estimate scene flow with high efficiency.

Comparison with Other Correlation Volumes: To further demonstrate the superiority of the proposed point-voxel correlation fields, we did comparison with correlation volume methods introduced in FlowNet3D [18] and PointPWC-Net [50]. To fairly compare, we applied their correlation volumes in our framework to substitute point-voxel correlation fields. Evaluation results are shown in Table 4. Leveraging all-pairs relations, our point-voxel correlation module outperforms other correlation volume methods.