|
преди 1 месец | |
---|---|---|
.. | ||
README.md | преди 1 месец | |
presets.py | преди 1 месец | |
train.py | преди 1 месец | |
transforms.py | преди 1 месец | |
utils.py | преди 1 месец |
This folder contains reference training scripts for optical flow. They serve as a log of how to train specific models, so as to provide baseline training and evaluation scripts to quickly bootstrap research.
The RAFT large model was trained on Flying Chairs and then on Flying Things. Both used 8 A100 GPUs and a batch size of 2 (so effective batch size is 16). The rest of the hyper-parameters are exactly the same as the original RAFT training recipe from https://github.com/princeton-vl/RAFT. The original recipe trains for 100000 updates (or steps) on each dataset - this corresponds to about 72 and 20 epochs on Chairs and Things respectively:
num_epochs = ceil(num_steps / number_of_steps_per_epoch)
= ceil(num_steps / (num_samples / effective_batch_size))
torchrun --nproc_per_node 8 --nnodes 1 train.py \
--dataset-root $dataset_root \
--name $name_chairs \
--model raft_large \
--train-dataset chairs \
--batch-size 2 \
--lr 0.0004 \
--weight-decay 0.0001 \
--epochs 72 \
--output-dir $chairs_dir
torchrun --nproc_per_node 8 --nnodes 1 train.py \
--dataset-root $dataset_root \
--name $name_things \
--model raft_large \
--train-dataset things \
--batch-size 2 \
--lr 0.000125 \
--weight-decay 0.0001 \
--epochs 20 \
--freeze-batch-norm \
--output-dir $things_dir\
--resume $chairs_dir/$name_chairs.pth
torchrun --nproc_per_node 1 --nnodes 1 train.py --val-dataset sintel --batch-size 1 --dataset-root $dataset_root --model raft_large --weights Raft_Large_Weights.C_T_SKHT_V2
This should give an epe of about 1.3822 on the clean pass and 2.7161 on the
final pass of Sintel-train. Results may vary slightly depending on the batch
size and the number of GPUs. For the most accurate results use 1 GPU and
--batch-size 1
:
Sintel val clean epe: 1.3822 1px: 0.9028 3px: 0.9573 5px: 0.9697 per_image_epe: 1.3822 f1: 4.0248
Sintel val final epe: 2.7161 1px: 0.8528 3px: 0.9204 5px: 0.9392 per_image_epe: 2.7161 f1: 7.5964
You can also evaluate on Kitti train:
torchrun --nproc_per_node 1 --nnodes 1 train.py --val-dataset kitti --batch-size 1 --dataset-root $dataset_root --model raft_large --weights Raft_Large_Weights.C_T_SKHT_V2
Kitti val epe: 4.7968 1px: 0.6388 3px: 0.8197 5px: 0.8661 per_image_epe: 4.5118 f1: 16.0679