comments

Created using spr 1.3.4

Author: boomanaiden154

gematria/model/python/loss_utils_test.py

@@ -320,8 +320,8 @@ def test_multi_task_unknown_shape(self):
 
   def test_single_task_unknown_shape(self):
     num_tasks = 1
-    actual_output = tf.reshape(self.actual_outputs_array, (-1, 1))
-    expected_output = tf.reshape(self.expected_outputs_array, (-1, 1))
+    actual_output = tf.constant(self.actual_outputs_array)
+    expected_output = tf.constant(self.expected_outputs_array)
     mask = tf.ones_like(actual_output, tf.dtypes.bool)
     percentile_ranks = (0, 50, 75, 100)
 

gematria/model/python/model_base.py

@@ -32,7 +32,7 @@
 import math
 import os
 import random
-from typing import Optional, TypeVar, Union
+from typing import Optional, TypeVar, Union, TypedDict
 
 from absl import logging
 from gematria.basic_block.python import basic_block
@@ -120,6 +120,12 @@ def after_run(self, run_context: ..., run_values: ...):
       )
 
 
+class OutputDict(TypedDict):
+  output: tf.Tensor | None
+  output_deltas: tf.Tensor | None
+  output_mask_deltas: tf.Tensor | None
+
+
 class ModelBase(tf.Module, metaclass=abc.ABCMeta):
   """Base class for Gematria basic block processing models.
 
@@ -441,42 +447,43 @@ def output_tensor_names(self) -> Sequence[str]:
     return (ModelBase.OUTPUT_TENSOR_NAME,)
 
   @abc.abstractmethod
-  def _forward(self, feed_dict: FeedDict) -> dict[str, tf.Tensor]:
+  def _forward(self, feed_dict: FeedDict) -> OutputDict:
     """Implements the forward pass of the model.
 
     This function should be implemented in downstream models and calculate the
     outputs of the model given the inputs specified in feed_dict.
+
+    Returns:
+      A dictionary containing tensors. This should contain a key, 'output' in
+      seq2num mode and a key named 'output_deltas' in seq2seq mode.
     """
     pass
 
-  def __call__(self, feed_dict, train=False):
+  def __call__(self, feed_dict: FeedDict, train=False) -> OutputDict:
     """Implements the non-model specific part of the forward pass.
 
     This function wraps the _forward method and does relevant calculations
     when working with models that use deltas.
     """
-    if self._use_deltas:
-      output_dict = {}
-
-      if train:
-        output_dict['output_mask_deltas'] = tf.nn.embedding_lookup(
-            feed_dict['output_mask'],
-            feed_dict['delta_block_index'],
-            name='ModelBase.output_mask_deltas',
-        )
+    if not self._use_deltas:
+      return self._forward(feed_dict)
 
-      output = self._forward(feed_dict)
+    output = self._forward(feed_dict)
 
-      output_dict['output'] = tf.math.segment_sum(
-          output['output_deltas'],
+    if train:
+      output['output_mask_deltas'] = tf.nn.embedding_lookup(
+          feed_dict['output_mask'],
           feed_dict['delta_block_index'],
-          name=ModelBase.OUTPUT_TENSOR_NAME,
+          name='ModelBase.output_mask_deltas',
       )
-      output_dict['output_deltas'] = output['output_deltas']
 
-      return output_dict
-    else:
-      return self._forward(feed_dict)
+    output['output'] = tf.math.segment_sum(
+        output['output_deltas'],
+        feed_dict['delta_block_index'],
+        name=ModelBase.OUTPUT_TENSOR_NAME,
+    )
+
+    return output
 
   @abc.abstractmethod
   def _make_model_name(self) -> str:
@@ -1151,9 +1158,9 @@ def predict(
       batch_output_blocks = []
       with timer.scoped('ModelBase.predict - one batch'):
         schedule = self.schedule_batch(batch)
+        output_dict = self(schedule)
+        output = output_dict['output']
         if self._use_deltas:
-          output_dict = self(schedule)
-          output = output_dict['output']
           output_deltas = output_dict['output_deltas']
           output_index = 0
           for block_index, block in enumerate(batch):
@@ -1273,9 +1280,9 @@ def run_one_epoch():
     for _ in range(0, num_epochs):
       stats = run_one_epoch()
       logging.info('Training: %s', stats)
-      return stats
+    return stats
 
-  def compute_loss(self, schedule: FeedDict):
+  def _compute_loss(self, schedule: FeedDict) -> loss_utils.LossComputation:
     output = self(schedule, train=True)
     loss = loss_utils.LossComputation(
         output['output'],
@@ -1301,7 +1308,7 @@ def compute_loss(self, schedule: FeedDict):
 
   def compute_loss_tensor(self, schedule: FeedDict):
     return tf.reduce_mean(
-        self.compute_loss(schedule).loss_tensor(
+        self._compute_loss(schedule).loss_tensor(
             self._loss_normalization, self._loss_type
         )
     )
@@ -1327,27 +1334,25 @@ def train_batch(
       # TrainingEpochStats.__init__() as keyword arguments.
       with tf.GradientTape() as tape:
         stats = {}
-        loss = self.compute_loss(schedule)
+        loss = self._compute_loss(schedule)
         loss_tensor_per_task = loss.loss_tensor(
             self._loss_normalization, self._loss_type
         )
         loss_tensor = tf.reduce_mean(loss_tensor_per_task)
 
         # The list of variables to optimize. By default, the list is empty which
         # means optimize all trainable variables.
-        variables = set()
+        requested_variables = set()
         for variable_group in self._trained_variable_groups:
-          variables.update(
-              [
-                  variable.ref()
-                  for variable in self._variable_groups.get(variable_group)
-              ]
+          requested_variables.update(
+              variable.ref()
+              for variable in self._variable_groups.get(variable_group)
           )
 
         trainable_variables = self._get_trainable_variables()
         variables = (
-            [variable.deref() for variable in variables]
-            if variables
+            [variable.deref() for variable in requested_variables]
+            if requested_variables
             else trainable_variables
         )