chore: 更新项目名称

astrai/trainer/metric_util.py

@@ -0,0 +1,103 @@
+import torch.nn as nn
+from typing import Dict
+
+
+def grad_norm(model: nn.Module, norm_type: int = 2) -> Dict[str, float]:
+    """Compute gradient norm for each parameter in the model."""
+    norms = {}
+    for name, param in model.named_parameters():
+        norms[name] = 0.0
+        if param.grad:
+            norm = param.grad.data.norm(norm_type).item()
+            norms[name] = norm
+    return norms
+
+
+def grad_std(model: nn.Module) -> Dict[str, float]:
+    """Compute standard deviation of gradients for each parameter."""
+    stds = {}
+    for name, param in model.named_parameters():
+        stds[name] = 0.0
+        if param.grad:
+            std = param.grad.data.std().item()
+            stds[name] = std
+    return stds
+
+
+def grad_max(model: nn.Module) -> Dict[str, float]:
+    """Find the maximum absolute gradient value for each parameter."""
+    max_vals = {}
+    for name, param in model.named_parameters():
+        max_vals[name] = -float("inf")
+        if param.grad:
+            max_val = param.grad.data.max().item()
+            max_vals[name] = max_val
+
+    return max_vals
+
+
+def grad_min(model: nn.Module) -> Dict[str, float]:
+    """Find the minimum absolute gradient value for each parameter."""
+    min_vals = {}
+    for name, param in model.named_parameters():
+        min_vals[name] = float("inf")
+        if param.grad:
+            min_val = param.grad.data.min().item()
+            min_vals[name] = min_val
+
+    return min_vals
+
+
+def grad_mean(model: nn.Module) -> Dict[str, float]:
+    """Compute mean of gradients for each parameter."""
+    means = {}
+    for name, param in model.named_parameters():
+        means[name] = 0.0
+        if param.grad:
+            mean = param.grad.data.mean().item()
+            means[name] = mean
+
+    return means
+
+
+def grad_nan_num(model: nn.Module) -> Dict[str, int]:
+    """Count the number of NaNs in gradients for each parameter."""
+    nan_nums = {}
+    for name, param in model.named_parameters():
+        nan_nums[name] = 0
+        if param.grad:
+            nan_num = param.grad.isnan().sum().item()
+            nan_nums[name] = nan_num
+    return nan_nums
+
+
+def ctx_get_loss(ctx):
+    return ctx.loss
+
+
+def ctx_get_lr(ctx):
+    return ctx.optimizer.param_groups[-1]["lr"]
+
+
+def ctx_get_grad_norm(ctx):
+    return grad_norm(ctx.model)
+
+
+def ctx_get_grad_std(ctx):
+    return grad_std(ctx.model)
+
+
+def ctx_get_grad_max(ctx):
+    return grad_max(ctx.model)
+
+
+def ctx_get_grad_min(ctx):
+    return grad_min(ctx.model)
+
+
+def ctx_get_grad_mean(ctx):
+    return grad_mean(ctx.model)
+
+
+def ctx_get_grad_nan_num(ctx):
+    return grad_nan_num(ctx.model)