Make word counters creation a static function
@@ -50,23 +50,13 @@ def learning_rate_decay_fn(decay_type,
return decay_fn
+def register_word_counters(self, features, labels):
+ """Stores word counter operators for sequences (if any) of `features`
+ and `labels`.
-@six.add_metaclass(abc.ABCMeta)
-class Model(object):
-
- def __init__(self, name):
- self.name = name
-
- def __call__(self, features, labels, params, mode):
- """Creates the model. See `tf.estimator.Estimator`'s `model_fn` argument
- for more details about arguments and the returned value.
+ See also `onmt.utils.misc.WordCounterHook` that fetches these counters
+ to log their value in TensorBoard.
"""
- self._count_words(features, labels)
- with tf.variable_scope(self.name):
- return self._build(features, labels, params, mode)
-
- def _count_words(self, features, labels):
- """Stores a word counter operator for sequences of features and labels."""
def _add_counter(word_count, name):
word_count = tf.cast(word_count, tf.int64)
total_word_count = tf.Variable(
@@ -88,6 +78,21 @@ class Model(object):
if labels_length is not None:
_add_counter(tf.reduce_sum(labels_length), "labels")
+
+@six.add_metaclass(abc.ABCMeta)
+class Model(object):
+
+ def __init__(self, name):
+ self.name = name
+
+ def __call__(self, features, labels, params, mode):
+ """Creates the model. See `tf.estimator.Estimator`'s `model_fn` argument
+ for more details about arguments and the returned value.
+ """
+ register_word_counters(features, labels)
+ with tf.variable_scope(self.name):
+ return self._build(features, labels, params, mode)
+
@abc.abstractmethod
def _build(self, features, labels, params, mode):
"""Creates the model. Subclasses should override this function."""
