Package: tf-raft
tf-raftRAFT (Recurrent All Pairs Field Transforms for Optical Flow, Teed et. al., ECCV2020) implementation via tf.kerasOriginal resourcesRAFT: Recurrent All Pairs Field Transforms for Optical Flowhttps://github.com/princeton-vl/RAFTInstallation$ pip install tf-raftor you can simply clone this repository.DependenciesTensorFlowTensorFlow-addonsalbumentationssee details inpyoroject.tomlOptical flow datasetsMPI-SintelorFlyingChairsdatasets are relatively light. See more datasets in theoirignal repositoryUsagefromtf_raft.modelimportRAFT,SmallRAFTfromtf_raft.lossesimportsequence_loss,end_point_error# iters/iters_pred are the number of recurrent update of flow in training/predictionraft=RAFT(iters=iters,iters_pred=iters_pred)raft.compile(optimizer=optimizer,clip_norm=clip_norm,loss=sequence_loss,epe=end_point_error)raft.fit(ds_train,epochs=epochs,callbacks=callbacks,steps_per_epoch=train_size//batch_size,validation_data=ds_val,validation_steps=val_size)In practice, you are required to prepare dataset, optimizer, callbacks etc, check details intrain_sintel.pyortrain_chairs.py.Train via YAML configurationtrain_chairs.pyandtrain_sintel.pytrain RAFT model via YAML configuration. Sample configs are inconfigsdirectory. Run;$pythontrain_chairs.py/path/to/config.ymlPre-trained modelsI made the pre-trained weights (on both FlyingChairs and MPI-Sintel) public.
You can download them viagsutilorcurl.Trained weights on FlyingChairs$gsutilcp-rgs://tf-raft-pretrained/2020-09-26T18-38/checkpoints.or$mkdircheckpoints
$curl-OLhttps://storage.googleapis.com/tf-raft-pretrained/2020-09-26T18-38/checkpoints/model.data-00000-of-00001
$curl-OLhttps://storage.googleapis.com/tf-raft-pretrained/2020-09-26T18-38/checkpoints/model.index
$mvmodel*checkpoints/Trained weights on MPI-Sintel (Clean path)$gsutilcp-rgs://tf-raft-pretrained/2020-09-26T08-51/checkpoints.or$mkdircheckpoints
$curl-OLhttps://storage.googleapis.com/tf-raft-pretrained/2020-09-26T08-51/checkpoints/model.data-00000-of-00001
$curl-OLhttps://storage.googleapis.com/tf-raft-pretrained/2020-09-26T08-51/checkpoints/model.index
$mvmodel*checkpoints/Load weightsraft=RAFT(iters=iters,iters_pred=iters_pred)raft.load_weights('checkpoints/model')# forward (with dummy inputs)x1=np.random.uniform(0,255,(1,448,512,3)).astype(np.float32)x2=np.random.uniform(0,255,(1,448,512,3)).astype(np.float32)flow_predictions=model([x1,x2],training=False)print(flow_predictions[-1].shape)# >> (1, 448, 512, 2)NoteThough I have tried to reproduce the original implementation faithfully, there is some difference between the original one and mine (mainly because of used framework: PyTorch/TensorFlow);The original implementations provides cuda-based correlation function but I don't. My TF-based implementation works well, but cuda-based one may run faster.I have trained my model on FlyingChairs and MPI-Sintel separately in my private environment (GCP with P100 accelerator). The model has been trained well, but not reached the best score reported in the paper (trained on multiple datasets).The original one uses mixed-precision. This may get training much faster, but I don't. TensorFlow also enables mixed-precision with few additional lines, seehttps://www.tensorflow.org/guide/mixed_precisionif interested.Additional, global gradient clipping seems to be essential for stable training though it is not emphasized in the original paper. This operation can be done viatorch.nn.utils.clip_grad_norm_(model.parameters(), clip)in PyTorch,tf.clip_by_global_norm(grads, clip_norm)in TF (coded atself.train_stepintf_raft/model.py).Referenceshttps://github.com/princeton-vl/RAFThttps://github.com/NVIDIA/flownet2-pytorchhttps://github.com/NVlabs/PWC-Net