File size: 8,744 Bytes
8e04e6f | 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 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 | import time
import torch
from lightning.pytorch.callbacks import Callback
# from pytorch_lightning.utilities.rank_zero import rank_zero_only # Do not use here, breaks clsuter
from loguru import logger
def log_metrics(pl_module, metrics):
for m in metrics:
pl_module.log(
m, metrics[m], on_step=True, on_epoch=True, prog_bar=False, logger=True
)
class UnusedParametersCallback(Callback):
def on_before_optimizer_step(self, trainer, pl_module, optimizer):
for name, param in pl_module.named_parameters():
if param.grad is None:
print(f"Unused parameter: {name}")
class LogEpochTimeCallback(Callback):
"""Simple callback that logs how long each epoch takes, in seconds, to a pytorch lightning log"""
def on_train_epoch_start(self, trainer, pl_module):
self.epoch_start = time.time()
def on_train_epoch_end(self, trainer, pl_module):
curr_time = time.time()
duration = curr_time # - self.epoch_start
pl_module.log(
"train_info/epoch_duration_secs",
duration,
on_step=False,
on_epoch=True,
prog_bar=False,
logger=True,
sync_dist=True,
)
if pl_module.current_epoch % 10 == 0:
logger.info(
f"Done training epoch {pl_module.current_epoch}, epoch took {duration} seconds"
)
class LogSetpTimeCallback(Callback):
"""Simple callback that logs how long each training step takes, in seconds, to a pytorch lightning log"""
def on_train_batch_start(self, trainer, pl_module, batch, batch_idx):
self.step_start = time.time()
def on_train_batch_end(self, trainer, pl_module, outputs, batch, batch_idx):
curr_time = time.time()
duration = curr_time - self.step_start
pl_module.log(
"train_info/step_duration_secs",
duration,
on_step=True,
on_epoch=False,
prog_bar=False,
logger=True,
sync_dist=True,
)
class GradAndWeightAnalysisCallback(Callback):
"""Some functionality to observe how are weights and gradientsbehaving during trainnig."""
def __init__(self, debug=True, moving_avg_size=100):
super(GradAndWeightAnalysisCallback, self).__init__()
self.debug = debug
self.moving_avg_size = moving_avg_size
self.avg_grad_history = []
self.max_grad_history = []
def _get_avg_and_max_w(self, pl_module):
"""Computes average and max weight in module."""
with torch.no_grad():
params = torch.nn.utils.parameters_to_vector(pl_module.parameters()).abs()
return params.sum() / params.numel(), params.max()
def on_before_optimizer_step(self, trainer, pl_module, optimizer):
"""Before updating log max and average weight."""
avg_w, max_w = self._get_avg_and_max_w(pl_module)
avg_g, max_g = self._get_avg_and_max_grad(pl_module)
metrics = {}
# Gradient metrics
metrics["avg_w_bef_step"] = avg_w
metrics["max_w_bef_step"] = max_w
metrics["avg_g_bef_step"] = avg_g
metrics["max_g_bef_step"] = max_g
# Add gradient metrics
if len(self.max_grad_history) > 1:
metrics["moving_avg_max_grad_bef_step"] = sum(self.max_grad_history) / len(
self.max_grad_history
)
metrics["moving_avg_avg_grad_bef_step"] = sum(self.avg_grad_history) / len(
self.avg_grad_history
)
metrics["max_g_over_avg_max_g_bef_step"] = (
max_g / metrics["moving_avg_max_grad_bef_step"]
)
metrics["avg_g_over_avg_avg_g_bef_step"] = (
avg_g / metrics["moving_avg_avg_grad_bef_step"]
)
# Update history of gradient statistics
if len(self.max_grad_history) >= self.moving_avg_size:
self.max_grad_history.pop(0)
if not (max_g.isnan().any() or max_g.isinf().any()):
self.max_grad_history.append(max_g.item())
if len(self.avg_grad_history) >= self.moving_avg_size:
self.avg_grad_history.pop(0)
if not (avg_g.isnan().any() or avg_g.isinf().any()):
self.avg_grad_history.append(avg_g.item())
if self.debug and (avg_w.isnan().any() or max_w.isnan().any()):
params = torch.nn.utils.parameters_to_vector(pl_module.parameters()).abs()
log_metrics(pl_module, metrics)
def on_train_batch_start(self, trainer, pl_module, batch, batch_idx):
"""First thing on training step."""
avg_w, max_w = self._get_avg_and_max_w(pl_module)
def _get_avg_and_max_grad(self, pl_module):
"""Computes average and max grad in module, if no grad computed zero."""
grad_sum = torch.tensor(0.0, device=pl_module.device)
max_grad = torch.tensor(0.0, device=pl_module.device)
count = 0
for p in pl_module.parameters():
if p.grad is not None:
abs_grad = p.grad.abs()
grad_sum += abs_grad.sum()
max_grad = torch.max(max_grad, abs_grad.max())
count += p.grad.numel()
if count == 0:
return torch.tensor(0.0), torch.tensor(0.0)
return grad_sum / count, max_grad
def _count_nan_grad(self, pl_module):
numels, num_nans = 0, 0
for p in pl_module.parameters():
if p.grad is not None:
numels += p.grad.numel()
num_nans += p.grad.isnan().sum().item()
return numels, num_nans
def on_before_zero_grad(self, trainer, pl_module, optimizer):
"""Before zero-ing grad log max and average grad
(should be shifted one back from next)."""
avg_g, max_g = self._get_avg_and_max_grad(pl_module)
metrics = {
"avg_g_bef_zerog": avg_g,
"max_g_bef_zerog": max_g,
}
log_metrics(pl_module, metrics)
def on_after_backward(self, trainer, pl_module):
"""After computing gradient log max and average grad.
This same value should be returned by <on_before_zxero_grad>
in the next iteration."""
avg_g, max_g = self._get_avg_and_max_grad(pl_module)
numels, num_nans = self._count_nan_grad(pl_module)
metrics = {
"avg_g_after_bwd": avg_g,
"max_g_after_bwd": max_g,
}
log_metrics(pl_module, metrics)
class SkipNanGradCallback(Callback):
"""Callback to skip gradient updates with NaN in them."""
def __init__(self, debug=True):
super(SkipNanGradCallback, self).__init__()
self.count = 0
self.iter = 0
def on_after_backward(self, trainer, pl_module):
nan_flag = False
self.iter += 1
for p in pl_module.parameters():
if p.grad is not None:
# has_grad = True
if p.grad.isnan().any():
nan_flag = True
if nan_flag:
self.count += 1
pl_module.zero_grad()
class SkipLargeGradients(Callback):
"""
This callback tracks the average max gradient for the last `moving_avg_size` steps and checks if the current max gradient
"""
def __init__(
self,
moving_avg_size: int = 100,
factor_threshold: int = 5,
min_opt_steps: int = 2000,
):
self.max_g_history = []
self.moving_avg_size = moving_avg_size
self.factor_threshold = factor_threshold
self.min_opt_steps = min_opt_steps
def on_before_optimizer_step(self, trainer, pl_module, optimizer):
# calculate the maximum gradient
max_g = max(
param.grad.data.abs().max()
for param in pl_module.parameters()
if param.grad is not None
)
# Update the moving average history
if len(self.max_g_history) >= self.moving_avg_size:
self.max_g_history.pop(0)
if not (max_g.isnan().any() or max_g.isinf().any()):
self.max_g_history.append(max_g.item())
if trainer.global_step > self.min_opt_steps:
if len(self.max_g_history) > 1:
moving_average = sum(self.max_g_history) / len(self.max_g_history)
# Check if the current max gradient exceeds factor_threshold * moving average
if max_g > self.factor_threshold * moving_average:
# Zero out gradients
for param in pl_module.parameters():
if param.grad is not None:
param.grad.data.zero_() # Set gradients to zero
|