Spaces:
Sleeping
Sleeping
File size: 2,548 Bytes
e45d058 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 |
import pytorch_lightning as pl
from pytorch_lightning import Callback
from pytorch_lightning.utilities import rank_zero_only
import torch
from torch.autograd import grad
class CausalityMonitor(Callback):
r"""Monitor causality of a model by tracking gradient leakage forward in time.
In a fully causal model, dy[k]du[s] ~= 0 for all k < s.
Args:
seq_len (int): Length of the sequence to monitor.
input_dim (int): Dimension of the input to monitor. If 0, the callback assumes
the task to be language modeling, and skips the embedding layer. If > 0,
input_dim is interpreted as the input channel dimension, i.e. D with
dummy input of dimension [B, L, D].
Notes:
This callback assumes that `pl_module.model` has a `net` or `s4seq` attribute,
indicating the primary model to monitor. For LMs, `net` or `s4seq` should
be after the embedding layer.
"""
def __init__(self, seq_len: int = 10, input_dim: int = 0):
super().__init__()
self.seq_len = seq_len
self.input_dim = input_dim
@rank_zero_only
def on_train_epoch_end(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None:
model = pl_module.model
with torch.enable_grad():
if self.input_dim == 0:
# [MP] LongTensors cannot have gradients - we start from post
# embedding in the LM case
input_dim = model.d_model
x = torch.randn((2, self.seq_len, input_dim), \
requires_grad=True).to(pl_module.device)
# [DF] HACK: we need to get the layer that comes after the embedding
if hasattr(model, 'net'):
y = model.net(x)
else:
y = model.s4seq(x)
else:
x = torch.randn(1, self.seq_len, self.input_dim, \
requires_grad=True).to(pl_module.device)
y = model(x)
stats = {}
for i in range(self.seq_len):
# total gradients flowing from y_i to x
g = grad(y[0,0,i].mean(), x, retain_graph=True, allow_unused=True)[0]
g = g[0,i+1:,:].abs().mean()
stats[f'stats/causality_{i}'] = g.item()
if trainer.loggers is not None:
for logger in trainer.loggers:
logger.log_metrics(stats, step=trainer.global_step)
|