| 256 | |
| 257 | class GradientAccumulationScheduler: |
| 258 | def __init__(self, scheduling: dict): |
| 259 | if scheduling == {}: # empty dict error |
| 260 | raise TypeError("Empty dict cannot be interpreted correct") |
| 261 | |
| 262 | for key in scheduling.keys(): |
| 263 | if not isinstance(key, int) or not isinstance(scheduling[key], int): |
| 264 | raise TypeError("All epoches and accumulation factor must be integers") |
| 265 | |
| 266 | minimal_epoch = min(scheduling.keys()) |
| 267 | if minimal_epoch < 1: |
| 268 | msg = f"Epochs indexing from 1, epoch {minimal_epoch} cannot be interpreted correct" |
| 269 | raise IndexError(msg) |
| 270 | elif minimal_epoch != 1: # if user didnt define first epoch accumulation factor |
| 271 | scheduling.update({1: 1}) |
| 272 | |
| 273 | self.scheduling = scheduling |
| 274 | self.epochs = sorted(scheduling.keys()) |
| 275 | |
| 276 | def on_epoch_begin(self, epoch, trainer): |
| 277 | epoch += 1 # indexing epochs from 1 |