TypeError: Metric ConfusionMatrixMetric is not constructable from the `get_config()` method

Viewed 67

I am building a multiclass classification federated learning model using TensorFlow and I want to measure the model performance using a confusion matrix. I found this article and applied the following two functions:

class ConfusionMatrixMetric(tf.keras.metrics.Metric):
    """
    A custom Keras metric to compute the running average of the confusion matrix
    """
    def __init__(self, num_classes, **kwargs):
        super(ConfusionMatrixMetric,self).__init__(name='confusion_matrix_metric',**kwargs) # handles base args (e.g., dtype)
        self.num_classes=num_classes
        self.total_cm = self.add_weight("total", shape=(num_classes,num_classes), initializer="zeros")
        
    def reset_states(self):
        for s in self.variables:
            s.assign(tf.zeros(shape=s.shape))
            
    def update_state(self, y_true, y_pred,sample_weight=None):
        self.total_cm.assign_add(self.confusion_matrix(y_true,y_pred))
        return self.total_cm
        
    def result(self):
        return self.process_confusion_matrix()
    
    def confusion_matrix(self,y_true, y_pred):
        """
        Make a confusion matrix
        """
        y_pred=tf.argmax(y_pred,1)
        cm=tf.math.confusion_matrix(y_true,y_pred,dtype=tf.float32,num_classes=self.num_classes)
        return cm
    
    def process_confusion_matrix(self):
        "returns precision, recall and f1 along with overall accuracy"
        cm=self.total_cm
        diag_part=tf.linalg.diag_part(cm)
        precision=diag_part/(tf.reduce_sum(cm,0)+tf.constant(1e-15))
        recall=diag_part/(tf.reduce_sum(cm,1)+tf.constant(1e-15))
        f1=2*precision*recall/(precision+recall+tf.constant(1e-15))
        return precision,recall,f1
    
    def fill_output(self,output):
        results=self.result()
        for i in range(self.num_classes):
            output['precision_{}'.format(i)]=results[0][i]
            output['recall_{}'.format(i)]=results[1][i]
            output['F1_{}'.format(i)]=results[2][i]
class MySequential(keras.Sequential):
    
    def train_step(self, data):
        # Unpack the data. Its structure depends on your model and
        # on what you pass to `fit()`.

        x, y = data

        with tf.GradientTape() as tape:
            y_pred = self(x, training=True)  # Forward pass
            # Compute the loss value.
            # The loss function is configured in `compile()`.
            loss = self.compiled_loss(
                y,
                y_pred,
                regularization_losses=self.losses,
            )

        # Compute gradients
        trainable_vars = self.trainable_variables
        gradients = tape.gradient(loss, trainable_vars)

        # Update weights
        self.optimizer.apply_gradients(zip(gradients, trainable_vars))
        self.compiled_metrics.update_state(y, y_pred)
        output={m.name: m.result() for m in self.metrics[:-1]}
        if 'confusion_matrix_metric' in self.metrics_names:
            self.metrics[-1].fill_output(output)
        return output
        
        
    def test_step(self, data):
        # Unpack the data. Its structure depends on your model and
        # on what you pass to `fit()`.

        x, y = data

        y_pred = self(x, training=False)  # Forward pass
        # Compute the loss value.
        # The loss function is configured in `compile()`.
        loss = self.compiled_loss(
            y,
            y_pred,
            regularization_losses=self.losses,
        )

        self.compiled_metrics.update_state(y, y_pred)
        output={m.name: m.result() for m in self.metrics[:-1]}
        if 'confusion_matrix_metric' in self.metrics_names:
            self.metrics[-1].fill_output(output)    
        return output

Here is the model:

def create_keras_model():
  initializer = tf.keras.initializers.Zeros()
  return tf.keras.models.Sequential([
      tf.keras.layers.Input(shape=(7,)),
      tf.keras.layers.Dense(32),
      tf.keras.layers.Dense(4, kernel_initializer=initializer),
      tf.keras.layers.Softmax(),
  ])
def model_fn():
  keras_model = create_keras_model()
  return tff.learning.from_keras_model(
      keras_model,
      input_spec=train_data[0].element_spec,
      loss=tf.keras.losses.SparseCategoricalCrossentropy(),
      metrics=[tf.keras.metrics.SparseCategoricalAccuracy(),
               ConfusionMatrixMetric(4)])

but when I compile the model, I am getting this error: (provided wit the traceback)

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-10-628f39042da4> in <module>()
      2     model_fn,
      3     client_optimizer_fn=lambda: tf.keras.optimizers.SGD(learning_rate=0.00001),
----> 4     server_optimizer_fn=lambda: tf.keras.optimizers.SGD(learning_rate=1.99))

11 frames
/usr/local/lib/python3.7/dist-packages/tensorflow_federated/python/learning/federated_averaging.py in build_federated_averaging_process(model_fn, client_optimizer_fn, server_optimizer_fn, client_weighting, broadcast_process, model_update_aggregation_factory, metrics_aggregator, use_experimental_simulation_loop)
    289       broadcast_process=broadcast_process,
    290       model_update_aggregation_factory=model_update_aggregation_factory,
--> 291       metrics_aggregator=metrics_aggregator)
    292 
    293   server_state_type = iter_proc.state_type.member

/usr/local/lib/python3.7/dist-packages/tensorflow_federated/python/learning/framework/optimizer_utils.py in build_model_delta_optimizer_process(model_fn, model_to_client_delta_fn, server_optimizer_fn, broadcast_process, model_update_aggregation_factory, metrics_aggregator)
    622       broadcast_process=broadcast_process,
    623       aggregation_process=aggregation_process,
--> 624       metrics_aggregator=metrics_aggregator)
    625 
    626   return iterative_process.IterativeProcess(

/usr/local/lib/python3.7/dist-packages/tensorflow_federated/python/learning/framework/optimizer_utils.py in _build_one_round_computation(model_fn, server_optimizer_fn, model_to_client_delta_fn, broadcast_process, aggregation_process, metrics_aggregator)
    321       metrics_aggregation_computation = aggregator.sum_then_finalize(
    322           whimsy_model_for_metadata.metric_finalizers(),
--> 323           unfinalized_metrics_type)
    324 
    325   @computations.tf_computation(model_weights_type, model_weights_type.trainable,

/usr/local/lib/python3.7/dist-packages/tensorflow_federated/python/learning/metrics/aggregator.py in sum_then_finalize(metric_finalizers, local_unfinalized_metrics_type)
    170 
    171   @computations.federated_computation(
--> 172       computation_types.at_clients(local_unfinalized_metrics_type))
    173   def aggregator_computation(client_local_unfinalized_metrics):
    174     unfinalized_metrics_sum = intrinsics.federated_sum(

/usr/local/lib/python3.7/dist-packages/tensorflow_federated/python/core/impl/wrappers/computation_wrapper.py in __call__(self, tff_internal_types, *args)
    493       parameter_type = _parameter_type(parameters, parameter_types)
    494       wrapped_func = self._strategy(
--> 495           fn_to_wrap, fn_name, parameter_type, unpack=None)
    496 
    497     # Copy the __doc__ attribute with the documentation in triple-quotes from

/usr/local/lib/python3.7/dist-packages/tensorflow_federated/python/core/impl/wrappers/computation_wrapper.py in __call__(self, fn_to_wrap, fn_name, parameter_type, unpack)
    220     try:
    221       args, kwargs = unpack_arguments_fn(packed_args)
--> 222       result = fn_to_wrap(*args, **kwargs)
    223       if result is None:
    224         raise ComputationReturnedNoneError(fn_to_wrap)

/usr/local/lib/python3.7/dist-packages/tensorflow_federated/python/learning/metrics/aggregator.py in aggregator_computation(client_local_unfinalized_metrics)
    175         client_local_unfinalized_metrics)
    176 
--> 177     @computations.tf_computation(local_unfinalized_metrics_type)
    178     def finalizer_computation(unfinalized_metrics):
    179       finalized_metrics = collections.OrderedDict()

/usr/local/lib/python3.7/dist-packages/tensorflow_federated/python/core/impl/wrappers/computation_wrapper.py in __call__(self, tff_internal_types, *args)
    493       parameter_type = _parameter_type(parameters, parameter_types)
    494       wrapped_func = self._strategy(
--> 495           fn_to_wrap, fn_name, parameter_type, unpack=None)
    496 
    497     # Copy the __doc__ attribute with the documentation in triple-quotes from

/usr/local/lib/python3.7/dist-packages/tensorflow_federated/python/core/impl/wrappers/computation_wrapper.py in __call__(self, fn_to_wrap, fn_name, parameter_type, unpack)
    220     try:
    221       args, kwargs = unpack_arguments_fn(packed_args)
--> 222       result = fn_to_wrap(*args, **kwargs)
    223       if result is None:
    224         raise ComputationReturnedNoneError(fn_to_wrap)

/usr/local/lib/python3.7/dist-packages/tensorflow_federated/python/learning/metrics/aggregator.py in finalizer_computation(unfinalized_metrics)
    180       for metric_name, metric_finalizer in metric_finalizers.items():
    181         finalized_metrics[metric_name] = metric_finalizer(
--> 182             unfinalized_metrics[metric_name])
    183       return finalized_metrics
    184 

/usr/local/lib/python3.7/dist-packages/tensorflow/python/util/traceback_utils.py in error_handler(*args, **kwargs)
    151     except Exception as e:
    152       filtered_tb = _process_traceback_frames(e.__traceback__)
--> 153       raise e.with_traceback(filtered_tb) from None
    154     finally:
    155       del filtered_tb

/usr/local/lib/python3.7/dist-packages/tensorflow/python/framework/func_graph.py in autograph_handler(*args, **kwargs)
   1145           except Exception as e:  # pylint:disable=broad-except
   1146             if hasattr(e, "ag_error_metadata"):
-> 1147               raise e.ag_error_metadata.to_exception(e)
   1148             else:
   1149               raise

TypeError: in user code:

    File "/usr/local/lib/python3.7/dist-packages/tensorflow_federated/python/learning/metrics/finalizer.py", line 65, in finalizer  *
        keras_metric = create_keras_metric(metric)
    File "/usr/local/lib/python3.7/dist-packages/tensorflow_federated/python/learning/metrics/finalizer.py", line 148, in create_keras_metric  *
        _check_keras_metric_config_constructable(metric)
    File "/usr/local/lib/python3.7/dist-packages/tensorflow_federated/python/learning/metrics/finalizer.py", line 112, in _check_keras_metric_config_constructable  *
        f'Metric {metric_type_str} is not constructable from the '

    TypeError: Metric ConfusionMatrixMetric is not constructable from the `get_config()` method, because `__init__` takes extra arguments that are not included in the `get_config()`: ['num_classes']. Pass the metric constructor instead, or update the `get_config()` in the metric class to include these extra arguments.
    Example:
    class CustomMetric(tf.keras.metrics.Metric):
      def __init__(self, arg1):
        self._arg1 = arg1
    
      def get_config(self)
        config = super().get_config()
        config['arg1'] = self._arg1
        return config

I do not know where the problem is? or if there is another simpler way to perform the confusion matrix with a federated learning model! Thanks

0 Answers
Related