[ViT] Refactor forward function

[yiwen-song]

We would like to simplify the forward function by moving the image reshaping and batch dimension permuting part into another function process_input .

Testing

Run this command:

PYTHONPATH=$PYTHONPATH:`pwd` python -u ~/workspace/scripts/run_with_submitit.py --timeout 3000 --ngpus 1 --nodes 1 --partition train --model vit_b_16 --batch-size 1 --test-only --weights ViT_B_16_Weights.ImageNet1K_V1

Job ID: 14451

Make sure nothing is broken.

@datumbox

cc @datumbox

[facebook-github-bot]

💊 CI failures summary and remediations

As of commit 52fe79c (more details on the Dr. CI page):

---

• 1/1 failures introduced in this PR

---

🕵️ 1 new failure recognized by patterns

The following CI failures do not appear to be due to upstream breakages:

unittest_linux_cpu_py3.7 (1/1)

Step: "Run tests" (full log | diagnosis details | 🔁 rerun)

/root/project/torchvision/io/video.py:406: Runt...log: [mov,mp4,m4a,3gp,3g2,mj2] moov atom not found

test/test_internet.py::TestDatasetUtils::test_download_url_dispatch_download_from_google_drive
  /root/project/env/lib/python3.7/unittest/mock.py:1925: ResourceWarning: unclosed <ssl.SSLSocket [closed] fd=13, family=AddressFamily.AF_INET, type=SocketKind.SOCK_STREAM, proto=6>
    setattr(_type, entry, MagicProxy(entry, self))

test/test_io.py::TestVideo::test_probe_video_from_memory
  /root/project/torchvision/io/_video_opt.py:429: UserWarning: The given buffer is not writable, and PyTorch does not support non-writable tensors. This means you can write to the underlying (supposedly non-writable) buffer using the tensor. You may want to copy the buffer to protect its data or make it writable before converting it to a tensor. This type of warning will be suppressed for the rest of this program. (Triggered internally at  /opt/conda/conda-bld/pytorch_1641456308618/work/torch/csrc/utils/tensor_new.cpp:998.)
    video_data = torch.frombuffer(video_data, dtype=torch.uint8)

test/test_io.py::TestVideo::test_read_video_timestamps_corrupted_file
  /root/project/torchvision/io/video.py:406: RuntimeWarning: Failed to open container for /tmp/tmp_p9bpnlc.mp4; Caught error: [Errno 1094995529] Invalid data found when processing input: '/tmp/tmp_p9bpnlc.mp4'; last error log: [mov,mp4,m4a,3gp,3g2,mj2] moov atom not found
    warnings.warn(msg, RuntimeWarning)

test/test_models.py::test_memory_efficient_densenet[densenet121]
test/test_models.py::test_memory_efficient_densenet[densenet169]
test/test_models.py::test_memory_efficient_densenet[densenet201]
test/test_models.py::test_memory_efficient_densenet[densenet161]
  /root/project/env/lib/python3.7/site-packages/torch/autocast_mode.py:162: UserWarning: User provided device_type of 'cuda', but CUDA is not available. Disabling
    warnings.warn('User provided device_type of \'cuda\', but CUDA is not available. Disabling')

test/test_models.py::test_inception_v3_eval

---

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.

[datumbox]

LGTM, thanks!

[github-actions]

Hey @sallysyw!

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