mirror of
https://github.com/wassname/pytorch-lightning.git
synced 2026-09-09 11:32:07 +08:00
3c2fd560aa
* first pass for LightningModule typehints * fix return types * add missing types * add type annotations to grads.py * add type annotations to hooks.py * add type annotation to memory.py * proper docstring quotation marks * add type annotations to saving.py * fix cyclic import problem * fix cyclic import problem * add missing whitespace * finish type hints for load_from_ methods * docs: prepare_data does not return anything * fix auto types in docs * revert typehint for trainer in hook * remove unnecessary return docs * some fixes for memory docs * revert typing for args kwargs * added all missing None return types * remove unused import * add more details to dict/list return types * fix line too long * optimize imports * linted * Revert "linted" This reverts commit 85559611e84e312bce64f4e73b638d4999a8439e. * remove whitespace * update * update * update * update * update * changelog Co-authored-by: Jirka Borovec <Borda@users.noreply.github.com> Co-authored-by: William Falcon <waf2107@columbia.edu>
31 lines
966 B
Python
31 lines
966 B
Python
"""
|
|
Module to describe gradients
|
|
"""
|
|
from typing import Dict
|
|
|
|
from torch import nn
|
|
|
|
|
|
class GradInformation(nn.Module):
|
|
|
|
def grad_norm(self, norm_type: float) -> Dict[str, int]:
|
|
results = {}
|
|
total_norm = 0
|
|
for name, p in self.named_parameters():
|
|
if p.requires_grad:
|
|
try:
|
|
param_norm = p.grad.data.norm(norm_type)
|
|
total_norm += param_norm ** norm_type
|
|
norm = param_norm ** (1 / norm_type)
|
|
|
|
grad = round(norm.data.cpu().numpy().flatten()[0], 3)
|
|
results['grad_{}_norm_{}'.format(norm_type, name)] = grad
|
|
except Exception:
|
|
# this param had no grad
|
|
pass
|
|
|
|
total_norm = total_norm ** (1. / norm_type)
|
|
grad = round(total_norm.data.cpu().numpy().flatten()[0], 3)
|
|
results['grad_{}_norm_total'.format(norm_type)] = grad
|
|
return results
|