Как реализовать планировщик batch_size, используя обратный вызов на керасе

Question

Я хочу создать планировщик batch_size, используя обратный вызов на keras (способом, аналогичным планировщику скорости обучения, который я пытаюсь имитировать в своем коде), но я не могу понять, как это сделать. Вот что я попробовал:

class BatchSizeScheduler(Callback):
      """Batch size scheduler.
      # Arguments
        schedule: a function that takes an epoch index as input
        (integer, indexed from 0) and current batch size
        and returns a new batch size as output (int).
        verbose: int. 0: quiet, 1: update messages.
     """
    def __init__(self, schedule, verbose=1):
         super(BatchSizeScheduler, self).__init__()
         self.schedule = schedule
         self.verbose = verbose

     def on_epoch_begin(self, epoch, logs=None):


         print(self.params)
         bs = self.params['batch_size']

         try:  # new API
            bs = self.schedule(epoch, bs=bs)
         except TypeError:  # old API for backward compatibility
            bs = self.schedule(epoch)
         if not isinstance(bs, (int, np.int32, np.int64)):
            raise ValueError('The output of the "schedule" function '
                         'should be integer.')

         self.params['batch_size'] =  bs 

         if self.verbose > 0:
             print('\nEpoch %05d: BatchSizeScheduler increasing batch '
                 'size to %s.' % (epoch + 1, self.params['batch_size']))

Проблема в том, что это не обновляет размер партии во время обучения. Печать (self.params) показывает увеличивающийся размер пакета, но я могу сказать по индикатору хода обучения, что алгоритм продолжает использовать тот же, исходный размер пакета. Я подозреваю, что способ обновления параметра (self.params['batch_size']=bs) неправильный?

Любая помощь могла бы быть полезна!