SMURF: Self-Teaching Multi-Frame Unsupervised RAFT with Full-Image Warping

Introduction

Optical flow describes a dense pixel-wise correspondence between two images, specifying for each pixel in the first image, where that pixel is in the second image. The resulting vector field of relative pixel locations represents apparent motion or “flow” between the two images. Estimating this flow field is a fundamental problem in computer vision and any advances in flow estimation benefit many downstream tasks such as visual odometry, multi-view depth estimation, and video object tracking.

Classical methods formulate optical flow estimation as an optimization problem . They infer, for a given image pair, a flow field that maximizes smoothness and the similarity of matched pixels. Recent supervised learning approaches instead train deep neural networks to estimate optical flow from examples of ground-truth annotated image pairs . Since obtaining ground truth flow is extremely difficult for real images, supervised learning is generally limited to synthetic data . While these methods have produced excellent results in the training domain, generalization is challenging when the gap between the target domain ant the synthetic training data is too wide.

Unsupervised learning is a promising direction to address this issue as it allows training optical flow models from unlabeled videos of any domain. The unsupervised approach works by combining ideas from classical methods and supervised-learning – training the same neural networks as in supervised approaches but optimizing them with objectives such as smoothness and photometric similarity from classical methods. Unlike those classical methods, unsupervised approaches perform optimization not per image pair but jointly for the entire training set.

Since unsupervised optical flow takes inspiration from classical and supervised learning methods, we can make substantial progress by properly combining novel ideas with insights from these two directions. In this paper, we do exactly that and make the following three contributions:

We integrate the current best supervised model, RAFT with unsupervised learning and perform key changes to the loss functions and data augmentation to properly regularize this model for unsupervised learning.

We perform unsupervised learning on image crops while using the full image to compute unsupervised losses. This technique, which we refer to as full-image warping, improves flow quality near image boundaries.

We leverage a classical method for multi-frame flow refinement to generate better labels for self-supervision from multi-frame input. This technique improves performance especially in occluded regions without requiring more than two frames for inference.

Our method Self-Teaching Multi-frame Unsupervised RAFT with Full-Image Warping (SMURF) combines these three contributions and improves the state of the art (SOTA) in unsupervised optical flow in all major benchmarks, i.e. it reduces errors by 40 / 36 / 39 % in the Sintel Clean / Sintel Final / KITTI 2015 benchmarks relative to the prior SOTA set by UFlow . These improvements also reduce the gap to supervised approaches, as SMURF is the first unsupervised optical flow method that outperforms supervised FlowNet2 and PWC-Net on all benchmarks.

Related Work

Optical flow was first studied in psychology to describe motion in the visual field . Later, the rise of computing in the 1980s led to analytical techniques to estimate optical flow from images . These techniques introduced photometric consistency and smoothness assumptions. These early methods do not perform any learning but instead solve a system of equations to find flow vectors that minimize the objective function for a given image pair. Follow-up work continued to improve flow accuracy, e.g. through better optimization techniques and occlusion reasoning .

Machine learning has helped to improve results substantially. First approaches used supervised convolutional neural networks that had relatively little flow-specific structure . Others introduced additional inductive biases from classical approaches such as coarse-to-fine search . The current best network architecture is the Recurrent All-Pairs Field Transforms (RAFT) model. It follows classical work that breaks with the coarse-to-fine assumption and computes the cost volume between all pairs of pixels and uses that information to iteratively refine the flow field . All supervised methods rely heavily on synthetic labeled data for training. While producing excellent results in the supervised setting, RAFT has not previously been used in unsupervised learning.

Unsupervised approaches appeared after supervised methods and showed that even without labels, deep learning can greatly outperform classical flow methods . Besides being more accurate than classical methods, learned methods are also faster at inference because all optimization occurs during training instead of during inference time.The reason why learned unsupervised methods outperform classical methods is not obvious as both approaches optimize similar objectives. We speculate that the network structure introduces helpful inductive biases and that optimizing a single network for all training images provides a better regularization than independently optimizing a flow field for each image.

A recent study performed an extensive comparison of the many proposed advances in unsupervised flow estimation and amalgamated these different works into a state of the art method called UFlow. We take this work as a starting point and build on the techniques suggested there such as range-map based occlusion reasoning , the Census loss for photometric consistency , edge aware smoothness , and self supervision .

We build on and substantially extend this prior work by enabling the RAFT model to work in an unsupervised learning setting through changes in the loss function and data augmentation. We also propose full-image warping to make the photometric loss useful for pixels that leave the (cropped) image plane. And we utilize a flow refinement technique to leverage multi-frame input during training to self-generate improved labels for self-supervision.

SMURF for Unsupervised Optical Flow

This section describes our method Self-Teaching Multi-frame Unsupervised RAFT with Full-Image Warping (SMURF) in two parts. The first part covers preliminaries, such as the problem definition and components adopted from prior work. The second part describes the three main improvements in our method.

We build on the loss function $\mathcal{L}_{\mathit{}}$ from UFlow ,

which is a weighted combination of three terms: occlusion aware photometric consistency $\mathcal{L}_{\mathit{photo}}$, edge-aware smoothness $\mathcal{L}_{\mathit{smooth}}$, and self-supervision $\mathcal{L}_{\mathit{self}}$.

The photometric consistency term is defined as

The $k$-th order edge-aware smoothness term is defined as $\displaystyle\mathcal{L}_{\mathit{smooth(k)}}(D,\theta)$ \displaystyle=\frac{1}{HW}\sum\Bigg{(}\exp{\left(-\frac{\lambda}{3}\sum_{c}\left\lvert\frac{\partial I_{1_{c}}}{\partial x}\right\rvert\right)}\odot\left\lvert\frac{\partial^{k}V_{1}}{\partial x^{k}}\right\rvert \displaystyle+\exp{\left(-\frac{\lambda}{3}\sum_{c}\left\lvert\frac{\partial I_{1_{c}}}{\partial y}\right\rvert\right)}\odot\left\lvert\frac{\partial^{k}V_{1}}{\partial y^{k}}\right\rvert\Bigg{)}\,. The expression $\frac{1}{HW}\sum$ again computes the mean over all pixels, here for two components that compute the $k$-th derivative of the flow field $V_{1}$ along the $x$ and $y$ axes. These derivatives are each weighted by an exponential of the mean image derivative across color channels $c$. This weighting provides edge-awareness by ignoring non-smooth transitions in the flow field at visual edges . The sensitivity to visual edges is controlled by $\lambda$.

The self-supervision term is defined as \displaystyle\mathcal{L}_{\mathit{self}}(D,\theta)=\frac{1}{HW}\sum\hat{M}_{1}\odot\Big{(}1-M_{1}\Big{)}\odot c\Big{(}\hat{V}_{1},V_{1}\Big{)}\,. where a self-generated flow label $\hat{V}_{1}$ is compared to the flow prediction $V_{1}$ using the generalized Charbonnier function $c(A,B)=\left((A-B)^{2}+\epsilon^{2}\right)^{\alpha}$ with all operations being applied elementwise. Our experiments use $\epsilon=0.001$ and $\alpha=0.5$. Prior work has modulated this comparison by masks $\hat{M}$ and $M$ for the generated labels and the predicted flow respectively, which are computed with a forward-backward consistency check , so that self-supervision is only applied to image regions where the label is forward-backward consistent but the prediction is not. However, as we will describe below, our method actually improves when we remove this masking.

The main reason to perform self-supervision is to improve flow estimation near the image boundaries where pixels can easily get occluded by moving outside of the image frame . To address this problem via self-supervision, we apply the model on an image pair to produce a flow label and use this self-generated label to supervise flow predictions from a cropped version of the image pair. This way the model can transfer its own predictions from an easier to a more difficult setting. We build on and extend this technique with additional augmentations as described below.

2 SMURF’s Improvements

After having covered the foundation that our method builds on, we will now explain our three major improvements: 1) enabling the RAFT architecture to work with unsupervised learning, 2) performing full-image warping while training on image crops, and 3) introducing a new method for multi-frame self-supervision.

In supervised learning of optical flow, it is common to apply losses not only to the final output but also to intermediate flow predictions . In unsupervised learning this is not typically done – presumably because intermediate outputs of other models are often at lower resolutions which might not work well with photometric and other unsupervised losses . However, RAFT produces intermediate predictions at full resolution and does not pass gradients between prediction steps. Thus, we apply the unsupervised losses at the entire sequence of RAFT’s intermediate predictions, which we found to be essential to prevent divergence and to achieve good results. Similarly to supervised methods, we exponentially decay the weight of these losses at earlier iterations . $\mathcal{L}_{\mathit{sequence}}=\sum_{i=1}^{n}\gamma^{n-i}\mathcal{L}_{\mathit{i}}$, where $n$ is the number of flow iterations, $\gamma$ is the decay factor, and $\mathcal{L}_{\mathit{i}}$ is the loss at iteration $i$. Our experiments use $n=12$ and $\gamma=0.8$.

For self-supervision, we implement the sequence loss by applying the model to a full image, taking its final output, and using that to supervise all iterations of the model applied to the cropped image (see Figure 1). Applying self-supervision from full to cropped images has been shown to provide a learning signal to pixels near the image border that move out of the image and therefore receive no signal from photometric losses . Performing this self-supervision in a sequence-aware manner allows the earlier stages of the model to learn from the model’s more refined final output. Potentially because of this inherent quality difference between the “student” and the “teacher” flow, we found that simplifying the self-supervision loss by removing masks $M$ and $\hat{M}$ improves performance further. The self-supervision loss without masks is \mathcal{L}_{\mathit{self}}(D,\theta)=\frac{1}{HW}\sum c\Big{(}\hat{V}_{1},V_{1}\Big{)}. Note that this change does not affect the photometric loss $\mathcal{L}_{\mathit{photo}}$, where we continue to use $O$ to mask out occlusions.

To regularize RAFT, we use the same augmentation as supervised RAFT which is much stronger than what has typically been used in unsupervised optical flow, except for the recent ARFlow . We randomly vary hue, brightness, saturation, stretching, scaling, random cropping, random flipping left/right and up/down, and we apply a random eraser augmentation that removes random parts of each image. All augmentations are applied to the model inputs, but not to the images used to compute the photometric and smoothness losses. The self-generated labels for self-supervision are computed from un-augmented images, which has the benefit of training the model to ignore these augmentations (see Figure 1).

2.2 Full-Image Warping

The photometric loss, which is essential for unsupervised optical flow estimation, is generally limited to flow vectors that stay inside the image frame because vectors that point outside of the frame have no pixels to compare their photometric appearance to. We address this limitation by computing the flow field from a cropped version of the images $I_{1}$ and $I_{2}$ while referencing the full, uncropped image $I_{2}$ when warping it with the estimated flow $V_{1}$ before computing the photometric loss (see Figure 2). As we also no longer mark these flow vectors that move outside the image frame as occluded, they now provide the model with a learning signal. We use full-image warping for all datasets except Flying Chairs, where we found that cropping the already small images hurt performance.

2.3 Multi-Frame Self-Supervision

Finally, we propose to leverage multi-frame information for self-supervision to generate better labels in occluded areas, inspired by work that used a similar technique for inference . For multi-frame self-supervision, we take a frame $t$ and compute the forward flow to the next frame ($t\rightarrow t+1$) and the backward flow to the previous frame ($t\rightarrow t-1$). We then use the backward flow to predict the forward flow through a tiny learned inversion model and use that prediction to inpaint areas that were occluded in the original forward flow but were not occluded in the backward flow – which is why the estimate from the backward flow is more accurate (see Figure 3). The tiny model for the backward-forward inversion consists of three layers of 3$\times$3 convolutions with channels that are applied on the backward flow and the image coordinates normalized to $$. The model is re-initialized and trained per frame using the non-occluded forward flow as supervision, after which the in-painted flow field is stored and used for self-supervision. We apply multi-frame self-supervision only at the final stage of training. Importantly, we use multiple frames only during training and not for inference.

Experiments

Our hyperparameters are based on UFlow with slight modifications based on further hyperparameter search. In all experiments, we use weights $\omega_{\mathit{photo}}=1$ and $\omega_{\mathit{self}}=0.3$. Regarding smoothness parameters, we set edge sensitivity $\lambda=150$ and $\omega_{\mathit{smooth}}=4$ for KITTI and Chairs and $\omega_{\mathit{smooth}}=2.5$ for Sintel. We use 2nd order smoothness ($k=2$) for KITTI and 1st order smoothness ($k=1$) for all other datasets. We train for 75K iterations with batch size 8, except for Sintel where we train for only 15K iterations to avoid overfitting. The self-supervision weight is set to 0 for the first 40% of gradient steps, then linearly increased to $0.3$ during the next 10% of steps and then kept constant. All of our training uses Adam ($\beta_{1}=0.9$, $\beta_{2}=0.999$, $\epsilon=10^{-8}$) with learning rate of $0.0002$, which takes about 1 day to converge on 8 GPUs running synchronous SGD. On all datasets, the learning rate is exponentially decayed to $\frac{1}{1000}$ of its original value over the last 20% of steps. On Sintel and on Flying Chairs we train with random crops of $368\times 496$, and on KITTI we train with random crops of size $296\times 696$. For all datasets we evaluated our model on the input resolution that had the lowest average end point error on the training set; $488\times 1144$ for KITTI and $480\times 928$ for Sintel. All test images were bilinearly resized to these resolutions during inference, and the resulting flow field was bilinearly resized and rescaled back to the native image size to compute evaluation metrics.

During the second stage of training when multi-frame self-supervision is applied, we generate labels for all images using the model trained according to the procedure described above. We then continue training the same model with only the self-supervision loss for an additional 30K iterations using the multiframe generated labels. We use the same hyperparameters as in the first stage but exponentially decay the learning rate for the last 5K iterations. For our best performing Sintel model, we train with the KITTI self-supervision labels mixed in at a ratio of 50%.

We train and evaluate our model according to the conventions in the literature using the following optical flow datasets: Flying Chairs , Sintel and KITTI 2015 . We pretrain on Flying Chairs before fine tuning on Sintel or KITTI. Similar to UFlow , we did not find a benefit to pretraining on more out-of-domain data, e.g. Flying Things . None of the training techniques in our method uses any ground truth labels. We train on the “training” portion of Flying Chairs, and divide the Sintel dataset according to its standard train / test split. For KITTI, we train on the multi-view extension following the split used in prior work : We train two models, one on the multi-view extension of the training set and one on the extension of the test set, evaluate these models appropriately. For ablations, we report metrics after training on the “test” portion of the dataset (which does not include labels) and evaluating on the training set, and for final benchmark numbers we report results after training on the training portion only. For our benchmark result on Sintel, we train on a 50% mixture of the KITTI and Sintel multi-frame self-supervision labels. For all datasets we report endpoint error (“EPE”), and for KITTI, we additionally report error rates (“ER”), where a prediction is considered erroneous if its EPE is $>3$ pixels or $>5\%$ of the length of the true flow vector. We generally compute these metrics for all pixels, except for “EPE (noc)” where only non-occluded pixels are considered.

Results

In this section, we compare SMURF to related methods, ablate the proposed improvements, and show limitations.

Qualitative results for SMURF are shown in Figure 4 and a comparison to other methods can be found in Table 1. This comparison shows that our model substantially outperforms all prior published methods on unsupervised optical flow on all benchmarks. Compared to the the previous state of the art method UFlow , our method reduces benchmark test errors by 40 / 36 / 39 % for Sintel Clean / Sintel Final / KITTI 2015. When we compare SMURF trained in domain to supervised methods trained out of domain, SMURF outperforms all supervised methods except RAFT on all benchmarksThis is a fair comparison as obtaining labels for a given domain is extremely difficult while our method trains on readily available video data.. Compared to supervised RAFT, it performs a bit worse on Sintel (1.99 vs 1.43 on Sintel Clean, 2.80 vs. 2.71 on Sintel Final) but much better on KITTI 2015 (2.01 vs. 5.04 EPE, 6.72% vs. 17.4% ER). SMURF even outperforms some supervised methods when they are finetuned on the test domain, e.g. FlowNet2-ft , PWC-Net-ft , and SelFlow-ft . Only VCN-ft and RAFT-ft achieve better performance here. The inference time of our model is the same as the supervised RAFT model, approximately 500ms when using 12 recurrent iterations. A qualitative comparison to supervised RAFT is shown in Figure 5.

In Table 2, we compare generalization across domains to prior unsupervised methods and find substantial improvements for every combination of training and test domain. Even after only training on Flying Chairs, our model already achieves an EPE of 2.19 on Sinel Clean and 3.35 on Sinel Final, which is superior to all prior unsupervised techniques even if fine tuned on Sintel data (first and third column in Table 1).

We also tested generalization from optical flow to stereo depth estimation. Here we evaluated our exact trained KITTI-15 flow model in the KITTI 2015 stereo depth benchmark . Despite neither tailoring the architecture or losses to stereo depth estimation nor training or fine tuning on any stereo data, our method achieves results that are competitive with the best unsupervised stereo depth estimation methods (see Table 3).

2 Ablation Study

To determine which aspects of our model are responsible for its improved performance over prior work, we perform an extensive ablation study. In these ablations, we always train one model per domain (on KITTI-2015-test and Sintel-test after pretraining on Flying Chairs), and evaluate those on the corresponding validation split of the same domain.

Our first ablation investigates how much improvement can be obtained by taking the prior state of the art method UFlow and replacing its PWC model with RAFT. As the results in Table 4 show, replacing the model without additional changes to the unsupervised learning method surprisingly does not improve but instead decreases performance. Through extensive experimentation, we identified and added the techniques presented here that enable superior unsupervised learning with RAFT. The gains from these techniques are much smaller with the PWC model, potentially because of the more constrained architecture.

Next, we test different combinations of the novel components in our method. The results in Table 5 show that every component has a significant impact on performance. Sequence losses and heavy augmentations are necessary to achieve good results, especially on Sintel. And when using heavy augmentations, we need to apply sequence losses to prevent divergence of the model. Full-Image warping and multi-frame self-supervision have the strongest effect on KITTI, which makes sense as larger camera motion in that dataset causes more occlusions at the image boundaries which these components help to address.

We also ablated our proposed changes to self-supervision (subsubsection 3.2.1). The ablations in Table 6 show that even with full-image warping, self-supervision remains an important component, and that it works best when not masking the loss as in prior work and when we use the final output (not intermediate model predictions) to generate the self-supervision labels.

Lastly, we provide a qualitative ablation of multi-frame self-supervision (subsubsection 3.2.3) and a comparison to other unsupervised methods that use multi-frame information . Figure 6 shows that multi-frame self-supervision substantially improves flow accuracy in occluded areas and does this much better than related multi-frame methods while being the only approach that requires multiple frames only during training and not for inference.

3 Limitations

A major limitation of unsupervised optical flow is that it estimates apparent visual motion rather than motion of physical objects (see Figure 7). Overcoming this limitation requires some form of supervision, reasoning about the 3-D space as in scene flow , reasoning about semantics, or a combination of these. Future work could try to transfer the techniques from our method to such approaches.

Conclusion

We have presented SMURF, an effective method for unsupervised learning of optical flow that reduces the gap to supervised approaches and shows excellent generalization across datasets and even to “zero-shot” depth estimation. SMURF brings key improvements, most importantly (1) enabling the RAFT architecture to work in an unsupervised setting via modifications to the unsupervised losses and data augmentation, (2) full-image warping for learning to predict out of frame motion, and (3) multi-frame self-supervision for improved flow estimates in occluded regions. We believe that these contributions are a step towards making unsupervised optical flow truly practical, so that optical flow models trained on unlabeled videos can provide high quality pixel-matching in domains without labeled data.

References