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))