from typing import List, Union
import probflow.utils.ops as O
from probflow.parameters import DeterministicParameter, Parameter
from .module import Module
[docs]class Embedding(Module):
r"""A categorical embedding layer.
Maps an input variable containing non-negative integers to dense vectors.
The length of the vectors (the dimensionality of the embedding) can be set
with the ``dims`` keyword argument. The embedding is learned over the
course of training: if there are N unique integers in the input, and the
embedding dimensionality is M, a matrix of NxM free parameters is created
and optimized to minimize the loss.
By default, a :class:`.Deterministic` distribution is used for the
embedding variables' posterior distributions, with :class:`.Normal`
``(0, 1)`` priors. This corresponds to normal non-probabilistic embedding
with L2 regularization.
The embeddings can be non-probabilistic (each integer corresponds to a
single point in M-dimensional space, the default), or probabilistic (each
integer corresponds to a M-dimensional multivariate distribution). Set the
`probabilistic` kwarg to True to use probabilistic embeddings.
Parameters
----------
k : int > 0 or List[int]
Number of categories to embed.
d : int > 0 or List[int]
Number of embedding dimensions.
posterior : |Distribution| class
Probability distribution class to use to approximate the posterior.
Default = :class:`.Deterministic`
prior : |Distribution| object
Prior probability distribution which has been instantiated
with parameters.
Default = :class:`.Normal` ``(0,1)``
initializer : dict of callables
Initializer functions to use for each variable of the variational
posterior distribution. Keys correspond to variable names (arguments
to the distribution), and values contain functions to initialize those
variables given ``shape`` as the single argument.
probabilistic : bool
Whether variational posteriors for the weights and biases should be
probabilistic. If False (the default), will use
:class:`.Deterministic` distributions for the variational posteriors.
If True, will use :class:`.Normal` distributions.
name : str
Name for this layer.
Default = 'Embeddings'
kwargs
Additional keyword arguments are passed to the :class:`.Parameter`
constructor which creates the embedding variables.
Examples
--------
Embed 10k word IDs into a 50-dimensional space:
.. code-block:: python
emb = Embedding(k=10000, d=50)
ids = tf.random.uniform([1000000], minval=1, maxval=10000,
dtype=tf.dtypes.int64)
embeddings = emb(ids)
TODO: fuller example
"""
def __init__(
self,
k: Union[int, List[int]],
d: Union[int, List[int]],
probabilistic: bool = False,
name: str = "Embedding",
**kwargs
):
# Convert to list if not already
if isinstance(k, int):
k = [k]
if isinstance(d, int):
d = [d]
# Check values
if len(k) != len(d):
raise ValueError("d and k must be the same length")
if any(e < 1 for e in k):
raise ValueError("k must be >0")
if any(e < 1 for e in d):
raise ValueError("d must be >0")
# Override posterior and initializer for probabilistic embedding
if probabilistic:
ParameterClass = Parameter
else:
ParameterClass = DeterministicParameter
# Create the parameters
self.embeddings = [
ParameterClass(
shape=[k[i], d[i]], name=name + "_" + str(i), **kwargs
)
for i in range(len(d))
]
def __call__(self, x):
"""Perform the forward pass"""
embs = [
O.gather(self.embeddings[i](), x[:, i])
for i in range(len(self.embeddings))
]
return O.cat(embs, -1)