dmlc--dgl
Minjie Wang
5dd35580f7
* improve performance of sample_neighbors * some more improve * test script * benchmarks * multi process * update more tests * WIP * adding two API for state saving * add create from state * upd test * missing file * wip: pickle/unpickle * more c apis * find the problem of empty data array * add null array; pickling speed is bad * still bad perf * still bad perf * wip * fix the pickle speed test; now everything looks good * minor fix * bugfix * some lint fix * address comments * more fix * fix lint * add utest for random.choice * add utest for dgl.rand_graph * fix cpp utests * try fix ci * fix bug in TF backend * upd choice docstring * address comments * upd * try fix compile * add comment
92 行
2.7 KiB
Python
92 行
2.7 KiB
Python
"""Python interfaces to DGL random number generators."""
|
|
import numpy as np
|
|
|
|
from ._ffi.function import _init_api
|
|
from . import backend as F
|
|
from . import ndarray as nd
|
|
|
|
def seed(val):
|
|
"""Set the seed of randomized methods in DGL.
|
|
|
|
The randomized methods include various samplers and random walk routines.
|
|
|
|
Parameters
|
|
----------
|
|
val : int
|
|
The seed
|
|
"""
|
|
_CAPI_SetSeed(val)
|
|
|
|
def choice(a, size, replace=True, prob=None): # pylint: disable=invalid-name
|
|
"""An equivalent to :func:`numpy.random.choice`.
|
|
|
|
Use this function if you:
|
|
|
|
* Perform a non-uniform sampling (probability tensor is given).
|
|
* Sample a small set from a very large population (ratio <5%) uniformly
|
|
*without* replacement.
|
|
* Have a backend tensor on hand and does not want to convert it to numpy
|
|
back and forth.
|
|
|
|
Compared to :func:`numpy.random.choice`, it is slower when replace is True
|
|
and is comparable when replace is False. It wins when the population is
|
|
very large and the number of draws are quite small (e.g., draw <5%). The
|
|
reasons are two folds:
|
|
|
|
* When ``a`` is a large integer, it avoids creating a large range array as
|
|
numpy does.
|
|
* When draw ratio is small, it switches to a hashmap based implementation.
|
|
|
|
It out-performs numpy for non-uniform sampling in general cases.
|
|
|
|
TODO(minjie): support RNG as one of the arguments.
|
|
|
|
Parameters
|
|
----------
|
|
a : 1-D tensor or int
|
|
If an ndarray, a random sample is generated from its elements. If an int,
|
|
the random sample is generated as if a were F.arange(a)
|
|
size : int or tuple of ints
|
|
Output shape. E.g., for size ``(m, n, k)``, then ``m * n * k`` samples are drawn.
|
|
replace : bool, optional
|
|
If true, sample with replacement.
|
|
prob : 1-D tensor, optional
|
|
The probabilities associated with each entry in a.
|
|
If not given the sample assumes a uniform distribution over all entries in a.
|
|
|
|
Returns
|
|
-------
|
|
samples : 1-D tensor
|
|
The generated random samples
|
|
"""
|
|
if isinstance(size, tuple):
|
|
num = np.prod(size)
|
|
else:
|
|
num = size
|
|
|
|
if F.is_tensor(a):
|
|
population = F.shape(a)[0]
|
|
else:
|
|
population = a
|
|
|
|
if prob is None:
|
|
prob = nd.null()
|
|
else:
|
|
prob = F.zerocopy_to_dgl_ndarray(prob)
|
|
|
|
bits = 64 # index array is in 64-bit
|
|
chosen_idx = _CAPI_Choice(int(num), int(population), prob, bool(replace), bits)
|
|
chosen_idx = F.zerocopy_from_dgl_ndarray(chosen_idx)
|
|
|
|
if F.is_tensor(a):
|
|
chosen = F.gather_row(a, chosen_idx)
|
|
else:
|
|
chosen = chosen_idx
|
|
|
|
if isinstance(size, tuple):
|
|
return F.reshape(chosen, size)
|
|
else:
|
|
return chosen
|
|
|
|
_init_api('dgl.rng', __name__)
|