wehub-resource-sync
46 lines
1.9 KiB
Python
46 lines
1.9 KiB
Python
"""Callback that uses the outputs of language models to add AR and TAR regularization
|
|
|
|
Docs: https://docs.fast.ai/callback.rnn.html.md"""
|
|
|
|
# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/34_callback.rnn.ipynb.
|
|
|
|
# %% auto #0
|
|
__all__ = ['ModelResetter', 'RNNCallback', 'RNNRegularizer', 'rnn_cbs']
|
|
|
|
# %% ../../nbs/34_callback.rnn.ipynb #cfe14518
|
|
from ..basics import *
|
|
|
|
# %% ../../nbs/34_callback.rnn.ipynb #9d838495
|
|
@docs
|
|
class ModelResetter(Callback):
|
|
"`Callback` that resets the model at each validation/training step"
|
|
def before_train(self): self.model.reset()
|
|
def before_validate(self): self.model.reset()
|
|
def after_fit(self): self.model.reset()
|
|
_docs = dict(before_train="Reset the model before training",
|
|
before_validate="Reset the model before validation",
|
|
after_fit="Reset the model after fitting")
|
|
|
|
# %% ../../nbs/34_callback.rnn.ipynb #a494c9f3
|
|
class RNNCallback(Callback):
|
|
"Save the raw and dropped-out outputs and only keep the true output for loss computation"
|
|
def after_pred(self): self.learn.pred,self.raw_out,self.out = [o[-1] if is_listy(o) else o for o in self.pred]
|
|
|
|
# %% ../../nbs/34_callback.rnn.ipynb #26f53ba9
|
|
class RNNRegularizer(Callback):
|
|
"Add AR and TAR regularization"
|
|
order,run_valid = RNNCallback.order+1,False
|
|
def __init__(self, alpha=0., beta=0.): store_attr()
|
|
def after_loss(self):
|
|
if not self.training: return
|
|
if self.alpha: self.learn.loss_grad += self.alpha * self.rnn.out.float().pow(2).mean()
|
|
if self.beta:
|
|
h = self.rnn.raw_out
|
|
if len(h)>1: self.learn.loss_grad += self.beta * (h[:,1:] - h[:,:-1]).float().pow(2).mean()
|
|
|
|
# %% ../../nbs/34_callback.rnn.ipynb #2fdc7811
|
|
def rnn_cbs(alpha=0., beta=0.):
|
|
"All callbacks needed for (optionally regularized) RNN training"
|
|
reg = [RNNRegularizer(alpha=alpha, beta=beta)] if alpha or beta else []
|
|
return [ModelResetter(), RNNCallback()] + reg
|