def f(log_probs, advantages, old_log_probs, mask):
if reweight: # Use new policy weights for sampled actions instead.
mask *= jnp.exp(fastmath.stop_gradient(log_probs) - old_log_probs)
if sampled_all_discrete: # Actions were sampled uniformly; weight them.
mask *= jnp.exp(old_log_probs)
weights = jnp.minimum(awr_weights(advantages, beta), w_max)
return -jnp.sum(log_probs * weights * mask) / jnp.sum(mask)
def DotProductAttention(query, key, value, mask):
"""Dot product self-attention.
Args:
query (jax.interpreters.xla.DeviceArray): array of query representations with shape (L_q by d)
key (jax.interpreters.xla.DeviceArray): array of key representations with shape (L_k by d)
value (jax.interpreters.xla.DeviceArray): array of value representations with shape (L_k by d) where L_v = L_k
mask (jax.interpreters.xla.DeviceArray): attention-mask, gates attention with shape (L_q by L_k)
Returns:
jax.interpreters.xla.DeviceArray: Self-attention array for q, k, v arrays. (L_q by L_k)
"""
assert query.shape[-1] == key.shape[-1] == value.shape[
-1], "Embedding dimensions of q, k, v aren't all the same"
depth = query.shape[-1]
# Calculate scaled query key dot product according to formula above
dots = jnp.matmul(query, jnp.swapaxes(key, -1, -2)) / jnp.sqrt(depth)
if mask is not None: # The 'None' in this line does not need to be replaced
dots = jnp.where(mask, dots, jnp.full_like(dots, -1e9))
# Softmax formula implementation
logsumexp = trax.fastmath.logsumexp(dots, axis=-1, keepdims=True)
dots = jnp.exp(dots - logsumexp)
attention = jnp.matmul(dots, value)
return attention
def _aggregate_values(self, values, aggregate_max, act_log_probs):
if self._q_value:
if aggregate_max:
values = jnp.max(values, axis=1)
elif self._sample_all_discrete_actions:
values = jnp.sum(values * jnp.exp(act_log_probs), axis=1)
else:
values = jnp.mean(values, axis=1)
return np.array(values) # Move the values to CPU.
def DotProductAttention(queries, keys, values, pos_emb, context_bias,
location_bias, mask, separate_cls, dropout, mode, rng):
"""Computes new activations via masked attention-weighted sum of values.
This function is the core of the attention mechanism. It:
- computes per-head attention weights from per-head `queries` and `keys`,
- applies `mask` to screen out positions that come from padding tokens,
- optionally applies dropout to attention weights, and
- uses attention weights to combine per-head `values` vectors.
Args:
queries: Per-head activations representing attention queries.
keys: Per-head activations representing attention keys.
values: Per-head activations to be combined by computed attention weights.
pos_emb: Per-head activations representing positional embeddings.
context_bias: Global context bias from Transformer XL's attention.
location_bias: Global location bias from Transformer XL's attention.
mask: Mask that distinguishes positions with real content vs. padding.
separate_cls: True/False if we separate_cls in calculations.
dropout: Probabilistic rate for dropout applied to attention strengths
(based on query-key pairs) before applying them to values.
mode: One of `'train'`, `'eval'`, or `'predict'`.
rng: Single-use random number generator (JAX PRNG key).
Returns:
Per-head activations resulting from masked per-head attention-weighted
sum of per-head values.
"""
d_feature = queries.shape[-1]
keys_len, queries_len = keys.shape[-2], queries.shape[-2]
funnel_factor, is_upsampling = calc_funnel_ratio(keys_len, queries_len)
ac = jnp.einsum('bnid,bnjd->bnij', queries + context_bias, keys)
bd = jnp.einsum('bnid,jnd->bnij', queries + location_bias, pos_emb)
if mode != 'predict':
bd = _fast_matrix_shift(bd, funnel_factor, is_upsampling)
if separate_cls:
# Masking out location part of attention for cls token
bd = bd.at[:, :, :, 0].set(0)
bd = bd.at[:, :, 0, :].set(0)
dots = (ac + bd) / jnp.sqrt(d_feature)
if mask is not None:
dots = jnp.where(mask, dots, jnp.full_like(dots, -1e9))
# Softmax.
dots = jnp.exp(dots - fastmath.logsumexp(dots, axis=-1, keepdims=True))
if dropout >= 1.0:
raise ValueError('Dropout rates must be lower than 1.')
if dropout is not None and dropout > 0.0 and mode == 'train':
keep = fastmath.random.bernoulli(rng, 1.0 - dropout, dots.shape)
dots = jnp.where(keep, dots / (1.0 - dropout), jnp.zeros_like(dots))
out = jnp.matmul(dots, values)
out = out.astype(jnp.float32)
dots = dots.astype(jnp.float32)
return out, dots
def _sincos(self, start, length, d_feature):
"""Create the sin-cos tensor of shape [1, length, d_feature]."""
position = jnp.arange(0, length)[:, None] + start
div_term = jnp.exp(
jnp.arange(0, d_feature, 2) * -(jnp.log(10000.0) / d_feature))
sin = jnp.sin(position * div_term)
cos = jnp.cos(position * div_term)
pe = jnp.concatenate([sin, cos], axis=1)
return pe[None, :, :] # [1, length, d_feature]
def _per_head_attention(queries, keys, values, mask, dropout, mode, rng):
"""Computes new per-head activations via scaled dot-product attention.
This function is the core of the attention mechanism. Given per-head
``queries`` (Q), ``keys`` (K), ``values`` (V), and ``mask``, it:
- computes the scaled dot product of each Q-K pair;
- applies ``mask`` to screen out positions that come from padding tokens
(indicated by 0 value);
- [in ``'train'`` mode] applies dropout to Q-K dot products;
- computes Q-K attention strengths using a per-query softmax of the Q-K dot
products; and
- for each query position, combines V vectors according to the Q-K
attention strengths.
Args:
queries: Per-head activations representing attention queries.
keys: Per-head activations representing attention keys.
values: Per-head activations to be combined by computed attention strengths.
mask: Mask that distinguishes positions with real content vs. padding.
dropout: Probababilistic rate for attention dropout, which overrides
(sets to zero) some attention strengths derived from query-key
matching. As a result, on a given forward pass, some value vectors
don't contribute to the output, analogous to how regular dropout can
cause some node activations to be ignored. Applies only in ``'train'``
mode.
mode: One of ``'train'``, ``'eval'``, or ``'predict'``.
rng: Single-use random number generator (JAX PRNG key).
Returns:
Tuple of (activations, attn_strengths), where activations are new per-head
activation vectors and attn_strengths is a matrix of per-head attention
strengths.
"""
if dropout >= 1.0:
raise ValueError(f'Dropout rate ({dropout}) must be lower than 1.')
d_feature = queries.shape[-1]
dots = jnp.matmul(queries, jnp.swapaxes(keys, -1, -2)) / jnp.sqrt(d_feature)
if mask is not None:
dots = jnp.where(mask,
dots,
jnp.full_like(dots, -1e9))
attn_strengths = (
jnp.exp(dots - fastmath.logsumexp(dots, axis=-1, keepdims=True)))
if dropout is not None and dropout > 0.0 and mode == 'train':
keep = fastmath.random.bernoulli(rng, 1.0 - dropout, attn_strengths.shape)
attn_strengths = jnp.where(keep,
attn_strengths / (1.0 - dropout),
jnp.zeros_like(attn_strengths))
activations = jnp.matmul(attn_strengths, values).astype(jnp.float32)
attn_strengths = attn_strengths.astype(jnp.float32)
return activations, attn_strengths
def Softmax(axis=-1):
"""Returns a layer that applies softmax along one tensor axis.
`Softmax` acts on a group of values and normalizes them to look like a set
of probability values. (Probability values must be non-negative, and as a
set must sum to 1.)
Args:
axis: Axis along which values are grouped for computing softmax.
"""
return Fn('Softmax', lambda x: jnp.exp(log_softmax(x, axis=axis)))
def fn(dist_inputs, actions, q_values, act_log_probs, mask):
del dist_inputs, actions, mask
q_values = jnp.swapaxes(q_values, 0, 1)
act_log_probs = jnp.swapaxes(act_log_probs, 0, 1)
if self._sample_all_discrete_actions:
values = jnp.sum(q_values * jnp.exp(act_log_probs), axis=0)
else:
values = jnp.mean(q_values, axis=0)
advantages = q_values - values # Broadcasting values over n_samples
if preprocess:
advantages = self._preprocess_advantages(advantages)
return advantages
def ProbsRatio(dist_inputs, actions, old_log_probs, log_prob_fun):
"""Probability Ratio from the PPO algorithm."""
# dist_inputs of the shape float32[128,1,18]
# actions of the shape int32[128,1]
# and old_log_probs of the shape float32[128,1]
new_log_probs = NewLogProbs(dist_inputs, actions, log_prob_fun)
assert new_log_probs.shape == old_log_probs.shape, (
f'new_log_probs.shape was {new_log_probs.shape} and'
f'old_log_probs.shape was {old_log_probs.shape}')
# The ratio between new_probs and old_probs expressed
# using log_probs and exponentiation
probs_ratio = jnp.exp(new_log_probs - old_log_probs)
return probs_ratio
def f(new_log_probs, advantages, old_log_probs, mask):
# new_log_probs of the shape float32[128,1]
# advantages of the shape int32[128,1]
# old_log_probs of the shape int32[128,1]
# mask of the shape int32[128,1]
if new_log_probs.shape != advantages.shape:
raise ValueError('New log-probs and advantages shapes '
'should be the same, %s != %s' % (new_log_probs.shape,
advantages.shape))
if new_log_probs.shape != old_log_probs.shape:
raise ValueError('New log-probs and old log-probs shapes '
'should be the same, %s != %s' % (new_log_probs.shape,
old_log_probs.shape))
if new_log_probs.shape != mask.shape:
raise ValueError('New log-probs and mask shapes should be the same'
', %s != %s' % (new_log_probs.shape, mask.shape))
# The ratio between new_probs and old_probs expressed
# using log_probs and exponentaion
probs_ratio = jnp.exp(new_log_probs - old_log_probs)
if advantages.shape != probs_ratio.shape:
raise ValueError('New log-probs and old log probs shapes '
'should be the same, %s != %s' % (advantages.shape,
probs_ratio.shape))
unclipped_objective = probs_ratio * advantages
clipped_objective = jnp.clip(probs_ratio,
1 - self._epsilon,
1 + self._epsilon) * advantages
if unclipped_objective.shape != probs_ratio.shape:
raise ValueError('unclipped_objective and clipped_objective shapes '
'should be the same, %s != %s' % (
unclipped_objective.shape,
clipped_objective.shape))
ppo_objective = jnp.minimum(unclipped_objective, clipped_objective)
if ppo_objective.shape != mask.shape:
raise ValueError('ppo_objective and mask shapes '
'should be the same, %s != %s' % (
ppo_objective.shape,
mask.shape))
ppo_loss = -jnp.sum(ppo_objective * mask) / jnp.sum(mask)
entropy_vec = self._policy_dist.entropy(
new_log_probs) * self._entropy_coeff
entropy_loss = jnp.mean(entropy_vec)
combined_loss = ppo_loss - entropy_loss
return combined_loss
def DotProductAttention(queries, keys, values, mask, dropout, mode, rng):
"""Computes new activations via masked attention-weighted sum of values.
This function is the core of the attention mechanism. It:
- computes per-head attention weights from per-head `queries` and `keys`,
- applies `mask` to screen out positions that come from padding tokens,
- optionally applies dropout to attention weights, and
- uses attention weights to combine per-head `values` vectors.
Args:
queries: Per-head activations representing attention queries.
keys: Per-head activations representing attention keys.
values: Per-head activations to be combined by computed attention weights.
mask: Mask that distinguishes positions with real content vs. padding.
dropout: Probababilistic rate for dropout applied to attention strengths
(based on query-key pairs) before applying them to values.
mode: One of `'train'`, `'eval'`, or `'predict'`.
rng: Single-use random number generator (JAX PRNG key).
Returns:
Per-head activations resulting from masked per-head attention-weighted
sum of per-head values.
"""
d_feature = queries.shape[-1]
dots = jnp.matmul(queries, jnp.swapaxes(keys, -1, -2)) / jnp.sqrt(d_feature)
if mask is not None:
# TODO(kitaev): workaround for https://github.com/google/jax/issues/850
# We must ensure that both mask and the -1e9 constant have a data dependency
# on the input. Broadcasted copies of these use a lot of memory, so they
# should be computed at runtime (rather than being global constants).
if fastmath.is_backend(fastmath.Backend.JAX):
mask = jax.lax.tie_in(dots, mask)
# JAX's `full_like` already ties in -1e9 to dots.
dots = jnp.where(mask, dots, jnp.full_like(dots, -1e9))
# Softmax.
dots = jnp.exp(dots - fastmath.logsumexp(dots, axis=-1, keepdims=True))
if dropout >= 1.0:
raise ValueError('Dropout rates must be lower than 1.')
if dropout is not None and dropout > 0.0 and mode == 'train':
keep = fastmath.random.bernoulli(rng, 1.0 - dropout, dots.shape)
dots = jnp.where(keep, dots / (1.0 - dropout), jnp.zeros_like(dots))
out = jnp.matmul(dots, values)
out = out.astype(jnp.float32)
dots = dots.astype(jnp.float32)
return out, dots
def rotate(x):
"""Rotate function."""
_, l, d = x.shape
inv_freq = jnp.exp(jnp.arange(0, d, 2) * -(jnp.log(10000.0) / d))
positions = jnp.arange(l)
freqs = jnp.einsum('i,j->ij', positions, inv_freq)
emb = jnp.concatenate((freqs, freqs), axis=-1)
cos = jnp.cos(emb)
sin = jnp.sin(emb)
def mul(vecs, pos_emb):
return jnp.einsum('bld,ld->bld', vecs, pos_emb)
def rotate_half(x):
x1, x2 = x[..., :x.shape[-1] // 2], x[..., x.shape[-1] // 2:]
return jnp.concatenate((-x2, x1), axis=x1.ndim - 1)
return mul(x, cos) + mul(rotate_half(x), sin)
def _calc_attn_scores(q, k):
ac = jnp.einsum('bnid,bnjd->bnij', q + context_bias, k)
bd = jnp.einsum('bnid,jnd->bnij', q + location_bias, pos_emb)
if mode != 'predict':
bd = _fast_matrix_shift(bd)
dots = (ac + bd) / jnp.sqrt(d_feature)
dots = jnp.where(mask, dots, jnp.full_like(dots, -1e9))
# Softmax.
dots = jnp.exp(dots - fastmath.logsumexp(dots, axis=-1, keepdims=True))
if dropout >= 1.0:
raise ValueError('Dropout rates must be lower than 1.')
if dropout is not None and dropout > 0.0 and mode == 'train':
keep = fastmath.random.bernoulli(rng, 1.0 - dropout, dots.shape)
dots = jnp.where(keep, dots / (1.0 - dropout), jnp.zeros_like(dots))
return dots
def DotProductAttention(queries, keys, values, mask, dropout, mode, rng):
"""Computes new activations via masked attention-weighted sum of values.
This function is the core of the attention mechanism. It:
- computes per-head attention weights from per-head ``queries`` and
``keys``,
- applies ``mask`` to screen out positions that come from padding tokens,
- optionally applies dropout to attention weights, and
- uses attention weights to combine per-head ``values`` vectors.
Args:
queries: Per-head activations representing attention queries.
keys: Per-head activations representing attention keys.
values: Per-head activations to be combined by computed attention weights.
mask: Mask that distinguishes positions with real content vs. padding.
dropout: Probababilistic rate for attention dropout, which overrides
(sets to zero) some attention strengths derived from query-key
matching. As a result, on a given forward pass, some value vectors
don't contribute to the output, analogous to how regular dropout can
cause some node activations to be ignored.
mode: One of ``'train'``, ``'eval'``, or ``'predict'``.
rng: Single-use random number generator (JAX PRNG key).
Returns:
Per-head activations resulting from masked per-head attention-weighted
sum of per-head values.
"""
d_feature = queries.shape[-1]
dots = jnp.matmul(queries, jnp.swapaxes(keys, -1,
-2)) / jnp.sqrt(d_feature)
if mask is not None:
dots = jnp.where(mask, dots, jnp.full_like(dots, -1e9))
# Softmax.
dots = jnp.exp(dots - fastmath.logsumexp(dots, axis=-1, keepdims=True))
if dropout >= 1.0:
raise ValueError('Dropout rates must be lower than 1.')
if dropout is not None and dropout > 0.0 and mode == 'train':
keep = fastmath.random.bernoulli(rng, 1.0 - dropout, dots.shape)
dots = jnp.where(keep, dots / (1.0 - dropout), jnp.zeros_like(dots))
out = jnp.matmul(dots, values)
out = out.astype(jnp.float32)
dots = dots.astype(jnp.float32)
return out, dots
def DotProductAttention(query, key, value, mask):
assert query.shape[-1] == key.shape[-1] == value.shape[-1]
depth = query.shape[-1]
dots = jnp.matmul(query, jnp.swapaxes(key, -1, -2)) / jnp.sqrt(
depth) # Part of dot product formula
# Apply mask
if mask is not None:
dots = jnp.where(mask, dots, jnp.full_like(dots, -1e9))
# Rest of dot product attention formula
logsumexp = trax.fastmath.logsumexp(dots, axis=-1, keepdims=True)
dots = jnp.exp(dots - logsumexp)
attention = jnp.matmul(dots, value)
return attention
def LossInput(dist_inputs, actions, q_values, act_log_probs, mask): # pylint: disable=invalid-name
"""Calculates action log probabilities and normalizes advantages."""
# (batch_size, n_samples, ...) -> (n_samples, batch_size, ...)
q_values = jnp.swapaxes(q_values, 0, 1)
mask = jnp.swapaxes(mask, 0, 1)
actions = jnp.swapaxes(actions, 0, 1)
act_log_probs = jnp.swapaxes(act_log_probs, 0, 1)
# TODO(pkozakowski,lukaszkaiser): Try max here, or reweighting?
if self._sample_all_discrete_actions:
values = jnp.sum(q_values * jnp.exp(act_log_probs), axis=0)
else:
values = jnp.mean(q_values, axis=0)
advantages = q_values - values # Broadcasting values over n_samples
advantages = self._preprocess_advantages(advantages)
# Broadcast inputs and calculate log-probs
dist_inputs = jnp.broadcast_to(
dist_inputs, (self._q_value_n_samples, ) + dist_inputs.shape)
log_probs = self._policy_dist.log_prob(dist_inputs, actions)
return (log_probs, advantages, act_log_probs, mask)
def _aggregate_values(self, values, aggregate, act_log_probs):
# Normalize the Q-values before aggragetion, so it can adapt to the scale
# of the returns. This does not affect mean and max aggregation.
scale = 1
epsilon = 1e-5
if self._q_value_normalization == 'std':
scale = jnp.std(values) + epsilon
elif self._q_value_normalization == 'abs':
scale = jnp.mean(jnp.abs(values - jnp.mean(values))) + epsilon
values /= scale
temp = self._q_value_temperature
if self._q_value:
assert values.shape[:2] == (self._value_batch_size,
self._q_value_n_samples)
if aggregate == 'max':
# max_a Q(s, a)
values = jnp.max(values, axis=1)
elif aggregate == 'softmax':
# sum_a (Q(s, a) * w(s, a))
# where w(s, .) = softmax (Q(s, .) / T)
weights = tl.Softmax(axis=1)(values / temp)
values = jnp.sum(values * weights, axis=1)
elif aggregate == 'logsumexp':
# log(mean_a exp(Q(s, a) / T)) * T
n = values.shape[1]
values = (fastmath.logsumexp(values / temp, axis=1) -
jnp.log(n)) * temp
else:
assert aggregate == 'mean'
# mean_a Q(s, a)
if self._sample_all_discrete_actions:
values = jnp.sum(values * jnp.exp(act_log_probs), axis=1)
else:
values = jnp.mean(values, axis=1)
# Re-scale the Q-values after aggregation.
values *= scale
return np.array(values) # Move the values to CPU.
def DotProductAttention(query, key, value, mask):
"""Dot product self-attention.
Args:
query (jax.interpreters.xla.DeviceArray): array of query representations with shape (L_q by d)
key (jax.interpreters.xla.DeviceArray): array of key representations with shape (L_k by d)
value (jax.interpreters.xla.DeviceArray): array of value representations with shape (L_k by d) where L_v = L_k
mask (jax.interpreters.xla.DeviceArray): attention-mask, gates attention with shape (L_q by L_k)
Returns:
jax.interpreters.xla.DeviceArray: Self-attention array for q, k, v arrays. (L_q by L_k)
"""
assert query.shape[-1] == key.shape[-1] == value.shape[
-1], "Embedding dimensions of q, k, v aren't all the same"
# scaling down (Q. K) dot product with square root of depth
depth = query.shape[-1]
# Calculate scaled query key dot product according to formula above
dots = jnp.matmul(query, jnp.swapaxes(key, -1, -2)) / jnp.sqrt(depth)
# Apply the mask
if mask is not None:
dots = jnp.where(mask, dots, jnp.full_like(dots, -1e9))
# Softmax formula implementation
# Use trax.fastmath.logsumexp of dots to avoid underflow by division by large numbers
logsumexp = trax.fastmath.logsumexp(dots, axis=-1, keepdims=True)
# Note: softmax = e^(dots - logsumexp(dots)) = E^dots / sumexp(dots)
dots = jnp.exp(dots - logsumexp)
# Multiply dots by value to get self-attention
# Use jnp.matmul()
attention = jnp.matmul(dots, value)
return attention
def entropy(self, log_probs):
probs = jnp.exp(log_probs)
return -jnp.sum(probs * log_probs, axis=-1)
def entropy(self, log_probs):
del log_probs # would be helpful if self._std was learnable
return jnp.exp(self._std) + .5 * jnp.log(2.0 * jnp.pi * jnp.e)
def awr_weights(advantages, beta):
return jnp.exp(advantages / beta)
### END CODE HERE ###
return -log_ppx
# In[49]:
# UNQ_C6 (UNIQUE CELL IDENTIFIER, DO NOT EDIT)
# Testing
model = GRULM()
model.init_from_file('model.pkl.gz')
batch = next(data_generator(batch_size, max_length, lines, shuffle=False))
preds = model(batch[0])
log_ppx = test_model(preds, batch[1])
print('The log perplexity and perplexity of your model are respectively',
log_ppx, np.exp(log_ppx))
# **Expected Output:** The log perplexity and perplexity of your model are respectively around 1.9 and 7.2.
# <a name='5'></a>
# # Part 5: Generating the language with your own model
#
# We will now use your own language model to generate new sentences for that we need to make draws from a Gumble distribution.
# The Gumbel Probability Density Function (PDF) is defined as:
#
# $$ f(z) = {1\over{\beta}}e^{(-z+e^{(-z)})} $$
#
# where: $$ z = {(x - \mu)\over{\beta}}$$
#
# The maximum value, which is what we choose as the prediction in the last step of a Recursive Neural Network `RNN` we are using for text generation, in a sample of a random variable following an exponential distribution approaches the Gumbel distribution when the sample increases asymptotically. For that reason, the Gumbel distribution is used to sample from a categorical distribution.
def Exp():
"""Returns a layer that computes the element-wise exponential of a tensor."""
return Fn('Exp', lambda x: jnp.exp(x)) # pylint: disable=unnecessary-lambda
def forward(self, x):
"""Executes this layer as part of a forward pass through the model.
Args:
x: Tensor of same shape and dtype as the input signature used to
initialize this layer.
Returns:
Tensor of same shape and dtype as the input.
"""
m1, m2, mb, w1, w2, b2 = self.weights
if self._mode != 'predict':
w1 = jnp.reshape(w1.T, (-1, self._d_ff))
w2 = jnp.reshape(w2, (self._d_ff, -1))
x_shape = x.shape
x = jnp.reshape(x,
[-1, x_shape[-1]]) # Easier to operate on flattened x.
# Q: should we add bias and/or put relu after the low-rank m1 dot?
mask_logits = jnp.dot(jnp.dot(x, m1), m2) + mb
mask_logits = jnp.reshape(mask_logits, [-1, self._d1, self._d2])
# Softmax.
mask_logsumexp = fastmath.logsumexp(mask_logits,
axis=-1,
keepdims=True)
log_mask = mask_logits - mask_logsumexp
mask = jnp.exp(log_mask)
# Gumbel-softmax with straight-through discretization.
rng1, rng2 = fastmath.random.split(self.rng, 2)
u = fastmath.random.uniform(rng1, mask.shape, jnp.float32, 1e-6,
1.0 - 1e-6)
g = -jnp.log(-jnp.log(u))
quant_mask = jnp.argmax(log_mask + g * self._temperature, axis=-1)
if self._mode == 'train':
# Tricks from Section 2.1 in https://arxiv.org/abs/1801.09797
quant_mask = tl.one_hot(quant_mask, self._n_elements_in_block)
quant_mask = fastmath.stop_gradient(quant_mask)
quant_mask += mask - fastmath.stop_gradient(
mask) # straight-through
# We will sometimes (quant_prob of the batches) use the soft-mask instead
# of the quantized mask to improve training stability (see paper above).
select = fastmath.random.uniform(rng2, (), jnp.float32, 0.0, 1.0)
quant_mask = jnp.where(select < self._quant_prob, quant_mask, mask)
quant_mask = jnp.reshape(quant_mask, [-1, self._d_ff])
if self._mode == 'train':
# In training, run full matmul to get benefits from the above tricks.
mid = jnp.dot(x, w1) * quant_mask # [joint_batch, d_ff]
relu = jnp.where(mid <= 0, jnp.zeros_like(mid), mid)
res = jnp.dot(relu, w2) + b2
elif self._mode == 'predict':
# w1 = jnp.reshape(w1.T, (self._d1, self._d2, -1))
# w2 = jnp.reshape(w2, (self._d1, self._d2, -1))
# This implementation mimicks inference. It's not efficient for large
# size of joint_batch, but at inference that will be 1 most of the time.
# Shapes:
# quant_mask is [joint_batch, self._d1]
# w1 is [d_model, self._d1, self._d2]
# we'll index w1 with advanced numpy indexing, first range over
# self._d1 times the batch size, second range being quant_mask
batch_size = quant_mask.shape[0]
idx1 = jnp.array([jnp.arange(self._d1)] * batch_size)
# flatten indices and select from w1
idx1 = jnp.reshape(idx1, [-1])
idx2 = jnp.reshape(quant_mask, [-1])
w = w1[idx1,
idx2, :] # now we have per-element weights with batch dim
w = jnp.reshape(w, [batch_size, self._d1, -1])
mid = jnp.einsum('ai,aji->aj', x, w)
relu = jnp.where(mid <= 0, jnp.zeros_like(mid), mid)
# w2 is [self._d1, self._d2, d_model]
v = w2[idx1, idx2, :]
v = jnp.reshape(v, [batch_size, self._d1, -1])
res = jnp.einsum('ai,aij->aj', relu, v) + b2
else:
quant_mask = tl.one_hot(quant_mask, self._n_elements_in_block)
quant_mask = jnp.reshape(quant_mask, [-1, self._d_ff])
mid = jnp.dot(x, w1) * quant_mask # [joint_batch, d_ff]
relu = jnp.where(mid <= 0, jnp.zeros_like(mid), mid)
res = jnp.dot(relu, w2) + b2
return jnp.reshape(res, x_shape) # un-flatten if needed
def unittest_test_model(target, pretrained_model):
successful_cases = 0
failed_cases = []
test_cases = [
{
"name": "check_default_example",
"input": {
"file":
pickle.load(open("./test_support_files/test_batch_1.pkl",
"rb")),
},
"expected": 1.7646706,
}, # 1.7646706 5.8396482
{
"name": "check_example_2",
"input": {
"file":
pickle.load(open("./test_support_files/test_batch_2.pkl",
"rb")),
},
"expected": 1.5336857,
}, # 1.5336857 4.635229
{
"name": "check_example_3",
"input": {
"file":
pickle.load(open("./test_support_files/test_batch_3.pkl",
"rb")),
},
"expected": 1.5870862,
}, # 1.5870862 4.889481
]
for test_case in test_cases:
preds = pretrained_model(test_case["input"]["file"][0])
test_case["input"]["preds"] = preds
test_case["input"]["target"] = test_case["input"]["file"][1]
test_case["input"].pop("file", None)
result = target(**test_case["input"])
try:
assert jnp.isclose(result, test_case["expected"], atol=1e-5)
successful_cases += 1
except:
failed_cases.append({
"name": test_case["name"],
"expected": test_case["expected"],
"got": result,
})
print(
f"Your log perplexity does not match with expected value.\nCheck if you are getting rid of the padding or checking if the target equals 0.\nIf your result is an array instead of a float, check if you are using numpy to perform the sums.\n\t Expected value near: {failed_cases[-1].get('expected')}.\n\t Got {failed_cases[-1].get('got')}."
)
try:
assert not jnp.exp(result) == jnp.inf
successful_cases += 1
except:
failed_cases.append({
"name": test_case["name"],
"expected": jnp.exp(test_case["expected"]),
"got": jnp.exp(result),
})
print(
f"Your perplexity overflowed. Take a look to the axis you are using in np.sum() function.\n\t Expected value near: {failed_cases[-1].get('expected')}.\n\t Got {failed_cases[-1].get('got')}."
)
try:
assert jnp.isclose(jnp.exp(result),
jnp.exp(test_case["expected"]),
atol=1e-5)
successful_cases += 1
except:
failed_cases.append({
"name": test_case["name"],
"expected": jnp.exp(test_case["expected"]),
"got": jnp.exp(result),
})
print(
f"Your perplexity does not match with expected.\n\t Expected value near: {failed_cases[-1].get('expected')}.\n\t Got {failed_cases[-1].get('got')}."
)
if len(failed_cases) == 0:
print("\033[92m All tests passed")
else:
print("\033[92m", successful_cases, " Tests passed")
print("\033[91m", len(failed_cases), " Tests failed")
def entropy(self, inputs):
(_, std) = self._params(inputs)
return jnp.sum(jnp.exp(std) + .5 * jnp.log(2.0 * jnp.pi * jnp.e),
axis=-1)
def entropy(self, inputs):
log_probs = inputs
probs = jnp.exp(log_probs)
return -jnp.sum(probs * log_probs, axis=-1)
### END CODE HERE ###
return -log_ppx
# In[20]:
# UNQ_C6 (UNIQUE CELL IDENTIFIER, DO NOT EDIT)
# Testing
model = GRULM()
model.init_from_file('model.pkl.gz')
batch = next(data_generator(batch_size, max_length, lines, shuffle=False))
preds = model(batch[0])
log_ppx = test_model(preds, batch[1])
print('The log perplexity and perplexity of your model are respectively', log_ppx, np.exp(log_ppx))
# **Expected Output:** The log perplexity and perplexity of your model are respectively around 1.9 and 7.2.
# <a name='5'></a>
# # Part 5: Generating the language with your own model
#
# We will now use your own language model to generate new sentences for that we need to make draws from a Gumble distribution.
# The Gumbel Probability Density Function (PDF) is defined as:
#
# $$ f(z) = {1\over{\beta}}e^{(-z+e^{(-z)})} $$
#
# where: $$ z = {(x - \mu)\over{\beta}}$$
#
def forward(self, x):
"""Executes this layer as part of a forward pass through the model.
Args:
x: Tensor of same shape and dtype as the input signature used to
initialize this layer.
Returns:
Tensor of same shape and dtype as the input.
"""
m1, w1, w2, b2 = self.weights
x_shape = x.shape
x = jnp.reshape(x,
[-1, x_shape[-1]]) # Easier to operate on flattened x.
# Q: check if we need bias and/or put relu after the m1 dot?
mask_logits = jnp.dot(x, m1)
# Softmax.
mask_logsumexp = fastmath.logsumexp(mask_logits,
axis=-1,
keepdims=True)
log_mask = mask_logits - mask_logsumexp
mask = jnp.exp(log_mask)
# Gumbel-softmax with straight-through discretization.
# TODO(lukaszkaiser, chowdhery): Extract this block and share
rng1, rng2 = fastmath.random.split(self.rng, 2)
u = fastmath.random.uniform(rng1, mask.shape, jnp.float32, 1e-6,
1.0 - 1e-6)
g = -jnp.log(-jnp.log(u))
selected_experts = jnp.argmax(log_mask + g * self._temperature,
axis=-1)
if self._mode == 'train':
# Tricks from Section 2.1 in https://arxiv.org/abs/1801.09797
quant_mask = tl.one_hot(selected_experts, self._num_experts)
quant_mask = fastmath.stop_gradient(quant_mask)
quant_mask += mask - fastmath.stop_gradient(
mask) # straight-through
# We will sometimes (50% of the batches) use the soft-mask instead of
# the quantized mask to improve training stability (see the paper above).
# Q: is selecting 50% of batches the best? Other %? Mixed in-batch?
select = fastmath.random.uniform(rng2, (), jnp.float32, -1.0, 1.0)
quant_mask = jnp.where(select > 0.0, quant_mask, mask)
else:
quant_mask = tl.one_hot(selected_experts, self._num_experts)
quant_mask = jnp.reshape(quant_mask, [-1, self._num_experts, 1])
quant_mask_shape = quant_mask.shape
batch_size = quant_mask.shape[0]
if self._mode == 'predict' and batch_size == 1:
# This implementation mimicks inference for batch_size 1.
start_idx = selected_experts[0] * self._n_elements_in_block
# w1 is [d_model, d_ff], w is [d_model, n_elements_in_block]
w = fastmath.dynamic_slice(
w1, [0, start_idx], [w1.shape[0], self._n_elements_in_block])
mid = jnp.dot(x, w)
relu = jnp.where(mid <= 0, jnp.zeros_like(mid), mid)
# w2 is [d_ff, d_model], v is [n_elements_in_block, d_model]
v = fastmath.dynamic_slice(
w2, [start_idx, 0], [self._n_elements_in_block, w2.shape[-1]])
v = jnp.reshape(v, [self._n_elements_in_block, -1])
res = jnp.dot(relu, v) + b2
else:
expanded_mask = jnp.broadcast_to(
quant_mask, (quant_mask_shape[0], quant_mask.shape[1],
self._n_elements_in_block))
expanded_mask = jnp.reshape(expanded_mask, (-1, self._d_ff))
mid = jnp.dot(x, w1) * expanded_mask # [joint_batch, d_ff]
relu = jnp.where(mid <= 0, jnp.zeros_like(mid), mid)
res = jnp.dot(relu, w2) + b2
return jnp.reshape(res, x_shape) # un-flatten if needed
def entropy(self, inputs):
log_probs = tl.LogSoftmax()(inputs)
probs = jnp.exp(log_probs)
return -jnp.sum(probs * log_probs, axis=-1)