/
parameter_statistics.py
155 lines (131 loc) · 6.69 KB
/
parameter_statistics.py
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
import numpy
import six
import chainer
from chainer.backends import cuda
from chainer import reporter
from chainer.training import extension
from chainer.training import trigger as trigger_module
class ParameterStatistics(extension.Extension):
"""Trainer extension to report parameter statistics.
Statistics are collected and reported for a given :class:`~chainer.Link`
or an iterable of :class:`~chainer.Link`\\ s. If a link contains child
links, the statistics are reported separately for each child.
Any function that takes a one-dimensional :class:`numpy.ndarray` or a
:class:`cupy.ndarray` and outputs a single or multiple real numbers can be
registered to handle the collection of statistics, e.g.
:meth:`numpy.ndarray.mean`.
The keys of reported statistics follow the convention of link name
followed by parameter name, attribute name and function name, e.g.
``VGG16Layers/conv1_1/W/data/mean``. They are prepended with an optional
prefix and appended with integer indices if the statistics generating
function return multiple values.
Args:
links (~chainer.Link or iterable of ~chainer.Link): Link(s) containing
the parameters to observe. The link is expected to have a ``name``
attribute which is used as a part of the report key.
statistics (dict): Dictionary with function name to function mappings.
The name is a string and is used as a part of the report key. The
function is responsible for generating the statistics.
report_params (bool): If ``True``, report statistics for parameter
values such as weights and biases.
report_grads (bool): If ``True``, report statistics for parameter
gradients.
prefix (str): Optional prefix to prepend to the report keys.
trigger: Trigger that decides when to aggregate the results and report
the values.
skip_nan_params (bool): If ``True``, statistics are not computed for
parameters including NaNs and a single NaN value is immediately
reported instead. Otherwise, this extension will simply try to
compute the statistics without performing any checks for NaNs.
"""
default_name = 'parameter_statistics'
priority = extension.PRIORITY_WRITER
# prefix ends with a '/' and param_name is preceded by a '/'
report_key_template = ('{prefix}{link_name}{param_name}/{attr_name}/'
'{function_name}')
default_statistics = {
'mean': lambda x: cuda.get_array_module(x).mean(x),
'std': lambda x: cuda.get_array_module(x).std(x),
'min': lambda x: cuda.get_array_module(x).min(x),
'max': lambda x: cuda.get_array_module(x).max(x),
'zeros': lambda x: cuda.get_array_module(x).count_nonzero(x == 0),
'percentile': lambda x: cuda.get_array_module(x).percentile(
x, (0.13, 2.28, 15.87, 50, 84.13, 97.72, 99.87))
}
def __init__(self, links, statistics=default_statistics,
report_params=True, report_grads=True, prefix=None,
trigger=(1, 'epoch'), skip_nan_params=False):
if not isinstance(links, (list, tuple)):
links = links,
self._links = links
if statistics is None:
statistics = {}
self._statistics = statistics
attrs = []
if report_params:
attrs.append('data')
if report_grads:
attrs.append('grad')
self._attrs = attrs
self._prefix = prefix
self._trigger = trigger_module.get_trigger(trigger)
self._summary = reporter.DictSummary()
self._skip_nan_params = skip_nan_params
def __call__(self, trainer):
"""Execute the statistics extension.
Collect statistics for the current state of parameters.
Note that this method will merely update its statistic summary, unless
the internal trigger is fired. If the trigger is fired, the summary
will also be reported and then reset for the next accumulation.
Args:
trainer (~chainer.training.Trainer): Associated trainer that
invoked this extension.
"""
statistics = {}
for link in self._links:
link_name = getattr(link, 'name', 'None')
for param_name, param in link.namedparams():
for attr_name in self._attrs:
for function_name, function in \
six.iteritems(self._statistics):
# Get parameters as a flattend one-dimensional array
# since the statistics function should make no
# assumption about the axes
params = getattr(param, attr_name).ravel()
if (self._skip_nan_params
and (cuda.get_array_module(params).isnan(params)
.any())):
value = numpy.nan
else:
value = function(params)
key = self.report_key_template.format(
prefix=self._prefix + '/' if self._prefix else '',
link_name=link_name,
param_name=param_name,
attr_name=attr_name,
function_name=function_name
)
if (isinstance(value, chainer.get_array_types())
and value.size > 1):
# Append integer indices to the keys if the
# statistic function return multiple values
statistics.update({'{}/{}'.format(key, i): v for
i, v in enumerate(value)})
else:
statistics[key] = value
self._summary.add(statistics)
if self._trigger(trainer):
reporter.report(self._summary.compute_mean())
self._summary = reporter.DictSummary() # Clear summary
def register_statistics(self, name, function):
"""Register a function to compute a certain statistic.
The registered function will be called each time the extension runs and
the results will be included in the report.
Args:
name (str): Name of the statistic.
function: Function to generate the statistic. Any function that
takes a one-dimensional :class:`numpy.ndarray` or a
:class:`cupy.ndarray` and outputs a single or multiple real
numbers is allowed.
"""
self._statistics[name] = function