Avoid branching in BasicBlock and Bottleneck

[kevinMEH]

Avoid branching during forward pass in BasicBlock and Bottleneck block for better performance.

Completely backwards compatible; downsampled = nn.Identity(x) is equivalent to downsampled = x

Should be always more performant, unless nn.Identity adds some kind of penalty during backprop.

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/vision/8187

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

✅ You can merge normally! (10 Unrelated Failures)

As of commit 38fa921 with merge base ae6b134 ():

FLAKY - The following jobs failed but were likely due to flakiness present on trunk:

• Lint / python-source-and-configs / linux-job (gh)
  RuntimeError: Command docker exec -t 2eadfe873b4ca203a1f7599c2d2aa4cfd4bad3e3c9c87677b3835b8da8d321ed /exec failed with exit code 1
• Lint / python-types / linux-job (gh)
  RuntimeError: Command docker exec -t d20ec2cc00d9a6782e546a382bad56e78e29624995e3ad3f66900e047c8a15f9 /exec failed with exit code 1
• Tests / unittests-linux (3.10, linux.12xlarge, cpu) / linux-job (gh)
  RuntimeError: Command docker exec -t 87a084de88ffd189d9f71eeec087bf54ebc173ef214a75f5b064ce183ee0b686 /exec failed with exit code 1
• Tests / unittests-linux (3.11, linux.12xlarge, cpu) / linux-job (gh)
  RuntimeError: Command docker exec -t 9c0ced1ae25fe917eb4838a361ee05057397ad14165928732f9b9f46498915a8 /exec failed with exit code 1
• Tests / unittests-linux (3.8, linux.12xlarge, cpu) / linux-job (gh)
  RuntimeError: Command docker exec -t 02f859447c7317a951745aff32563c0c7df9587138996d66d4d300368e0f81ec /exec failed with exit code 1
• Tests / unittests-linux (3.9, linux.12xlarge, cpu) / linux-job (gh)
  RuntimeError: Command docker exec -t a51d661031e083186d1389878dc419641c0e1be8cd39072105f07f4b933efb76 /exec failed with exit code 1
• Tests / unittests-macos (3.10, macos-12) / macos-job (gh)
  RuntimeError: Command bash /Users/runner/work/_temp/exec_script failed with exit code 1
• Tests / unittests-macos (3.11, macos-12) / macos-job (gh)
  RuntimeError: Command bash /Users/runner/work/_temp/exec_script failed with exit code 1
• Tests / unittests-macos (3.8, macos-m1-12) / macos-job (gh)
  test/test_transforms_v2.py::TestLinearTransform::test_transform[cpu-dtype1-make_image_tensor]
• Tests / unittests-windows (3.11, windows.4xlarge, cpu) / windows-job (gh)
  test/test_datasets.py::INaturalistTestCase::test_num_examples

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[NicolasHug]

Thanks for the PR @kevinMEH.

I suspect the performance gain coming from removing a single if in the forward pass of resnet would be minimal. I'm interested in looking at some benchmark results, if you have any.

[kevinMEH]

Every single block layer has an if removed. I think that the performance gain is pretty substantial, however as of right now I cannot spare any GPU time to test it. I will test it once I regain access to a GPU.

[NicolasHug]

Thanks, let's wait for benchmarks before moving forward with this then.

We should also remember that this if condition is executed on the CPU, which (hopefully) is running way ahead of the GPU. So, evaluating that if condition is not only crazy cheap but probably doesn't even impact any of the time spent on the GPU - which is what dictates the perf here.