Spaces:
Running on Zero
Running on Zero
File size: 2,627 Bytes
2e1dc7f | 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 | from typing import Optional
from torch import Tensor, nn
MAXROWLEN = 50
N1 = 17
N2 = MAXROWLEN - N1 - 1
def print_params(model: nn.Module, depth: int = 0, summary: bool = False):
total_params = sum(p.numel() for p in model.parameters())
trainable_params = sum(
p.numel() for p in model.parameters() if p.requires_grad)
name = model._get_name()
s_header = f"{'MODEL:':<{N1}} {name:>{N2}}"
if not summary:
s_header += f" DEVICE GRAD NAN? SHAPE"
print(s_header)
print(f"{'TOTAL PARAMS:':<{N1}} {total_params:>{N2},}")
print(f"{'TRAINABLE PARAMS:':<{N1}} {trainable_params:>{N2},}")
print()
_print_params(model, depth, summary, tabs=0)
def _print_params(model: nn.Module,
depth: int = 0,
summary: bool = False,
tabs=0):
if depth == 0:
if not summary:
for name, p in model.named_parameters():
if len(name) > MAXROWLEN:
name = f"...{name[-37:]}"
print(f"{name:>{MAXROWLEN}} {str(p.device):>6} "
f"{'Grad' if p.requires_grad else '----'} "
f"{'NaNs!' if p.isnan().any() else 'clean'} "
f"{tuple(p.data.shape)}")
return
for childname, childmodule in model.named_children():
if len(list(childmodule.parameters())) == 0:
continue
total_count = sum(p.numel() for p in childmodule.parameters())
trainable_count = sum(
p.numel() for p in childmodule.parameters() if p.requires_grad)
print(" " * tabs + "-" * (len(childname) + 1))
print(" " * tabs + f"{childname.upper()}: {trainable_count:,} "
f"trainable of {total_count:,}")
print(" " * tabs + "-" * (len(childname) + 1))
_print_params(childmodule, depth - 1, summary=summary, tabs=tabs + 1)
return
def printshape(*x: Optional[Tensor]) -> None:
for item in x:
print(f"{list(item.shape) if item is not None else 'None'}")
return
def sanity_check(model: nn.Module, interrupt: bool = True):
clean: bool = True
if not all(p.isfinite().all() for p in model.parameters()):
print("Model has some inf or NaN parameters")
clean = False
if not all(p.grad.isfinite().all()
for p in model.parameters()
if p.requires_grad and p.grad is not None):
print("Model has some inf or NaN gradients")
clean = False
if interrupt and not clean:
raise RuntimeError("Model sanity check failed")
return clean
|