Enable to set custom number of steps in benchmarks

[kgajdamo]

• Added --num-steps argument, that allows to set number of the steps (batches two iterate over) in inference and training benchmarks.
• Added wrapper for NeighborLoader, that handles custom number of steps.

[codecov]

Codecov Report

Merging #5898 (d5813c8) into master (6eef677) will not change coverage.
The diff coverage is n/a.

@@           Coverage Diff           @@
##           master    #5898   +/-   ##
=======================================
  Coverage   84.49%   84.49%           
=======================================
  Files         361      361           
  Lines       19840    19840           
=======================================
  Hits        16763    16763           
  Misses       3077     3077           

📣 We’re building smart automated test selection to slash your CI/CD build times. Learn more

[yanbing-j]

Hi @kgajdamo , thanks for your hard work! For benchmark purpose, the function of setting num_steps is very useful, which decrases much time that wasted in running the entire epoch. LGTM!

[kgajdamo]

@rusty1s please do not merge it yet. I noticed that for the same number of steps the version with RandomSampler is x4 faster. While the time should be the same. I haven't had time to look into it yet.

[kgajdamo]

@rusty1s please do not merge it yet. I noticed that for the same number of steps the version with RandomSampler is x4 faster. While the time should be the same. I haven't had time to look into it yet.

[rusty1s]

Ok :)

[kgajdamo]

@rusty1s I checked it and found out what was the reason off such timing. I provided wrong first argument to the RandomSampler:

sampler = torch.utils.data.RandomSampler(
                     range(len(data)), num_samples=args.num_steps *
                     batch_size)

But when I put instead of range(len(data)) - range(len(data.train_mask)) I am receiving an index out of range error. What do You think about it?

[rusty1s]

I think you need to use range(int(data.train_mask.sum())).

[kgajdamo]

I think you need to use range(int(data.train_mask.sum())).

@rusty1s Thank You! Now it's working and You can merge.