Efficiently compute total number of steps

[RicardoDominguez]

When calculating the total number of steps, the current implementation required iterating through the entire dataset, which can be very slow for large datasets. Instead, take sampler.num_batches(), which I believe should correspond precisely to the length of the data loader (since when packing, the batch size of the data loader is 1).

[winglian]

@RicardoDominguez thanks! looks good so far. I need to just do some sanity checks on the validity of this logic. Do you have any ideas about how we might be able to unit test this?

[winglian]

I went ahead and rebased the PR onto main as there were some linting issues in the commit of main you were on.

[RicardoDominguez]

Hmm not sure how to unit test these changes beyond comparing sampler.num_batches() to len(dataloader).

[winglian]

I also think there is a minor discrepancy because we often pass drop_last=True to the dataloader.

[winglian]

I think this might not be necessary as checking the length of the sampler is already what happens under the hood.

Also, inspecting the torch DataLoader class, the __len__() method already uses return len(self._index_sampler), where

    @property
    def _index_sampler(self):
        # The actual sampler used for generating indices for `_DatasetFetcher`
        # (see _utils/fetch.py) to read data at each time. This would be
        # `.batch_sampler` if in auto-collation mode, and `.sampler` otherwise.
        # We can't change `.sampler` and `.batch_sampler` attributes for BC
        # reasons.
        if self._auto_collation:
            return self.batch_sampler
        else:
            return self.sampler