RAFT training reference Improvement

[YosuaMichael]

Do some of the task on: #5056

• Change function name from validate to evaluate
• Support --device and enable training on non-distributed mode
• Include optimizer and scheduler in the checkpoint

Sample script to run on non-distributed mode and on cpu:

 python train.py \
    --dataset-root $dataset_root \
    --name $name_chairs \
    --model raft_small \
    --train-dataset chairs \
    --batch-size 2 \
    --lr 0.0004 \
    --weight-decay 0.0001 \
    --epochs 2 \
    --output-dir $out_chairs \
    --device cpu

To test on CPU, I run on a mock dataset by replacing https://github.com/pytorch/vision/blob/main/torchvision/datasets/_optical_flow.py with https://gist.github.com/YosuaMichael/9c49729243ff9d467ece06ab8641680d.

Note that as of now, if we run on distributed mode using torchrun, then it must use --device cuda.

[facebook-github-bot]

💊 CI failures summary and remediations

As of commit 0e7ab27 (more details on the Dr. CI page):

---

💚 💚 Looks good so far! There are no failures yet. 💚 💚

---

This comment was automatically generated by Dr. CI (expand for details).

Please report bugs/suggestions to the (internal) Dr. CI Users group.

Click here to manually regenerate this comment.

[NicolasHug]

Thank you for the PR @YosuaMichael. There are 2 minor issues (see below), but otherwise this looks great!

[YosuaMichael]

Update:

Support saving of the optimizer and scheduler on the checkpoint.

[YosuaMichael]

Hi @NicolasHug , I decided to put the commit for saving optimizer and scheduler in this PR as well: 09d78d1
Could you also help to review this? Thanks!

[NicolasHug]

Thanks @YosuaMichael , we're almost there :) . I made a few comments below

[NicolasHug]

Thanks @YosuaMichael, nice work ! There was a minor issue left, which I fixed in 2857e21: when no trainset is specified we want to directly go to evaluate, without worrying about train_dataset - the previous code would fail because it's None.

[github-actions]

Hey @NicolasHug!

You merged this PR, but no labels were added. The list of valid labels is available at https://github.com/pytorch/vision/blob/main/.github/process_commit.py