def _replicated_train_step(inputs):
"""Replicated training step."""
features, y = inputs
per_replica_loss = self.optimizer.update(self.model, features, y, num_replicas)
per_replica_batchsz = tf.cast(get_shape_as_list(y)[0], tf.float32)
per_replica_report_loss = per_replica_loss * per_replica_batchsz
return per_replica_report_loss, per_replica_batchsz
def _train_step(inputs):
"""Replicated training step."""
features, y = inputs
loss = self.optimizer.update(self.model, features, y)
batchsz = get_shape_as_list(y)[0]
report_loss = loss * batchsz
return report_loss, batchsz
def smoothed_loss(self, logits, labels):
V = get_shape_as_list(logits)[-1]
one_hot = tf.one_hot(labels, V)
return tf.keras.losses.categorical_crossentropy(
y_true=one_hot,
y_pred=logits,
from_logits=True,
label_smoothing=self.label_smoothing)
def arc_label_loss(model, x, y):
arcs_pred, labels_pred = model(x)
arcs_gold, labels_gold = y
# First, trim gold labels to length of the input
B, T = get_shape_as_list(labels_pred)[:2]
flat_labels_pred = tf.reshape(labels_pred, (B * T, T, -1))
flat_arcs_pred = tf.reshape(arcs_pred, (B * T, T))
flat_labels_gold = tf.reshape(labels_gold[:, :T], [-1])
flat_arcs_gold = tf.reshape(arcs_gold[:, :T], [-1])
idx = tf.cast(tf.range(get_shape_as_list(flat_arcs_gold)[0]),
dtype=tf.int64)
indices = tf.stack([idx, flat_arcs_gold], 1)
flat_labels_gather = tf.gather_nd(flat_labels_pred, indices)
arc_loss = tf.compat.v1.losses.sparse_softmax_cross_entropy(
logits=flat_arcs_pred, labels=flat_arcs_gold)
rel_loss = tf.compat.v1.losses.sparse_softmax_cross_entropy(
logits=flat_labels_gather, labels=flat_labels_gold)
return arc_loss + rel_loss
def call(self, qkvm):
query, key, value, mask = qkvm
batchsz = get_shape_as_list(query)[0]
# (B, T, H, D) -> (B, H, T, D)
query = tf.transpose(
tf.reshape(self.w_Q(query), [batchsz, -1, self.h, self.d_k]),
[0, 2, 1, 3])
key = tf.transpose(
tf.reshape(self.w_K(key), [batchsz, -1, self.h, self.d_k]),
[0, 2, 1, 3])
x = self.attn_fn((query, key, value, mask))
x = tf.transpose(x, [0, 2, 1, 3])
x = tf.squeeze(x, axis=2)
return x
def _replica_test_step(inputs):
features, y = inputs
y = tf.cast(y, tf.int64)
##per_replica_cm = tf.zeros((nc, nc), dtype=tf.int64)
logits = self.model(features)
y_ = tf.argmax(logits, axis=1, output_type=tf.int64)
##indices = tf.stack((y, y_), axis=-1)
##dense_shape = tf.cast(tf.shape(per_replica_cm), tf.int64)
##sparse_ups = tf.SparseTensor(indices=indices, values=tf.ones(get_shape_as_list(indices)[0], dtype=tf.int64),
## dense_shape=dense_shape)
##per_replica_cm = tf.compat.v1.sparse_add(per_replica_cm, sparse_ups)
per_replica_acc = tf.reduce_sum(tf.cast(y == y_, tf.float32))
per_replica_loss = tf.compat.v1.losses.sparse_softmax_cross_entropy(labels=y, logits=logits)
per_replica_batchsz = tf.cast(get_shape_as_list(y)[0], tf.float32)
per_replica_report_loss = per_replica_loss * per_replica_batchsz
return per_replica_report_loss, per_replica_batchsz, per_replica_acc ##, per_replica_cm
def _test(self, loader, steps=0, **kwargs):
"""Test an epoch of data using either the input loader or using `tf.dataset`
In non-`tf.dataset` mode, we cycle the loader data feed, and pull a batch and feed it to the feed dict
When we use `tf.dataset`s under the hood, this function simply uses the loader to know how many steps
to train.
:param loader: A data feed
:param kwargs: See below
:Keyword Arguments:
* *dataset* (`bool`) Set to `True` if using `tf.dataset`s, defaults to `True`
* *reporting_fns* (`list`) A list of reporting hooks to use
* *verbose* (`dict`) A dictionary containing `console` boolean and `file` name if on
:return: Metrics
"""
metrics = [LAS(), UAS(), LCM(), UCM()]
pg = create_progress_bar(steps)
SET_TRAIN_FLAG(False)
for features, y in pg(loader):
heads_gold, labels_gold = y
greedy_heads_pred, greedy_labels_pred = self.model.decode(features)
B, T = get_shape_as_list(greedy_labels_pred)[:2]
labels_gold_trimmed = labels_gold[:, :T].numpy()
heads_gold_trimmed = heads_gold[:, :T].numpy()
for i in range(B):
for m in metrics:
if self.punct_eval is False:
labels_gold_trimmed[i] = masked_fill(
labels_gold_trimmed[i],
labels_gold_trimmed[i] == self.model.punct,
Offsets.PAD)
m.add(greedy_heads_pred[i], heads_gold_trimmed[i],
greedy_labels_pred[i], labels_gold_trimmed[i])
metrics = {m.name: m.score for m in metrics}
return metrics
def best_from(x_preds):
B, T, V = get_shape_as_list(x_preds)
sample_dist = tf.exp(tf.reshape(x_preds, (B * T, V)))
output = tf.reshape(tf.random.categorical(sample_dist, num_samples=1),
(B, T))
return output
def _num_toks(y):
return tf.reduce_prod(get_shape_as_list(y))
