74 lines
3.2 KiB
Python
74 lines
3.2 KiB
Python
import json
|
|
import logging
|
|
import math
|
|
import os
|
|
import torch
|
|
|
|
from plugins.plugins import BasePlugin
|
|
|
|
class Accumulnator(BasePlugin):
|
|
|
|
def __init__(self):
|
|
path = os.path.join(os.path.dirname(__file__), "accumulnator.json")
|
|
logging.info(f" * Accumulnator plugin instantiated, loading config from {path}")
|
|
with open(path, 'rt') as f:
|
|
config = json.load(f)
|
|
begin_epoch = config['begin_epoch']
|
|
begin_grad_accum = config['begin_grad_accum']
|
|
end_epoch = config['end_epoch']
|
|
end_grad_accum = config['end_grad_accum']
|
|
|
|
# spread the grad accums
|
|
curve = config['curve']
|
|
steps = end_epoch - begin_epoch
|
|
if curve == 'linear':
|
|
accums = torch.linspace(start=begin_grad_accum,
|
|
end=end_grad_accum,
|
|
steps=end_epoch-begin_epoch).tolist()
|
|
elif curve == 'log':
|
|
accums = torch.logspace(start=math.log(begin_grad_accum, 2),
|
|
end=math.log(end_grad_accum, 2),
|
|
base=2,
|
|
steps=steps).tolist()
|
|
else:
|
|
raise NotImplementedError(f"curve not {curve} not recognized")
|
|
#print(f"accums: {accums}")
|
|
accums_per_epoch = {}
|
|
for i in range(begin_epoch):
|
|
accums_per_epoch[i] = begin_grad_accum
|
|
for i in range(steps):
|
|
#print(f"took accum {accums[i]} for epoch {i+begin_epoch}")
|
|
accums_per_epoch[i+begin_epoch] = round(accums[i])
|
|
|
|
logging.info(f" * Accumulnator will set grad_accum as follows: {accums_per_epoch}")
|
|
|
|
self.per_epoch_grad_accum = accums_per_epoch
|
|
|
|
|
|
def on_epoch_end(self, **kwargs):
|
|
just_finished_epoch = kwargs['epoch']
|
|
epoch = just_finished_epoch + 1
|
|
grad_accum = self.per_epoch_grad_accum.get(epoch)
|
|
if grad_accum is None:
|
|
logging.warning(f" * Accumulnator has no grad_accum setting for epoch {epoch} - leaving as-is")
|
|
else:
|
|
logging.info(f" * Accumulnator setting grad_accum for epoch {epoch} to {grad_accum}")
|
|
arg_update_callback = kwargs['arg_update_callback']
|
|
arg_update_callback('grad_accum', grad_accum)
|
|
|
|
|
|
def _get_update_step_indices(self, epoch, epoch_length_steps: int) -> list[int]:
|
|
if self.every_n_epochs >= 1:
|
|
if ((epoch+1) % self.every_n_epochs) == 0:
|
|
# last step only
|
|
return [epoch_length_steps-1]
|
|
else:
|
|
return []
|
|
else:
|
|
# subdivide the epoch evenly, by rounding self.every_n_epochs to the nearest clean division of steps
|
|
num_divisions = max(1, min(epoch_length_steps, round(1/self.every_n_epochs)))
|
|
# validation happens after training:
|
|
# if an epoch has eg 100 steps and num_divisions is 2, then validation should occur after steps 49 and 99
|
|
validate_every_n_steps = epoch_length_steps / num_divisions
|
|
return [math.ceil((i+1)*validate_every_n_steps) - 1 for i in range(num_divisions)]
|