wehub-resource-sync
191 lines
9.3 KiB
Python
191 lines
9.3 KiB
Python
"""Basic callbacks for Learner
|
|
|
|
Docs: https://docs.fast.ai/callback.core.html.md"""
|
|
|
|
# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/13_callback.core.ipynb.
|
|
|
|
# %% auto #0
|
|
__all__ = ['Callback', 'TrainEvalCallback', 'GatherPredsCallback', 'FetchPredsCallback', 'CancelStepException',
|
|
'CancelBackwardException', 'CancelFitException', 'CancelEpochException', 'CancelTrainException',
|
|
'CancelValidException', 'CancelBatchException', 'event']
|
|
|
|
# %% ../../nbs/13_callback.core.ipynb #93fea1ce
|
|
from ..data.all import *
|
|
from ..optimizer import *
|
|
from ..losses import BaseLoss
|
|
|
|
# %% ../../nbs/13_callback.core.ipynb #80cdee5a
|
|
_all_ = ['CancelStepException','CancelBackwardException','CancelFitException','CancelEpochException','CancelTrainException','CancelValidException','CancelBatchException']
|
|
|
|
# %% ../../nbs/13_callback.core.ipynb #ab22dd60
|
|
_events = L.split('after_create before_fit before_epoch before_train before_batch after_pred after_loss \
|
|
before_backward after_cancel_backward after_backward before_step after_cancel_step after_step \
|
|
after_cancel_batch after_batch after_cancel_train after_train before_validate after_cancel_validate \
|
|
after_validate after_cancel_epoch after_epoch after_cancel_fit after_fit')
|
|
|
|
mk_class('event', **_events.map_dict(),
|
|
doc="All possible events as attributes to get tab-completion and typo-proofing")
|
|
|
|
# %% ../../nbs/13_callback.core.ipynb #0ed2aba5
|
|
_all_ = ['event']
|
|
|
|
# %% ../../nbs/13_callback.core.ipynb #05cd2800
|
|
_inner_loop = "before_batch after_pred after_loss before_backward after_cancel_backward after_backward before_step after_step after_cancel_batch after_batch".split()
|
|
|
|
# %% ../../nbs/13_callback.core.ipynb #9eba90b5
|
|
_ex_docs = dict(
|
|
CancelBatchException="Skip the rest of this batch and go to `after_batch`",
|
|
CancelTrainException="Skip the rest of the training part of the epoch and go to `after_train`",
|
|
CancelValidException="Skip the rest of the validation part of the epoch and go to `after_validate`",
|
|
CancelEpochException="Skip the rest of this epoch and go to `after_epoch`",
|
|
CancelStepException ="Skip stepping the optimizer",
|
|
CancelBackwardException="Skip the backward pass and go to `after_backward`",
|
|
CancelFitException ="Interrupts training and go to `after_fit`")
|
|
|
|
for c,d in _ex_docs.items(): mk_class(c,sup=Exception,doc=d)
|
|
|
|
# %% ../../nbs/13_callback.core.ipynb #011373d8
|
|
@funcs_kwargs(as_method=True)
|
|
class Callback(Stateful,GetAttr):
|
|
"Basic class handling tweaks of the training loop by changing a `Learner` in various events"
|
|
order,_default,learn,run,run_train,run_valid = 0,'learn',None,True,True,True
|
|
_methods = _events
|
|
|
|
def __init__(self, **kwargs): assert not kwargs, f'Passed unknown events: {kwargs}'
|
|
def __repr__(self): return type(self).__name__
|
|
|
|
def __call__(self, event_name):
|
|
"Call `self.{event_name}` if it's defined"
|
|
_run = (event_name not in _inner_loop or (self.run_train and getattr(self, 'training', True)) or
|
|
(self.run_valid and not getattr(self, 'training', False)))
|
|
res = None
|
|
if self.run and _run:
|
|
try: res = getcallable(self, event_name)()
|
|
except (CancelBatchException, CancelBackwardException, CancelEpochException, CancelFitException, CancelStepException, CancelTrainException, CancelValidException): raise
|
|
except Exception as e: raise modify_exception(e, f'Exception occured in `{self.__class__.__name__}` when calling event `{event_name}`:\n\t{e.args[0]}', replace=True)
|
|
if event_name=='after_fit': self.run=True #Reset self.run to True at each end of fit
|
|
return res
|
|
|
|
def __setattr__(self, name, value):
|
|
"Set an attribute for a `Callback`"
|
|
if hasattr(self.learn,name):
|
|
warn(f"You are shadowing an attribute ({name}) that exists in the learner. Use `self.learn.{name}` to avoid this")
|
|
super().__setattr__(name, value)
|
|
|
|
@property
|
|
def name(self):
|
|
"Name of the `Callback`, camel-cased and with '*Callback*' removed"
|
|
return class2attr(self, 'Callback')
|
|
|
|
# %% ../../nbs/13_callback.core.ipynb #d8968bed
|
|
class TrainEvalCallback(Callback):
|
|
"`Callback` that tracks the number of iterations done and properly sets training/eval mode"
|
|
order,run_valid = -10,False
|
|
def after_create(self): self.learn.n_epoch = 1
|
|
|
|
def before_fit(self):
|
|
"Set the iter and epoch counters to 0, put the model and the right device"
|
|
self.learn.epoch,self.learn.loss = 0,tensor(0.)
|
|
self.learn.train_iter,self.learn.pct_train = 0,0.
|
|
device = getattr(self.dls, 'device', default_device())
|
|
self.model.to(device)
|
|
if isinstance(self.loss_func, (nn.Module, BaseLoss)): self.loss_func.to(device)
|
|
if hasattr(self.model, 'reset'): self.model.reset()
|
|
|
|
def after_batch(self):
|
|
"Update the iter counter (in training mode)"
|
|
self.learn.pct_train += 1./(self.n_iter*self.n_epoch)
|
|
self.learn.train_iter += 1
|
|
|
|
def before_train(self):
|
|
"Set the model to training mode"
|
|
self.learn.pct_train=self.epoch/self.n_epoch
|
|
self.model.train()
|
|
self.learn.training=True
|
|
|
|
def before_validate(self):
|
|
"Set the model to validation mode"
|
|
self.model.eval()
|
|
self.learn.training=False
|
|
|
|
# %% ../../nbs/13_callback.core.ipynb #33328a26
|
|
if not hasattr(defaults, 'callbacks'): defaults.callbacks = [TrainEvalCallback]
|
|
|
|
# %% ../../nbs/13_callback.core.ipynb #a558f93f
|
|
class GatherPredsCallback(Callback):
|
|
"`Callback` that returns all predictions and targets, optionally `with_input` or `with_loss`"
|
|
_stateattrs=('preds','targets','inputs','losses')
|
|
def __init__(self,
|
|
with_input:bool=False, # Whether to return inputs
|
|
with_loss:bool=False, # Whether to return losses
|
|
save_preds:Path=None, # Path to save predictions
|
|
save_targs:Path=None, # Path to save targets
|
|
with_preds:bool=True, # Whether to return predictions
|
|
with_targs:bool=True, # Whether to return targets
|
|
concat_dim:int=0, # Dimension to concatenate returned tensors
|
|
pickle_protocol:int=2 # Pickle protocol used to save predictions and targets
|
|
):
|
|
store_attr()
|
|
|
|
def before_batch(self):
|
|
"If `with_input`, detach batch inputs"
|
|
if self.with_input: self.inputs.append((self.learn.to_detach(self.xb)))
|
|
|
|
def before_validate(self):
|
|
"Initialize containers"
|
|
self.preds,self.targets = [],[]
|
|
if self.with_input: self.inputs = []
|
|
if self.with_loss: self.losses = []
|
|
|
|
def after_batch(self):
|
|
"Save predictions, targets and potentially losses"
|
|
if not hasattr(self, 'pred'): return
|
|
preds,targs = self.learn.to_detach(self.pred),self.learn.to_detach(self.yb)
|
|
if self.with_preds: self.preds.append(preds)
|
|
if self.with_targs: self.targets.append(targs)
|
|
if self.save_preds is not None:
|
|
torch.save(preds, self.save_preds/str(self.iter), pickle_protocol=self.pickle_protocol)
|
|
if self.save_targs is not None:
|
|
torch.save(targs[0], self.save_targs/str(self.iter), pickle_protocol=self.pickle_protocol)
|
|
if self.with_loss:
|
|
bs = find_bs(self.yb)
|
|
loss = self.loss if self.loss.numel() == bs else self.loss.view(bs,-1).mean(1)
|
|
self.losses.append(self.learn.to_detach(loss))
|
|
|
|
def after_validate(self):
|
|
"Concatenate all recorded tensors"
|
|
if not hasattr(self, 'preds'): return
|
|
if self.with_input: self.inputs = detuplify(to_concat(self.inputs, dim=self.concat_dim))
|
|
if self.with_preds: self.preds = detuplify(to_concat(self.preds, dim=self.concat_dim))
|
|
if self.with_targs: self.targets = detuplify(to_concat(self.targets, dim=self.concat_dim))
|
|
if self.with_loss: self.losses = to_concat(self.losses)
|
|
|
|
def all_tensors(self) -> (Tensor, list):
|
|
"Returns all recorded tensors in the order [inputs, preds, targets, losses]"
|
|
res = [self.preds if self.with_preds else None, self.targets if self.with_targs else None]
|
|
if self.with_input: res = [self.inputs] + res
|
|
if self.with_loss: res.append(self.losses)
|
|
return res
|
|
|
|
# %% ../../nbs/13_callback.core.ipynb #1c76da8f
|
|
class FetchPredsCallback(Callback):
|
|
"A callback to fetch predictions during the training loop"
|
|
remove_on_fetch = True
|
|
def __init__(self,
|
|
ds_idx:int=1, # Index of dataset, 0 for train, 1 for valid, used if `dl` is not present
|
|
dl:DataLoader=None, # `DataLoader` used for fetching `Learner` predictions
|
|
with_input:bool=False, # Whether to return inputs in `GatherPredsCallback`
|
|
with_decoded:bool=False, # Whether to return decoded predictions
|
|
cbs:Callback|MutableSequence=None, # `Callback` to temporarily remove from `Learner`
|
|
reorder:bool=True # Whether to sort prediction results
|
|
):
|
|
self.cbs = L(cbs)
|
|
store_attr('ds_idx,dl,with_input,with_decoded,reorder')
|
|
|
|
def after_validate(self):
|
|
"Fetch predictions from `Learner` without `self.cbs` and `remove_on_fetch` callbacks"
|
|
to_rm = L(cb for cb in self.learn.cbs if getattr(cb, 'remove_on_fetch', False))
|
|
with self.learn.removed_cbs(to_rm + self.cbs) as learn:
|
|
self.preds = learn.get_preds(ds_idx=self.ds_idx, dl=self.dl,
|
|
with_input=self.with_input, with_decoded=self.with_decoded, inner=True, reorder=self.reorder)
|