[Templates] Unify the batch inference template with an existing Data example

[justinvyu]

This PR de-duplicates the batch inference template by making it the same as the existing pytorch gpu batch inference example. There still needs to be a copy due to relative references in the docs not generating correctly when pulling the notebook code directly.

This PR also fixes some typos in the Data example and changes some code to have no warnings show up when running through the example (increasing the model + dataset size for a reasonable batch size with 4 workers + using a kwarg when initializing the resnet model with weights).

Notes

GPU utilization after (no warnings about reducing the batch size):

Related issue number

Checks

• I've signed off every commit(by using the -s flag, i.e., git commit -s) in this PR.
• I've run scripts/format.sh to lint the changes in this PR.
• I've included any doc changes needed for https://docs.ray.io/en/master/.
  • I've added any new APIs to the API Reference. For example, if I added a
    method in Tune, I've added it in doc/source/tune/api/ under the
    corresponding .rst file.
• I've made sure the tests are passing. Note that there might be a few flaky tests, see the recent failures at https://flakey-tests.ray.io/
• Testing Strategy
  • Unit tests
  • Release tests
  • This PR is not tested :(

[justinvyu]

Template running as release tests: https://buildkite.com/ray-project/release-tests-pr/builds/42209

[amogkam]

Overall lgtm, but can we remove the explicit materialize call? It’s not necessary and prevents from streaming to the writes

[justinvyu]

@amogkam Done. I originally had a materialize because running take_batch and write_parquet would both seem to run the full dataset execution. However, this is not actually the case since take_batch actually only runs the prediction on a small amount of data, and write_parquet finishes the execution on the rest of the data.

So, the predictions are only computed 1x, rather than 2x as I originally thought. Is that correct?

[amogkam]

Yep that’s right it will only run 1x, barring a few extra samples