EasyEdit / easyeditor /util /runningstats.py
ZJUPeng's picture
initial commit
3494c6b
raw
history blame
65.3 kB
"""
To use a runningstats object,
1. Create the the desired stat object, e.g., `m = Mean()`
2. Feed it batches via the add method, e.g., `m.add(batch)`
3. Repeat step 2 any number of times.
4. Read out the statistic of interest, e.g., `m.mean()`
Built-in runningstats objects include:
Mean - produces mean().
Variance - mean() and variance() and stdev().
Covariance - mean(), covariance(), correlation(), variance(), stdev().
SecondMoment - moment() is the non-mean-centered covariance, E[x x^T].
Quantile - quantile(), min(), max(), median(), mean(), variance(), stdev().
TopK - topk() returns (values, indexes).
Bincount - bincount() histograms nonnegative integer data.
IoU - intersection(), union(), iou() tally binary co-occurrences.
History - history() returns concatenation of data.
CrossCovariance - covariance between two signals, without self-covariance.
CrossIoU - iou between two signals, without self-IoU.
CombinedStat - aggregates any set of stats.
Add more running stats by subclassing the Stat class.
These statistics are vectorized along dim>=1, so stat.add()
should supply a two-dimensional input where the zeroth
dimension is the batch/sampling dimension and the first
dimension is the feature dimension.
The data type and device used matches the data passed to add();
for example, for higher-precision covariances, convert to double
before calling add().
It is common to want to compute and remember a statistic sampled
over a Dataset, computed in batches, possibly caching the computed
statistic in a file. The tally(stat, dataset, cache) handles
this pattern. It takes a statistic, a dataset, and a cache filename
and sets up a data loader that can be run (or not, if cached) to
compute the statistic, adopting the convention that cached stats are
saved to and loaded from numpy npz files.
"""
import math
import os
import random
import struct
import numpy
import torch
from torch.utils.data.sampler import Sampler
def tally(stat, dataset, cache=None, quiet=False, **kwargs):
"""
To use tally, write code like the following.
stat = Mean()
ds = MyDataset()
for batch in tally(stat, ds, cache='mymean.npz', batch_size=50):
stat.add(batch)
mean = stat.mean()
The first argument should be the Stat being computed. After the
loader is exhausted, tally will bring this stat to the cpu and
cache it (if a cache is specified).
The dataset can be a torch Dataset or a plain Tensor, or it can
be a callable that returns one of those.
Details on caching via the cache= argument:
If the given filename cannot be loaded, tally will leave the
statistic object empty and set up a DataLoader object so that
the loop can be run. After the last iteration of the loop, the
completed statistic will be moved to the cpu device and also
saved in the cache file.
If the cached statistic can be loaded from the given file, tally
will not set up the data loader and instead will return a fully
loaded statistic object (on the cpu device) and an empty list as
the loader.
The `with cache_load_enabled(False):` context manager can
be used to disable loading from the cache.
If needed, a DataLoader will be created to wrap the dataset:
Keyword arguments of tally are passed to the DataLoader,
so batch_size, num_workers, pin_memory, etc. can be specified.
Subsampling is supported via sample_size= and random_sample=:
If sample_size=N is specified, rather than loading the whole
dataset, only the first N items are sampled. If additionally
random_sample=S is specified, the pseudorandom seed S will be
used to select a fixed psedorandom sample of size N to sample.
"""
assert isinstance(stat, Stat)
args = {}
for k in ["sample_size"]:
if k in kwargs:
args[k] = kwargs[k]
cached_state = load_cached_state(cache, args, quiet=quiet)
if cached_state is not None:
stat.load_state_dict(cached_state)
def empty_loader():
return
yield
return empty_loader()
loader = make_loader(dataset, **kwargs)
def wrapped_loader():
yield from loader
stat.to_(device="cpu")
if cache is not None:
save_cached_state(cache, stat, args)
return wrapped_loader()
class cache_load_enabled:
"""
When used as a context manager, cache_load_enabled(False) will prevent
tally from loading cached statsitics, forcing them to be recomputed.
"""
def __init__(self, enabled=True):
self.prev = False
self.enabled = enabled
def __enter__(self):
global global_load_cache_enabled
self.prev = global_load_cache_enabled
global_load_cache_enabled = self.enabled
def __exit__(self, exc_type, exc_value, traceback):
global global_load_cache_enabled
global_load_cache_enabled = self.prev
class Stat:
"""
Abstract base class for a running pytorch statistic.
"""
def __init__(self, state):
"""
By convention, all Stat subclasses can be initialized by passing
state=; and then they will initialize by calling load_state_dict.
"""
self.load_state_dict(resolve_state_dict(state))
def add(self, x, *args, **kwargs):
"""
Observes a batch of samples to be incorporated into the statistic.
Dimension 0 should be the batch dimension, and dimension 1 should
be the feature dimension of the pytorch tensor x.
"""
pass
def load_state_dict(self, d):
"""
Loads this Stat from a dictionary of numpy arrays as saved
by state_dict.
"""
pass
def state_dict(self):
"""
Saves this Stat as a dictionary of numpy arrays that can be
stored in an npz or reloaded later using load_state_dict.
"""
return {}
def save(self, filename):
"""
Saves this stat as an npz file containing the state_dict.
"""
save_cached_state(filename, self, {})
def load(self, filename):
"""
Loads this stat from an npz file containing a saved state_dict.
"""
self.load_state_dict(load_cached_state(filename, {}, quiet=True, throw=True))
def to_(self, device):
"""
Moves this Stat to the given device.
"""
pass
def cpu_(self):
"""
Moves this Stat to the cpu device.
"""
self.to_("cpu")
def cuda_(self):
"""
Moves this Stat to the default cuda device.
"""
self.to_("cuda")
def _normalize_add_shape(self, x, attr="data_shape"):
"""
Flattens input data to 2d.
"""
if not torch.is_tensor(x):
x = torch.tensor(x)
if len(x.shape) < 1:
x = x.view(-1)
data_shape = getattr(self, attr, None)
if data_shape is None:
data_shape = x.shape[1:]
setattr(self, attr, data_shape)
else:
assert x.shape[1:] == data_shape
return x.view(x.shape[0], int(numpy.prod(data_shape)))
def _restore_result_shape(self, x, attr="data_shape"):
"""
Restores output data to input data shape.
"""
data_shape = getattr(self, attr, None)
if data_shape is None:
return x
return x.view(data_shape * len(x.shape))
class Mean(Stat):
"""
Running mean.
"""
def __init__(self, state=None):
if state is not None:
return super().__init__(state)
self.count = 0
self.batchcount = 0
self._mean = None
self.data_shape = None
def add(self, a):
a = self._normalize_add_shape(a)
if len(a) == 0:
return
batch_count = a.shape[0]
batch_mean = a.sum(0) / batch_count
self.batchcount += 1
# Initial batch.
if self._mean is None:
self.count = batch_count
self._mean = batch_mean
return
# Update a batch using Chan-style update for numerical stability.
self.count += batch_count
new_frac = float(batch_count) / self.count
# Update the mean according to the batch deviation from the old mean.
delta = batch_mean.sub_(self._mean).mul_(new_frac)
self._mean.add_(delta)
def size(self):
return self.count
def mean(self):
return self._restore_result_shape(self._mean)
def to_(self, device):
if self._mean is not None:
self._mean = self._mean.to(device)
def load_state_dict(self, state):
self.count = state["count"]
self.batchcount = state["batchcount"]
self._mean = torch.from_numpy(state["mean"])
self.data_shape = (
None if state["data_shape"] is None else tuple(state["data_shape"])
)
def state_dict(self):
return dict(
constructor=self.__module__ + "." + self.__class__.__name__ + "()",
count=self.count,
data_shape=self.data_shape and tuple(self.data_shape),
batchcount=self.batchcount,
mean=self._mean.cpu().numpy(),
)
class NormMean(Mean):
"""
Running average of the norm of input vectors
"""
def __init__(self, state=None):
super().__init__(state)
def add(self, a):
super().add(a.norm(dim=-1))
class Variance(Stat):
"""
Running computation of mean and variance. Use this when you just need
basic stats without covariance.
"""
def __init__(self, state=None):
if state is not None:
return super().__init__(state)
self.count = 0
self.batchcount = 0
self._mean = None
self.v_cmom2 = None
self.data_shape = None
def add(self, a):
a = self._normalize_add_shape(a)
if len(a) == 0:
return
batch_count = a.shape[0]
batch_mean = a.sum(0) / batch_count
centered = a - batch_mean
self.batchcount += 1
# Initial batch.
if self._mean is None:
self.count = batch_count
self._mean = batch_mean
self.v_cmom2 = centered.pow(2).sum(0)
return
# Update a batch using Chan-style update for numerical stability.
oldcount = self.count
self.count += batch_count
new_frac = float(batch_count) / self.count
# Update the mean according to the batch deviation from the old mean.
delta = batch_mean.sub_(self._mean).mul_(new_frac)
self._mean.add_(delta)
# Update the variance using the batch deviation
self.v_cmom2.add_(centered.pow(2).sum(0))
self.v_cmom2.add_(delta.pow_(2).mul_(new_frac * oldcount))
def size(self):
return self.count
def mean(self):
return self._restore_result_shape(self._mean)
def variance(self, unbiased=True):
return self._restore_result_shape(
self.v_cmom2 / (self.count - (1 if unbiased else 0))
)
def stdev(self, unbiased=True):
return self.variance(unbiased=unbiased).sqrt()
def to_(self, device):
if self._mean is not None:
self._mean = self._mean.to(device)
if self.v_cmom2 is not None:
self.v_cmom2 = self.v_cmom2.to(device)
def load_state_dict(self, state):
self.count = state["count"]
self.batchcount = state["batchcount"]
self._mean = torch.from_numpy(state["mean"])
self.v_cmom2 = torch.from_numpy(state["cmom2"])
self.data_shape = (
None if state["data_shape"] is None else tuple(state["data_shape"])
)
def state_dict(self):
return dict(
constructor=self.__module__ + "." + self.__class__.__name__ + "()",
count=self.count,
data_shape=self.data_shape and tuple(self.data_shape),
batchcount=self.batchcount,
mean=self._mean.cpu().numpy(),
cmom2=self.v_cmom2.cpu().numpy(),
)
class Covariance(Stat):
"""
Running computation. Use this when the entire covariance matrix is needed,
and when the whole covariance matrix fits in the GPU.
Chan-style numerically stable update of mean and full covariance matrix.
Chan, Golub. LeVeque. 1983. http://www.jstor.org/stable/2683386
"""
def __init__(self, state=None):
if state is not None:
return super().__init__(state)
self.count = 0
self._mean = None
self.cmom2 = None
self.data_shape = None
def add(self, a):
a = self._normalize_add_shape(a)
if len(a) == 0:
return
batch_count = a.shape[0]
# Initial batch.
if self._mean is None:
self.count = batch_count
self._mean = a.sum(0) / batch_count
centered = a - self._mean
self.cmom2 = centered.t().mm(centered)
return
# Update a batch using Chan-style update for numerical stability.
self.count += batch_count
# Update the mean according to the batch deviation from the old mean.
delta = a - self._mean
self._mean.add_(delta.sum(0) / self.count)
delta2 = a - self._mean
# Update the variance using the batch deviation
self.cmom2.addmm_(mat1=delta.t(), mat2=delta2)
def to_(self, device):
if self._mean is not None:
self._mean = self._mean.to(device)
if self.cmom2 is not None:
self.cmom2 = self.cmom2.to(device)
def mean(self):
return self._restore_result_shape(self._mean)
def covariance(self, unbiased=True):
return self._restore_result_shape(
self.cmom2 / (self.count - (1 if unbiased else 0))
)
def correlation(self, unbiased=True):
cov = self.cmom2 / (self.count - (1 if unbiased else 0))
rstdev = cov.diag().sqrt().reciprocal()
return self._restore_result_shape(rstdev[:, None] * cov * rstdev[None, :])
def variance(self, unbiased=True):
return self._restore_result_shape(
self.cmom2.diag() / (self.count - (1 if unbiased else 0))
)
def stdev(self, unbiased=True):
return self.variance(unbiased=unbiased).sqrt()
def state_dict(self):
return dict(
constructor=self.__module__ + "." + self.__class__.__name__ + "()",
count=self.count,
data_shape=self.data_shape and tuple(self.data_shape),
mean=self._mean.cpu().numpy(),
cmom2=self.cmom2.cpu().numpy(),
)
def load_state_dict(self, state):
self.count = state["count"]
self._mean = torch.from_numpy(state["mean"])
self.cmom2 = torch.from_numpy(state["cmom2"])
self.data_shape = (
None if state["data_shape"] is None else tuple(state["data_shape"])
)
class SecondMoment(Stat):
"""
Running computation. Use this when the entire non-centered 2nd-moment
'covariance-like' matrix is needed, and when the whole matrix fits
in the GPU.
"""
def __init__(self, split_batch=True, state=None):
if state is not None:
return super().__init__(state)
self.count = 0
self.mom2 = None
self.split_batch = split_batch
def add(self, a):
a = self._normalize_add_shape(a)
if len(a) == 0:
return
# Initial batch reveals the shape of the data.
if self.count == 0:
self.mom2 = a.new(a.shape[1], a.shape[1]).zero_()
batch_count = a.shape[0]
# Update the covariance using the batch deviation
self.count += batch_count
self.mom2 += a.t().mm(a)
def to_(self, device):
if self.mom2 is not None:
self.mom2 = self.mom2.to(device)
def moment(self):
return self.mom2 / self.count
def state_dict(self):
return dict(
constructor=self.__module__ + "." + self.__class__.__name__ + "()",
count=self.count,
mom2=self.mom2.cpu().numpy(),
)
def load_state_dict(self, state):
self.count = int(state["count"])
self.mom2 = torch.from_numpy(state["mom2"])
class Bincount(Stat):
"""
Running bincount. The counted array should be an integer type with
non-negative integers.
"""
def __init__(self, state=None):
if state is not None:
return super().__init__(state)
self.count = 0
self._bincount = None
def add(self, a, size=None):
a = a.view(-1)
bincount = a.bincount()
if self._bincount is None:
self._bincount = bincount
elif len(self._bincount) < len(bincount):
bincount[: len(self._bincount)] += self._bincount
self._bincount = bincount
else:
self._bincount[: len(bincount)] += bincount
if size is None:
self.count += len(a)
else:
self.count += size
def to_(self, device):
self._bincount = self._bincount.to(device)
def size(self):
return self.count
def bincount(self):
return self._bincount
def state_dict(self):
return dict(
constructor=self.__module__ + "." + self.__class__.__name__ + "()",
count=self.count,
bincount=self._bincount.cpu().numpy(),
)
def load_state_dict(self, dic):
self.count = int(dic["count"])
self._bincount = torch.from_numpy(dic["bincount"])
class CrossCovariance(Stat):
"""
Covariance. Use this when an off-diagonal block of the covariance
matrix is needed (e.g., when the whole covariance matrix does
not fit in the GPU, this could use a quarter of the memory).
Chan-style numerically stable update of mean and full covariance matrix.
Chan, Golub. LeVeque. 1983. http://www.jstor.org/stable/2683386
"""
def __init__(self, split_batch=True, state=None):
if state is not None:
return super().__init__(state)
self.count = 0
self._mean = None
self.cmom2 = None
self.v_cmom2 = None
self.split_batch = split_batch
def add(self, a, b):
if len(a.shape) == 1:
a = a[None, :]
b = b[None, :]
assert a.shape[0] == b.shape[0]
if len(a.shape) > 2:
a, b = [
d.view(d.shape[0], d.shape[1], -1)
.permute(0, 2, 1)
.reshape(-1, d.shape[1])
for d in [a, b]
]
batch_count = a.shape[0]
# Initial batch.
if self._mean is None:
self.count = batch_count
self._mean = [d.sum(0) / batch_count for d in [a, b]]
centered = [d - bm for d, bm in zip([a, b], self._mean)]
self.v_cmom2 = [c.pow(2).sum(0) for c in centered]
self.cmom2 = centered[0].t().mm(centered[1])
return
# Update a batch using Chan-style update for numerical stability.
self.count += batch_count
# Update the mean according to the batch deviation from the old mean.
delta = [(d - bm) for d, bm in zip([a, b], self._mean)]
for m, d in zip(self._mean, delta):
m.add_(d.sum(0) / self.count)
delta2 = [(d - bm) for d, bm in zip([a, b], self._mean)]
# Update the cross-covariance using the batch deviation
self.cmom2.addmm_(mat1=delta[0].t(), mat2=delta2[1])
# Update the variance using the batch deviation
for vc2, d, d2 in zip(self.v_cmom2, delta, delta2):
vc2.add_((d * d2).sum(0))
def mean(self):
return self._mean
def variance(self, unbiased=True):
return [vc2 / (self.count - (1 if unbiased else 0)) for vc2 in self.v_cmom2]
def stdev(self, unbiased=True):
return [v.sqrt() for v in self.variance(unbiased=unbiased)]
def covariance(self, unbiased=True):
return self.cmom2 / (self.count - (1 if unbiased else 0))
def correlation(self):
covariance = self.covariance(unbiased=False)
rstdev = [s.reciprocal() for s in self.stdev(unbiased=False)]
cor = rstdev[0][:, None] * covariance * rstdev[1][None, :]
# Remove NaNs
cor[torch.isnan(cor)] = 0
return cor
def to_(self, device):
self._mean = [m.to(device) for m in self._mean]
self.v_cmom2 = [vcs.to(device) for vcs in self.v_cmom2]
self.cmom2 = self.cmom2.to(device)
def state_dict(self):
return dict(
constructor=self.__module__ + "." + self.__class__.__name__ + "()",
count=self.count,
mean_a=self._mean[0].cpu().numpy(),
mean_b=self._mean[1].cpu().numpy(),
cmom2_a=self.v_cmom2[0].cpu().numpy(),
cmom2_b=self.v_cmom2[1].cpu().numpy(),
cmom2=self.cmom2.cpu().numpy(),
)
def load_state_dict(self, state):
self.count = int(state["count"])
self._mean = [torch.from_numpy(state[f"mean_{k}"]) for k in "ab"]
self.v_cmom2 = [torch.from_numpy(state[f"cmom2_{k}"]) for k in "ab"]
self.cmom2 = torch.from_numpy(state["cmom2"])
def _float_from_bool(a):
"""
Since pytorch only supports matrix multiplication on float,
IoU computations are done using floating point types.
This function binarizes the input (positive to True and
nonpositive to False), and converts from bool to float.
If the data is already a floating-point type, it leaves
it keeps the same type; otherwise it uses float.
"""
if a.dtype == torch.bool:
return a.float()
if a.dtype.is_floating_point:
return a.sign().clamp_(0)
return (a > 0).float()
class IoU(Stat):
"""
Running computation of intersections and unions of all features.
"""
def __init__(self, state=None):
if state is not None:
return super().__init__(state)
self.count = 0
self._intersection = None
def add(self, a):
assert len(a.shape) == 2
a = _float_from_bool(a)
if self._intersection is None:
self._intersection = torch.mm(a.t(), a)
else:
self._intersection.addmm_(a.t(), a)
self.count += len(a)
def size(self):
return self.count
def intersection(self):
return self._intersection
def union(self):
total = self._intersection.diagonal(0)
return total[:, None] + total[None, :] - self._intersection
def iou(self):
return self.intersection() / (self.union() + 1e-20)
def to_(self, _device):
self._intersection = self._intersection.to(_device)
def state_dict(self):
return dict(
constructor=self.__module__ + "." + self.__class__.__name__ + "()",
count=self.count,
intersection=self._intersection.cpu().numpy(),
)
def load_state_dict(self, state):
self.count = int(state["count"])
self._intersection = torch.tensor(state["intersection"])
class CrossIoU(Stat):
"""
Running computation of intersections and unions of two binary vectors.
"""
def __init__(self, state=None):
if state is not None:
return super().__init__(state)
self.count = 0
self._intersection = None
self.total_a = None
self.total_b = None
def add(self, a, b):
assert len(a.shape) == 2 and len(b.shape) == 2
assert len(a) == len(b), f"{len(a)} vs {len(b)}"
a = _float_from_bool(a) # CUDA only supports mm on float...
b = _float_from_bool(b) # otherwise we would use integers.
intersection = torch.mm(a.t(), b)
asum = a.sum(0)
bsum = b.sum(0)
if self._intersection is None:
self._intersection = intersection
self.total_a = asum
self.total_b = bsum
else:
self._intersection += intersection
self.total_a += asum
self.total_b += bsum
self.count += len(a)
def size(self):
return self.count
def intersection(self):
return self._intersection
def union(self):
return self.total_a[:, None] + self.total_b[None, :] - self._intersection
def iou(self):
return self.intersection() / (self.union() + 1e-20)
def to_(self, _device):
self.total_a = self.total_a.to(_device)
self.total_b = self.total_b.to(_device)
self._intersection = self._intersection.to(_device)
def state_dict(self):
return dict(
constructor=self.__module__ + "." + self.__class__.__name__ + "()",
count=self.count,
total_a=self.total_a.cpu().numpy(),
total_b=self.total_b.cpu().numpy(),
intersection=self._intersection.cpu().numpy(),
)
def load_state_dict(self, state):
self.count = int(state["count"])
self.total_a = torch.tensor(state["total_a"])
self.total_b = torch.tensor(state["total_b"])
self._intersection = torch.tensor(state["intersection"])
class Quantile(Stat):
"""
Streaming randomized quantile computation for torch.
Add any amount of data repeatedly via add(data). At any time,
quantile estimates be read out using quantile(q).
Implemented as a sorted sample that retains at least r samples
(by default r = 3072); the number of retained samples will grow to
a finite ceiling as the data is accumulated. Accuracy scales according
to r: the default is to set resolution to be accurate to better than about
0.1%, while limiting storage to about 50,000 samples.
Good for computing quantiles of huge data without using much memory.
Works well on arbitrary data with probability near 1.
Based on the optimal KLL quantile algorithm by Karnin, Lang, and Liberty
from FOCS 2016. http://ieee-focs.org/FOCS-2016-Papers/3933a071.pdf
"""
def __init__(self, r=3 * 1024, buffersize=None, seed=None, state=None):
if state is not None:
return super().__init__(state)
self.depth = None
self.dtype = None
self.device = None
resolution = r * 2 # sample array is at least half full before discard
self.resolution = resolution
# Default buffersize: 128 samples (and smaller than resolution).
if buffersize is None:
buffersize = min(128, (resolution + 7) // 8)
self.buffersize = buffersize
self.samplerate = 1.0
self.data = None
self.firstfree = [0]
self.randbits = torch.ByteTensor(resolution)
self.currentbit = len(self.randbits) - 1
self.extremes = None
self.count = 0
self.batchcount = 0
def size(self):
return self.count
def _lazy_init(self, incoming):
self.depth = incoming.shape[1]
self.dtype = incoming.dtype
self.device = incoming.device
self.data = [
torch.zeros(
self.depth, self.resolution, dtype=self.dtype, device=self.device
)
]
self.extremes = torch.zeros(self.depth, 2, dtype=self.dtype, device=self.device)
self.extremes[:, 0] = float("inf")
self.extremes[:, -1] = -float("inf")
def to_(self, device):
"""Switches internal storage to specified device."""
if device != self.device:
old_data = self.data
old_extremes = self.extremes
self.data = [d.to(device) for d in self.data]
self.extremes = self.extremes.to(device)
self.device = self.extremes.device
del old_data
del old_extremes
def add(self, incoming):
if self.depth is None:
self._lazy_init(incoming)
assert len(incoming.shape) == 2
assert incoming.shape[1] == self.depth, (incoming.shape[1], self.depth)
self.count += incoming.shape[0]
self.batchcount += 1
# Convert to a flat torch array.
if self.samplerate >= 1.0:
self._add_every(incoming)
return
# If we are sampling, then subsample a large chunk at a time.
self._scan_extremes(incoming)
chunksize = int(math.ceil(self.buffersize / self.samplerate))
for index in range(0, len(incoming), chunksize):
batch = incoming[index : index + chunksize]
sample = sample_portion(batch, self.samplerate)
if len(sample):
self._add_every(sample)
def _add_every(self, incoming):
supplied = len(incoming)
index = 0
while index < supplied:
ff = self.firstfree[0]
available = self.data[0].shape[1] - ff
if available == 0:
if not self._shift():
# If we shifted by subsampling, then subsample.
incoming = incoming[index:]
if self.samplerate >= 0.5:
# First time sampling - the data source is very large.
self._scan_extremes(incoming)
incoming = sample_portion(incoming, self.samplerate)
index = 0
supplied = len(incoming)
ff = self.firstfree[0]
available = self.data[0].shape[1] - ff
copycount = min(available, supplied - index)
self.data[0][:, ff : ff + copycount] = torch.t(
incoming[index : index + copycount, :]
)
self.firstfree[0] += copycount
index += copycount
def _shift(self):
index = 0
# If remaining space at the current layer is less than half prev
# buffer size (rounding up), then we need to shift it up to ensure
# enough space for future shifting.
while self.data[index].shape[1] - self.firstfree[index] < (
-(-self.data[index - 1].shape[1] // 2) if index else 1
):
if index + 1 >= len(self.data):
return self._expand()
data = self.data[index][:, 0 : self.firstfree[index]]
data = data.sort()[0]
if index == 0 and self.samplerate >= 1.0:
self._update_extremes(data[:, 0], data[:, -1])
offset = self._randbit()
position = self.firstfree[index + 1]
subset = data[:, offset::2]
self.data[index + 1][:, position : position + subset.shape[1]] = subset
self.firstfree[index] = 0
self.firstfree[index + 1] += subset.shape[1]
index += 1
return True
def _scan_extremes(self, incoming):
# When sampling, we need to scan every item still to get extremes
self._update_extremes(
torch.min(incoming, dim=0)[0], torch.max(incoming, dim=0)[0]
)
def _update_extremes(self, minr, maxr):
self.extremes[:, 0] = torch.min(
torch.stack([self.extremes[:, 0], minr]), dim=0
)[0]
self.extremes[:, -1] = torch.max(
torch.stack([self.extremes[:, -1], maxr]), dim=0
)[0]
def _randbit(self):
self.currentbit += 1
if self.currentbit >= len(self.randbits):
self.randbits.random_(to=2)
self.currentbit = 0
return self.randbits[self.currentbit]
def state_dict(self):
state = dict(
constructor=self.__module__ + "." + self.__class__.__name__ + "()",
resolution=self.resolution,
depth=self.depth,
buffersize=self.buffersize,
samplerate=self.samplerate,
sizes=numpy.array([d.shape[1] for d in self.data]),
extremes=self.extremes.cpu().detach().numpy(),
size=self.count,
batchcount=self.batchcount,
)
for i, (d, f) in enumerate(zip(self.data, self.firstfree)):
state[f"data.{i}"] = d.cpu().detach().numpy()[:, :f].T
return state
def load_state_dict(self, state):
self.resolution = int(state["resolution"])
self.randbits = torch.ByteTensor(self.resolution)
self.currentbit = len(self.randbits) - 1
self.depth = int(state["depth"])
self.buffersize = int(state["buffersize"])
self.samplerate = float(state["samplerate"])
firstfree = []
buffers = []
for i, s in enumerate(state["sizes"]):
d = state[f"data.{i}"]
firstfree.append(d.shape[0])
buf = numpy.zeros((d.shape[1], s), dtype=d.dtype)
buf[:, : d.shape[0]] = d.T
buffers.append(torch.from_numpy(buf))
self.firstfree = firstfree
self.data = buffers
self.extremes = torch.from_numpy((state["extremes"]))
self.count = int(state["size"])
self.batchcount = int(state.get("batchcount", 0))
self.dtype = self.extremes.dtype
self.device = self.extremes.device
def min(self):
return self.minmax()[0]
def max(self):
return self.minmax()[-1]
def minmax(self):
if self.firstfree[0]:
self._scan_extremes(self.data[0][:, : self.firstfree[0]].t())
return self.extremes.clone()
def median(self):
return self.quantiles(0.5)
def mean(self):
return self.integrate(lambda x: x) / self.count
def variance(self, unbiased=True):
mean = self.mean()[:, None]
return self.integrate(lambda x: (x - mean).pow(2)) / (
self.count - (1 if unbiased else 0)
)
def stdev(self, unbiased=True):
return self.variance(unbiased=unbiased).sqrt()
def _expand(self):
cap = self._next_capacity()
if cap > 0:
# First, make a new layer of the proper capacity.
self.data.insert(
0, torch.zeros(self.depth, cap, dtype=self.dtype, device=self.device)
)
self.firstfree.insert(0, 0)
else:
# Unless we're so big we are just subsampling.
assert self.firstfree[0] == 0
self.samplerate *= 0.5
for index in range(1, len(self.data)):
# Scan for existing data that needs to be moved down a level.
amount = self.firstfree[index]
if amount == 0:
continue
position = self.firstfree[index - 1]
# Move data down if it would leave enough empty space there
# This is the key invariant: enough empty space to fit half
# of the previous level's buffer size (rounding up)
if self.data[index - 1].shape[1] - (amount + position) >= (
-(-self.data[index - 2].shape[1] // 2) if (index - 1) else 1
):
self.data[index - 1][:, position : position + amount] = self.data[
index
][:, :amount]
self.firstfree[index - 1] += amount
self.firstfree[index] = 0
else:
# Scrunch the data if it would not.
data = self.data[index][:, :amount]
data = data.sort()[0]
if index == 1:
self._update_extremes(data[:, 0], data[:, -1])
offset = self._randbit()
scrunched = data[:, offset::2]
self.data[index][:, : scrunched.shape[1]] = scrunched
self.firstfree[index] = scrunched.shape[1]
return cap > 0
def _next_capacity(self):
cap = int(math.ceil(self.resolution * (0.67 ** len(self.data))))
if cap < 2:
return 0
# Round up to the nearest multiple of 8 for better GPU alignment.
cap = -8 * (-cap // 8)
return max(self.buffersize, cap)
def _weighted_summary(self, sort=True):
if self.firstfree[0]:
self._scan_extremes(self.data[0][:, : self.firstfree[0]].t())
size = sum(self.firstfree)
weights = torch.FloatTensor(size) # Floating point
summary = torch.zeros(self.depth, size, dtype=self.dtype, device=self.device)
index = 0
for level, ff in enumerate(self.firstfree):
if ff == 0:
continue
summary[:, index : index + ff] = self.data[level][:, :ff]
weights[index : index + ff] = 2.0**level
index += ff
assert index == summary.shape[1]
if sort:
summary, order = torch.sort(summary, dim=-1)
weights = weights[order.view(-1).cpu()].view(order.shape)
summary = torch.cat(
[self.extremes[:, :1], summary, self.extremes[:, 1:]], dim=-1
)
weights = torch.cat(
[
torch.zeros(weights.shape[0], 1),
weights,
torch.zeros(weights.shape[0], 1),
],
dim=-1,
)
return (summary, weights)
def quantiles(self, quantiles):
if not hasattr(quantiles, "cpu"):
quantiles = torch.tensor(quantiles)
qshape = quantiles.shape
if self.count == 0:
return torch.full((self.depth,) + qshape, torch.nan)
summary, weights = self._weighted_summary()
cumweights = torch.cumsum(weights, dim=-1) - weights / 2
cumweights /= torch.sum(weights, dim=-1, keepdim=True)
result = torch.zeros(
self.depth, quantiles.numel(), dtype=self.dtype, device=self.device
)
# numpy is needed for interpolation
nq = quantiles.view(-1).cpu().detach().numpy()
ncw = cumweights.cpu().detach().numpy()
nsm = summary.cpu().detach().numpy()
for d in range(self.depth):
result[d] = torch.tensor(
numpy.interp(nq, ncw[d], nsm[d]), dtype=self.dtype, device=self.device
)
return result.view((self.depth,) + qshape)
def integrate(self, fun):
result = []
for level, ff in enumerate(self.firstfree):
if ff == 0:
continue
result.append(
torch.sum(fun(self.data[level][:, :ff]) * (2.0**level), dim=-1)
)
if len(result) == 0:
return None
return torch.stack(result).sum(dim=0) / self.samplerate
def readout(self, count=1001):
return self.quantiles(torch.linspace(0.0, 1.0, count))
def normalize(self, data):
"""
Given input data as taken from the training distirbution,
normalizes every channel to reflect quantile values,
uniformly distributed, within [0, 1].
"""
assert self.count > 0
assert data.shape[0] == self.depth
summary, weights = self._weighted_summary()
cumweights = torch.cumsum(weights, dim=-1) - weights / 2
cumweights /= torch.sum(weights, dim=-1, keepdim=True)
result = torch.zeros_like(data).float()
# numpy is needed for interpolation
ndata = data.cpu().numpy().reshape((data.shape[0], -1))
ncw = cumweights.cpu().numpy()
nsm = summary.cpu().numpy()
for d in range(self.depth):
normed = torch.tensor(
numpy.interp(ndata[d], nsm[d], ncw[d]),
dtype=torch.float,
device=data.device,
).clamp_(0.0, 1.0)
if len(data.shape) > 1:
normed = normed.view(*(data.shape[1:]))
result[d] = normed
return result
def sample_portion(vec, p=0.5):
"""
Subsamples a fraction (given by p) of the given batch. Used by
Quantile when the data gets very very large.
"""
bits = torch.bernoulli(
torch.zeros(vec.shape[0], dtype=torch.uint8, device=vec.device), p
)
return vec[bits]
class TopK:
"""
A class to keep a running tally of the the top k values (and indexes)
of any number of torch feature components. Will work on the GPU if
the data is on the GPU. Tracks largest by default, but tracks smallest
if largest=False is passed.
This version flattens all arrays to avoid crashes.
"""
def __init__(self, k=100, largest=True, state=None):
if state is not None:
return super().__init__(state)
self.k = k
self.count = 0
# This version flattens all data internally to 2-d tensors,
# to avoid crashes with the current pytorch topk implementation.
# The data is puffed back out to arbitrary tensor shapes on output.
self.data_shape = None
self.top_data = None
self.top_index = None
self.next = 0
self.linear_index = 0
self.perm = None
self.largest = largest
def add(self, data, index=None):
"""
Adds a batch of data to be considered for the running top k.
The zeroth dimension enumerates the observations. All other
dimensions enumerate different features.
"""
if self.top_data is None:
# Allocation: allocate a buffer of size 5*k, at least 10, for each.
self.data_shape = data.shape[1:]
feature_size = int(numpy.prod(self.data_shape))
self.top_data = torch.zeros(
feature_size, max(10, self.k * 5), out=data.new()
)
self.top_index = self.top_data.clone().long()
self.linear_index = (
0
if len(data.shape) == 1
else torch.arange(feature_size, out=self.top_index.new()).mul_(
self.top_data.shape[-1]
)[:, None]
)
size = data.shape[0]
sk = min(size, self.k)
if self.top_data.shape[-1] < self.next + sk:
# Compression: if full, keep topk only.
self.top_data[:, : self.k], self.top_index[:, : self.k] = self.topk(
sorted=False, flat=True
)
self.next = self.k
# Pick: copy the top sk of the next batch into the buffer.
# Currently strided topk is slow. So we clone after transpose.
# TODO: remove the clone() if it becomes faster.
cdata = data.reshape(size, numpy.prod(data.shape[1:])).t().clone()
td, ti = cdata.topk(sk, sorted=False, largest=self.largest)
self.top_data[:, self.next : self.next + sk] = td
if index is not None:
ti = index[ti]
else:
ti = ti + self.count
self.top_index[:, self.next : self.next + sk] = ti
self.next += sk
self.count += size
def size(self):
return self.count
def topk(self, sorted=True, flat=False):
"""
Returns top k data items and indexes in each dimension,
with channels in the first dimension and k in the last dimension.
"""
k = min(self.k, self.next)
# bti are top indexes relative to buffer array.
td, bti = self.top_data[:, : self.next].topk(
k, sorted=sorted, largest=self.largest
)
# we want to report top indexes globally, which is ti.
ti = self.top_index.view(-1)[(bti + self.linear_index).view(-1)].view(
*bti.shape
)
if flat:
return td, ti
else:
return (
td.view(*(self.data_shape + (-1,))),
ti.view(*(self.data_shape + (-1,))),
)
def to_(self, device):
if self.top_data is not None:
self.top_data = self.top_data.to(device)
if self.top_index is not None:
self.top_index = self.top_index.to(device)
if isinstance(self.linear_index, torch.Tensor):
self.linear_index = self.linear_index.to(device)
def state_dict(self):
return dict(
constructor=self.__module__ + "." + self.__class__.__name__ + "()",
k=self.k,
count=self.count,
largest=self.largest,
data_shape=self.data_shape and tuple(self.data_shape),
top_data=self.top_data.cpu().detach().numpy(),
top_index=self.top_index.cpu().detach().numpy(),
next=self.next,
linear_index=(
self.linear_index.cpu().numpy()
if isinstance(self.linear_index, torch.Tensor)
else self.linear_index
),
perm=self.perm,
)
def load_state_dict(self, state):
self.k = int(state["k"])
self.count = int(state["count"])
self.largest = bool(state.get("largest", True))
self.data_shape = (
None if state["data_shape"] is None else tuple(state["data_shape"])
)
self.top_data = torch.from_numpy(state["top_data"])
self.top_index = torch.from_numpy(state["top_index"])
self.next = int(state["next"])
self.linear_index = (
torch.from_numpy(state["linear_index"])
if len(state["linear_index"].shape) > 0
else int(state["linear_index"])
)
class History(Stat):
"""
Accumulates the concatenation of all the added data.
"""
def __init__(self, data=None, state=None):
if state is not None:
return super().__init__(state)
self._data = data
self._added = []
def _cat_added(self):
if len(self._added):
self._data = torch.cat(
([self._data] if self._data is not None else []) + self._added
)
self._added = []
def add(self, d):
self._added.append(d)
if len(self._added) > 100:
self._cat_added()
def history(self):
self._cat_added()
return self._data
def load_state_dict(self, state):
data = state["data"]
self._data = None if data is None else torch.from_numpy(data)
self._added = []
def state_dict(self):
self._cat_added()
return dict(
constructor=self.__module__ + "." + self.__class__.__name__ + "()",
data=None if self._data is None else self._data.cpu().numpy(),
)
def to_(self, device):
"""Switches internal storage to specified device."""
self._cat_added()
if self._data is not None:
self._data = self._data.to(device)
class CombinedStat(Stat):
"""
A Stat that bundles together multiple Stat objects.
Convenient for loading and saving a state_dict made up of a
hierarchy of stats, and for use with the tally() function.
Example:
cs = CombinedStat(m=Mean(), q=Quantile())
for [b] in tally(cs, MyDataSet(), cache=fn, batch_size=100):
cs.add(b)
print(cs.m.mean())
print(cs.q.median())
"""
def __init__(self, state=None, **kwargs):
self._objs = kwargs
if state is not None:
return super().__init__(state)
def __getattr__(self, k):
if k in self._objs:
return self._objs[k]
raise AttributeError()
def add(self, d, *args, **kwargs):
for obj in self._objs.values():
obj.add(d, *args, **kwargs)
def load_state_dict(self, state):
for prefix, obj in self._objs.items():
obj.load_state_dict(pull_key_prefix(prefix, state))
def state_dict(self):
result = {}
for prefix, obj in self._objs.items():
result.update(push_key_prefix(prefix, obj.state_dict()))
return result
def to_(self, device):
"""Switches internal storage to specified device."""
for v in self._objs.values():
v.to_(device)
def push_key_prefix(prefix, d):
"""
Returns a dict with the same values as d, but where each key
adds the prefix, followed by a dot.
"""
return {prefix + "." + k: v for k, v in d.items()}
def pull_key_prefix(prefix, d):
"""
Returns a filtered dict of all the items of d that start with
the given key prefix, plus a dot, with that prefix removed.
"""
pd = prefix + "."
lpd = len(pd)
return {k[lpd:]: v for k, v in d.items() if k.startswith(pd)}
# We wish to be able to save None (null) values in numpy npz files,
# yet do so without setting the unsecure 'allow_pickle' flag. To do
# that, we will encode null as a special kind of IEEE 754 NaN value.
# Inspired by https://github.com/zuiderkwast/nanbox/blob/master/nanbox.h
# we follow the same Nanboxing scheme used in JavaScriptCore
# (search for JSCJSValue.h#L435), which encodes null values in NaN
# as the NaN value with hex pattern 0xfff8000000000002.
null_numpy_value = numpy.array(
struct.unpack(">d", struct.pack(">Q", 0xFFF8000000000002))[0], dtype=numpy.float64
)
def is_null_numpy_value(v):
"""
True if v is a 64-bit float numpy scalar NaN matching null_numpy_value.
"""
return (
isinstance(v, numpy.ndarray)
and numpy.ndim(v) == 0
and v.dtype == numpy.float64
and numpy.isnan(v)
and 0xFFF8000000000002 == struct.unpack(">Q", struct.pack(">d", v))[0]
)
def box_numpy_null(d):
"""
Replaces None with null_numpy_value, leaving non-None values unchanged.
Recursively descends into a dictionary replacing None values.
"""
try:
return {k: box_numpy_null(v) for k, v in d.items()}
except Exception:
return null_numpy_value if d is None else d
def unbox_numpy_null(d):
"""
Reverses box_numpy_null, replacing null_numpy_value with None.
Recursively descends into a dictionary replacing None values.
"""
try:
return {k: unbox_numpy_null(v) for k, v in d.items()}
except Exception:
return None if is_null_numpy_value(d) else d
def resolve_state_dict(s):
"""
Resolves a state, which can be a filename or a dict-like object.
"""
if isinstance(s, str):
return unbox_numpy_null(numpy.load(s))
return s
global_load_cache_enabled = True
def load_cached_state(cachefile, args, quiet=False, throw=False):
"""
Resolves a state, which can be a filename or a dict-like object.
"""
if not global_load_cache_enabled or cachefile is None:
return None
try:
if isinstance(cachefile, dict):
dat = cachefile
cachefile = "state" # for printed messages
else:
dat = unbox_numpy_null(numpy.load(cachefile))
for a, v in args.items():
if a not in dat or dat[a] != v:
if not quiet:
print("%s %s changed from %s to %s" % (cachefile, a, dat[a], v))
return None
except (FileNotFoundError, ValueError) as e:
if throw:
raise e
return None
else:
if not quiet:
print("Loading cached %s" % cachefile)
return dat
def save_cached_state(cachefile, obj, args):
"""
Saves the state_dict of the given object in a dict or npz file.
"""
if cachefile is None:
return
dat = obj.state_dict()
for a, v in args.items():
if a in dat:
assert dat[a] == v
dat[a] = v
if isinstance(cachefile, dict):
cachefile.clear()
cachefile.update(dat)
else:
os.makedirs(os.path.dirname(cachefile), exist_ok=True)
numpy.savez(cachefile, **box_numpy_null(dat))
class FixedSubsetSampler(Sampler):
"""Represents a fixed sequence of data set indices.
Subsets can be created by specifying a subset of output indexes.
"""
def __init__(self, samples):
self.samples = samples
def __iter__(self):
return iter(self.samples)
def __len__(self):
return len(self.samples)
def __getitem__(self, key):
return self.samples[key]
def subset(self, new_subset):
return FixedSubsetSampler(self.dereference(new_subset))
def dereference(self, indices):
"""
Translate output sample indices (small numbers indexing the sample)
to input sample indices (larger number indexing the original full set)
"""
return [self.samples[i] for i in indices]
class FixedRandomSubsetSampler(FixedSubsetSampler):
"""Samples a fixed number of samples from the dataset, deterministically.
Arguments:
data_source,
sample_size,
seed (optional)
"""
def __init__(self, data_source, start=None, end=None, seed=1):
rng = random.Random(seed)
shuffled = list(range(len(data_source)))
rng.shuffle(shuffled)
self.data_source = data_source
super(FixedRandomSubsetSampler, self).__init__(shuffled[start:end])
def class_subset(self, class_filter):
"""
Returns only the subset matching the given rule.
"""
if isinstance(class_filter, int):
def rule(d):
return d[1] == class_filter
else:
rule = class_filter
return self.subset(
[i for i, j in enumerate(self.samples) if rule(self.data_source[j])]
)
def make_loader(
dataset, sample_size=None, batch_size=1, sampler=None, random_sample=None, **kwargs
):
"""Utility for creating a dataloader on fixed sample subset."""
import typing
if isinstance(dataset, typing.Callable):
# To support deferred dataset loading, support passing a factory
# that creates the dataset when called.
dataset = dataset()
if isinstance(dataset, torch.Tensor):
# The dataset can be a simple tensor.
dataset = torch.utils.data.TensorDataset(dataset)
if sample_size is not None:
assert sampler is None, "sampler cannot be specified with sample_size"
if sample_size > len(dataset):
print(
"Warning: sample size %d > dataset size %d"
% (sample_size, len(dataset))
)
sample_size = len(dataset)
if random_sample is None:
sampler = FixedSubsetSampler(list(range(sample_size)))
else:
sampler = FixedRandomSubsetSampler(
dataset, seed=random_sample, end=sample_size
)
return torch.utils.data.DataLoader(
dataset, sampler=sampler, batch_size=batch_size, **kwargs
)
# Unit Tests
def _unit_test():
import warnings
warnings.filterwarnings("error")
import argparse
import random
import shutil
import tempfile
import time
parser = argparse.ArgumentParser(description="Test things out")
parser.add_argument("--mode", default="cpu", help="cpu or cuda")
parser.add_argument("--test_size", type=int, default=1000000)
args = parser.parse_args()
testdir = tempfile.mkdtemp()
batch_size = random.randint(500, 1500)
# Test NaNboxing.
assert numpy.isnan(null_numpy_value)
assert is_null_numpy_value(null_numpy_value)
assert not is_null_numpy_value(numpy.nan)
# Test Covariance
goal = torch.tensor(numpy.random.RandomState(1).standard_normal(10 * 10)).view(
10, 10
)
data = (
torch.tensor(numpy.random.RandomState(2).standard_normal(args.test_size * 10))
.view(args.test_size, 10)
.mm(goal)
)
data += torch.randn(1, 10) * 999
dcov = data.t().cov()
dcorr = data.t().corrcoef()
rcov = Covariance()
rcov.add(data) # All one batch
assert (rcov.covariance() - dcov).abs().max() < 1e-16
cs = CombinedStat(cov=Covariance(), xcov=CrossCovariance())
ds = torch.utils.data.TensorDataset(data)
for [a] in tally(cs, ds, batch_size=9876):
cs.cov.add(a)
cs.xcov.add(a[:, :3], a[:, 3:])
assert (data.mean(0) - cs.cov.mean()).abs().max() < 1e-12
assert (dcov - cs.cov.covariance()).abs().max() < 2e-12
assert (dcov[:3, 3:] - cs.xcov.covariance()).abs().max() < 1e-12
assert (dcov.diagonal() - torch.cat(cs.xcov.variance())).abs().max() < 1e-12
assert (dcorr - cs.cov.correlation()).abs().max() < 2e-12
# Test CrossCovariance and CrossIoU
fn = f"{testdir}/cross_cache.npz"
ds = torch.utils.data.TensorDataset(
(
torch.arange(args.test_size)[:, None] % torch.arange(1, 6)[None, :] == 0
).double(),
(
torch.arange(args.test_size)[:, None] % torch.arange(5, 8)[None, :] == 0
).double(),
)
c = CombinedStat(c=CrossCovariance(), iou=CrossIoU())
riou = IoU()
count = 0
for [a, b] in tally(c, ds, cache=fn, batch_size=100):
count += 1
c.add(a, b)
riou.add(torch.cat([a, b], dim=1))
assert count == -(-args.test_size // 100)
cor = c.c.correlation()
iou = c.iou.iou()
assert cor.shape == iou.shape == (5, 3)
assert iou[4, 0] == 1.0
assert abs(iou[0, 2] + (-args.test_size // 7 / float(args.test_size))) < 1e-6
assert abs(cor[4, 0] - 1.0) < 1e-2
assert abs(cor[0, 2] - 0.0) < 1e-6
assert all((riou.iou()[:5, -3:] == iou).view(-1))
assert all(riou.iou().diagonal(0) == 1)
c = CombinedStat(c=CrossCovariance(), iou=CrossIoU())
count = 0
for [a, b] in tally(c, ds, cache=fn, batch_size=10):
count += 1
c.add(a, b)
assert count == 0
assert all((c.c.correlation() == cor).view(-1))
assert all((c.iou.iou() == iou).view(-1))
# Test Concatantaion, Mean, Bincount and tally.
fn = f"{testdir}/series_cache.npz"
count = 0
ds = torch.utils.data.TensorDataset(torch.arange(args.test_size))
c = CombinedStat(s=History(), m=Mean(), b=Bincount())
for [b] in tally(c, ds, cache=fn, batch_size=batch_size):
count += 1
c.add(b)
assert count == -(-args.test_size // batch_size)
assert len(c.s.history()) == args.test_size
assert c.s.history()[-1] == args.test_size - 1
assert all(c.s.history() == ds.tensors[0])
assert all(c.b.bincount() == torch.ones(args.test_size))
assert c.m.mean() == float(args.test_size - 1) / 2.0
c2 = CombinedStat(s=History(), m=Mean(), b=Bincount())
batches = tally(c2, ds, cache=fn)
assert len(c2.s.history()) == args.test_size
assert all(c2.s.history() == c.s.history())
assert all(c2.b.bincount() == torch.ones(args.test_size))
assert c2.m.mean() == c.m.mean()
count = 0
for b in batches:
count += 1
assert count == 0 # Shouldn't do anything when it's cached
# An adverarial case: we keep finding more numbers in the middle
# as the stream goes on.
amount = args.test_size
quantiles = 1000
data = numpy.arange(float(amount))
data[1::2] = data[-1::-2] + (len(data) - 1)
data /= 2
depth = 50
alldata = data[:, None] + (numpy.arange(depth) * amount)[None, :]
actual_sum = torch.FloatTensor(numpy.sum(alldata * alldata, axis=0))
amt = amount // depth
for r in range(depth):
numpy.random.shuffle(alldata[r * amt : r * amt + amt, r])
if args.mode == "cuda":
alldata = torch.cuda.FloatTensor(alldata)
device = torch.device("cuda")
else:
alldata = torch.FloatTensor(alldata)
device = None
starttime = time.time()
cs = CombinedStat(
qc=Quantile(),
m=Mean(),
v=Variance(),
c=Covariance(),
s=SecondMoment(),
t=TopK(),
i=IoU(),
)
# Feed data in little batches
i = 0
while i < len(alldata):
batch_size = numpy.random.randint(1000)
cs.add(alldata[i : i + batch_size])
i += batch_size
# Test state dict
saved = cs.state_dict()
# numpy.savez(f'{testdir}/saved.npz', **box_numpy_null(saved))
# saved = unbox_numpy_null(numpy.load(f'{testdir}/saved.npz'))
cs.save(f"{testdir}/saved.npz")
loaded = unbox_numpy_null(numpy.load(f"{testdir}/saved.npz"))
assert set(loaded.keys()) == set(saved.keys())
# Restore using state=saved in constructor.
cs2 = CombinedStat(
qc=Quantile(),
m=Mean(),
v=Variance(),
c=Covariance(),
s=SecondMoment(),
t=TopK(),
i=IoU(),
state=saved,
)
# saved = unbox_numpy_null(numpy.load(f'{testdir}/saved.npz'))
assert not cs2.qc.device.type == "cuda"
cs2.to_(device)
# alldata = alldata.cpu()
cs2.add(alldata)
actual_sum *= 2
# print(abs(alldata.mean(0) - cs2.m.mean()) / alldata.mean())
assert all(abs(alldata.mean(0) - cs2.m.mean()) / alldata.mean() < 1e-5)
assert all(abs(alldata.mean(0) - cs2.v.mean()) / alldata.mean() < 1e-5)
assert all(abs(alldata.mean(0) - cs2.c.mean()) / alldata.mean() < 1e-5)
# print(abs(alldata.var(0) - cs2.v.variance()) / alldata.var(0))
assert all(abs(alldata.var(0) - cs2.v.variance()) / alldata.var(0) < 1e-3)
assert all(abs(alldata.var(0) - cs2.c.variance()) / alldata.var(0) < 1e-2)
# print(abs(alldata.std(0) - cs2.v.stdev()) / alldata.std(0))
assert all(abs(alldata.std(0) - cs2.v.stdev()) / alldata.std(0) < 1e-4)
# print(abs(alldata.std(0) - cs2.c.stdev()) / alldata.std(0))
assert all(abs(alldata.std(0) - cs2.c.stdev()) / alldata.std(0) < 2e-3)
moment = (alldata.t() @ alldata) / len(alldata)
# print(abs(moment - cs2.s.moment()) / moment.abs())
assert all((abs(moment - cs2.s.moment()) / moment.abs()).view(-1) < 1e-2)
assert all(alldata.max(dim=0)[0] == cs2.t.topk()[0][:, 0])
assert cs2.i.iou()[0, 0] == 1
assert all((cs2.i.iou()[1:, 1:] == 1).view(-1))
assert all(cs2.i.iou()[1:, 0] < 1)
assert all(cs2.i.iou()[1:, 0] == cs2.i.iou()[0, 1:])
# Restore using cs.load() method.
cs = CombinedStat(
qc=Quantile(),
m=Mean(),
v=Variance(),
c=Covariance(),
s=SecondMoment(),
t=TopK(),
i=IoU(),
)
cs.load(f"{testdir}/saved.npz")
assert not cs.qc.device.type == "cuda"
cs.to_(device)
cs.add(alldata)
# actual_sum *= 2
# print(abs(alldata.mean(0) - cs.m.mean()) / alldata.mean())
assert all(abs(alldata.mean(0) - cs.m.mean()) / alldata.mean() < 1e-5)
assert all(abs(alldata.mean(0) - cs.v.mean()) / alldata.mean() < 1e-5)
assert all(abs(alldata.mean(0) - cs.c.mean()) / alldata.mean() < 1e-5)
# print(abs(alldata.var(0) - cs.v.variance()) / alldata.var(0))
assert all(abs(alldata.var(0) - cs.v.variance()) / alldata.var(0) < 1e-3)
assert all(abs(alldata.var(0) - cs.c.variance()) / alldata.var(0) < 1e-2)
# print(abs(alldata.std(0) - cs.v.stdev()) / alldata.std(0))
assert all(abs(alldata.std(0) - cs.v.stdev()) / alldata.std(0) < 1e-4)
# print(abs(alldata.std(0) - cs.c.stdev()) / alldata.std(0))
assert all(abs(alldata.std(0) - cs.c.stdev()) / alldata.std(0) < 2e-3)
moment = (alldata.t() @ alldata) / len(alldata)
# print(abs(moment - cs.s.moment()) / moment.abs())
assert all((abs(moment - cs.s.moment()) / moment.abs()).view(-1) < 1e-2)
assert all(alldata.max(dim=0)[0] == cs.t.topk()[0][:, 0])
assert cs.i.iou()[0, 0] == 1
assert all((cs.i.iou()[1:, 1:] == 1).view(-1))
assert all(cs.i.iou()[1:, 0] < 1)
assert all(cs.i.iou()[1:, 0] == cs.i.iou()[0, 1:])
# Randomized quantile test
qc = cs.qc
ro = qc.readout(1001).cpu()
endtime = time.time()
gt = (
torch.linspace(0, amount, quantiles + 1)[None, :]
+ (torch.arange(qc.depth, dtype=torch.float) * amount)[:, None]
)
maxreldev = torch.max(torch.abs(ro - gt) / amount) * quantiles
print("Randomized quantile test results:")
print("Maximum relative deviation among %d perentiles: %f" % (quantiles, maxreldev))
minerr = torch.max(
torch.abs(
qc.minmax().cpu()[:, 0] - torch.arange(qc.depth, dtype=torch.float) * amount
)
)
maxerr = torch.max(
torch.abs(
(qc.minmax().cpu()[:, -1] + 1)
- (torch.arange(qc.depth, dtype=torch.float) + 1) * amount
)
)
print("Minmax error %f, %f" % (minerr, maxerr))
interr = torch.max(
torch.abs(qc.integrate(lambda x: x * x).cpu() - actual_sum) / actual_sum
)
print("Integral error: %f" % interr)
medianerr = torch.max(
torch.abs(qc.median() - alldata.median(0)[0]) / alldata.median(0)[0]
).cpu()
print("Median error: %f" % medianerr)
meanerr = torch.max(torch.abs(qc.mean() - alldata.mean(0)) / alldata.mean(0)).cpu()
print("Mean error: %f" % meanerr)
varerr = torch.max(torch.abs(qc.variance() - alldata.var(0)) / alldata.var(0)).cpu()
print("Variance error: %f" % varerr)
counterr = (
(qc.integrate(lambda x: torch.ones(x.shape[-1]).cpu()) - qc.size())
/ (0.0 + qc.size())
).item()
print("Count error: %f" % counterr)
print("Time %f" % (endtime - starttime))
# Algorithm is randomized, so some of these will fail with low probability.
assert maxreldev < 1.0
assert minerr == 0.0
assert maxerr == 0.0
assert interr < 0.01
assert abs(counterr) < 0.001
shutil.rmtree(testdir, ignore_errors=True)
print("OK")
if __name__ == "__main__":
_unit_test()