bs-scheduler: A Batch Size Scheduler library compatible with PyTorch DataLoaders

Abstract

Deep learning models involve computationally intensive training experiments. Increasing the batch size improves the training speed and hardware efficiency by enabling deep neural networks to ingest and process more data in parallel. Inspired by learning rate adaptation policies that yield good results, methods that gradually adjust the batch size have been developed. These methods enhance hardware efficiency without compromising generalization performance. Despite their potential, such methods have not gained widespread popularity or adoption: unlike widely used learning rate policies, for which there is built-in support in most of the deep learning frameworks, the use of batch size adaptation policies requires custom implementations. We introduce an open-source package that implements batch size adaptation policies, which can be seamlessly integrated into deep learning training pipelines. This facilitates more efficient experimentation and accelerates research workflows.