269 lines
10 KiB
Python
269 lines
10 KiB
Python
# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
|
#
|
|
# NVIDIA CORPORATION and its licensors retain all intellectual property
|
|
# and proprietary rights in and to this software, related documentation
|
|
# and any modifications thereto. Any use, reproduction, disclosure or
|
|
# distribution of this software and related documentation without an express
|
|
# license agreement from NVIDIA CORPORATION is strictly prohibited.
|
|
|
|
"""Facilities for reporting and collecting training statistics across
|
|
multiple processes and devices. The interface is designed to minimize
|
|
synchronization overhead as well as the amount of boilerplate in user
|
|
code."""
|
|
|
|
import re
|
|
import numpy as np
|
|
import torch
|
|
import dnnlib
|
|
|
|
from . import misc
|
|
|
|
#----------------------------------------------------------------------------
|
|
|
|
_num_moments = 3 # [num_scalars, sum_of_scalars, sum_of_squares]
|
|
_reduce_dtype = torch.float32 # Data type to use for initial per-tensor reduction.
|
|
_counter_dtype = torch.float64 # Data type to use for the internal counters.
|
|
_rank = 0 # Rank of the current process.
|
|
_sync_device = None # Device to use for multiprocess communication. None = single-process.
|
|
_sync_called = False # Has _sync() been called yet?
|
|
_counters = dict() # Running counters on each device, updated by report(): name => device => torch.Tensor
|
|
_cumulative = dict() # Cumulative counters on the CPU, updated by _sync(): name => torch.Tensor
|
|
|
|
#----------------------------------------------------------------------------
|
|
|
|
def init_multiprocessing(rank, sync_device):
|
|
r"""Initializes `torch_utils.training_stats` for collecting statistics
|
|
across multiple processes.
|
|
|
|
This function must be called after
|
|
`torch.distributed.init_process_group()` and before `Collector.update()`.
|
|
The call is not necessary if multi-process collection is not needed.
|
|
|
|
Args:
|
|
rank: Rank of the current process.
|
|
sync_device: PyTorch device to use for inter-process
|
|
communication, or None to disable multi-process
|
|
collection. Typically `torch.device('cuda', rank)`.
|
|
"""
|
|
global _rank, _sync_device
|
|
assert not _sync_called
|
|
_rank = rank
|
|
_sync_device = sync_device
|
|
|
|
#----------------------------------------------------------------------------
|
|
|
|
@misc.profiled_function
|
|
def report(name, value):
|
|
r"""Broadcasts the given set of scalars to all interested instances of
|
|
`Collector`, across device and process boundaries.
|
|
|
|
This function is expected to be extremely cheap and can be safely
|
|
called from anywhere in the training loop, loss function, or inside a
|
|
`torch.nn.Module`.
|
|
|
|
Warning: The current implementation expects the set of unique names to
|
|
be consistent across processes. Please make sure that `report()` is
|
|
called at least once for each unique name by each process, and in the
|
|
same order. If a given process has no scalars to broadcast, it can do
|
|
`report(name, [])` (empty list).
|
|
|
|
Args:
|
|
name: Arbitrary string specifying the name of the statistic.
|
|
Averages are accumulated separately for each unique name.
|
|
value: Arbitrary set of scalars. Can be a list, tuple,
|
|
NumPy array, PyTorch tensor, or Python scalar.
|
|
|
|
Returns:
|
|
The same `value` that was passed in.
|
|
"""
|
|
if name not in _counters:
|
|
_counters[name] = dict()
|
|
|
|
elems = torch.as_tensor(value)
|
|
if elems.numel() == 0:
|
|
return value
|
|
|
|
elems = elems.detach().flatten().to(_reduce_dtype)
|
|
moments = torch.stack([
|
|
torch.ones_like(elems).sum(),
|
|
elems.sum(),
|
|
elems.square().sum(),
|
|
])
|
|
assert moments.ndim == 1 and moments.shape[0] == _num_moments
|
|
moments = moments.to(_counter_dtype)
|
|
|
|
device = moments.device
|
|
if device not in _counters[name]:
|
|
_counters[name][device] = torch.zeros_like(moments)
|
|
_counters[name][device].add_(moments)
|
|
return value
|
|
|
|
#----------------------------------------------------------------------------
|
|
|
|
def report0(name, value):
|
|
r"""Broadcasts the given set of scalars by the first process (`rank = 0`),
|
|
but ignores any scalars provided by the other processes.
|
|
See `report()` for further details.
|
|
"""
|
|
report(name, value if _rank == 0 else [])
|
|
return value
|
|
|
|
#----------------------------------------------------------------------------
|
|
|
|
class Collector:
|
|
r"""Collects the scalars broadcasted by `report()` and `report0()` and
|
|
computes their long-term averages (mean and standard deviation) over
|
|
user-defined periods of time.
|
|
|
|
The averages are first collected into internal counters that are not
|
|
directly visible to the user. They are then copied to the user-visible
|
|
state as a result of calling `update()` and can then be queried using
|
|
`mean()`, `std()`, `as_dict()`, etc. Calling `update()` also resets the
|
|
internal counters for the next round, so that the user-visible state
|
|
effectively reflects averages collected between the last two calls to
|
|
`update()`.
|
|
|
|
Args:
|
|
regex: Regular expression defining which statistics to
|
|
collect. The default is to collect everything.
|
|
keep_previous: Whether to retain the previous averages if no
|
|
scalars were collected on a given round
|
|
(default: True).
|
|
"""
|
|
def __init__(self, regex='.*', keep_previous=True):
|
|
self._regex = re.compile(regex)
|
|
self._keep_previous = keep_previous
|
|
self._cumulative = dict()
|
|
self._moments = dict()
|
|
self.update()
|
|
self._moments.clear()
|
|
|
|
def names(self):
|
|
r"""Returns the names of all statistics broadcasted so far that
|
|
match the regular expression specified at construction time.
|
|
"""
|
|
return [name for name in _counters if self._regex.fullmatch(name)]
|
|
|
|
def update(self):
|
|
r"""Copies current values of the internal counters to the
|
|
user-visible state and resets them for the next round.
|
|
|
|
If `keep_previous=True` was specified at construction time, the
|
|
operation is skipped for statistics that have received no scalars
|
|
since the last update, retaining their previous averages.
|
|
|
|
This method performs a number of GPU-to-CPU transfers and one
|
|
`torch.distributed.all_reduce()`. It is intended to be called
|
|
periodically in the main training loop, typically once every
|
|
N training steps.
|
|
"""
|
|
if not self._keep_previous:
|
|
self._moments.clear()
|
|
for name, cumulative in _sync(self.names()):
|
|
if name not in self._cumulative:
|
|
self._cumulative[name] = torch.zeros([_num_moments], dtype=_counter_dtype)
|
|
delta = cumulative - self._cumulative[name]
|
|
self._cumulative[name].copy_(cumulative)
|
|
if float(delta[0]) != 0:
|
|
self._moments[name] = delta
|
|
|
|
def _get_delta(self, name):
|
|
r"""Returns the raw moments that were accumulated for the given
|
|
statistic between the last two calls to `update()`, or zero if
|
|
no scalars were collected.
|
|
"""
|
|
assert self._regex.fullmatch(name)
|
|
if name not in self._moments:
|
|
self._moments[name] = torch.zeros([_num_moments], dtype=_counter_dtype)
|
|
return self._moments[name]
|
|
|
|
def num(self, name):
|
|
r"""Returns the number of scalars that were accumulated for the given
|
|
statistic between the last two calls to `update()`, or zero if
|
|
no scalars were collected.
|
|
"""
|
|
delta = self._get_delta(name)
|
|
return int(delta[0])
|
|
|
|
def mean(self, name):
|
|
r"""Returns the mean of the scalars that were accumulated for the
|
|
given statistic between the last two calls to `update()`, or NaN if
|
|
no scalars were collected.
|
|
"""
|
|
delta = self._get_delta(name)
|
|
if int(delta[0]) == 0:
|
|
return float('nan')
|
|
return float(delta[1] / delta[0])
|
|
|
|
def std(self, name):
|
|
r"""Returns the standard deviation of the scalars that were
|
|
accumulated for the given statistic between the last two calls to
|
|
`update()`, or NaN if no scalars were collected.
|
|
"""
|
|
delta = self._get_delta(name)
|
|
if int(delta[0]) == 0 or not np.isfinite(float(delta[1])):
|
|
return float('nan')
|
|
if int(delta[0]) == 1:
|
|
return float(0)
|
|
mean = float(delta[1] / delta[0])
|
|
raw_var = float(delta[2] / delta[0])
|
|
return np.sqrt(max(raw_var - np.square(mean), 0))
|
|
|
|
def as_dict(self):
|
|
r"""Returns the averages accumulated between the last two calls to
|
|
`update()` as an `dnnlib.EasyDict`. The contents are as follows:
|
|
|
|
dnnlib.EasyDict(
|
|
NAME = dnnlib.EasyDict(num=FLOAT, mean=FLOAT, std=FLOAT),
|
|
...
|
|
)
|
|
"""
|
|
stats = dnnlib.EasyDict()
|
|
for name in self.names():
|
|
stats[name] = dnnlib.EasyDict(num=self.num(name), mean=self.mean(name), std=self.std(name))
|
|
return stats
|
|
|
|
def __getitem__(self, name):
|
|
r"""Convenience getter.
|
|
`collector[name]` is a synonym for `collector.mean(name)`.
|
|
"""
|
|
return self.mean(name)
|
|
|
|
#----------------------------------------------------------------------------
|
|
|
|
def _sync(names):
|
|
r"""Synchronize the global cumulative counters across devices and
|
|
processes. Called internally by `Collector.update()`.
|
|
"""
|
|
if len(names) == 0:
|
|
return []
|
|
global _sync_called
|
|
_sync_called = True
|
|
|
|
# Collect deltas within current rank.
|
|
deltas = []
|
|
device = _sync_device if _sync_device is not None else torch.device('cpu')
|
|
for name in names:
|
|
delta = torch.zeros([_num_moments], dtype=_counter_dtype, device=device)
|
|
for counter in _counters[name].values():
|
|
delta.add_(counter.to(device))
|
|
counter.copy_(torch.zeros_like(counter))
|
|
deltas.append(delta)
|
|
deltas = torch.stack(deltas)
|
|
|
|
# Sum deltas across ranks.
|
|
if _sync_device is not None:
|
|
torch.distributed.all_reduce(deltas)
|
|
|
|
# Update cumulative values.
|
|
deltas = deltas.cpu()
|
|
for idx, name in enumerate(names):
|
|
if name not in _cumulative:
|
|
_cumulative[name] = torch.zeros([_num_moments], dtype=_counter_dtype)
|
|
_cumulative[name].add_(deltas[idx])
|
|
|
|
# Return name-value pairs.
|
|
return [(name, _cumulative[name]) for name in names]
|
|
|
|
#----------------------------------------------------------------------------
|