How to: custom metric, but the aggregation is chosen through arguments?

[viktor-ktorvi]

Hi,

I want to implement a lot of custom metrics and for each to have min, mean, max, or concat aggregations. The way I've done this before involves a lot of boilerplate, and with the amount of different metrics I'm implementing it's getting tedious and leaves a lot of room for error. I'd like to know if there's an easy way to set the aggregation by providing an argument to the constructor of the custom metric.

This is roughly how I do things now:

class CustomMetric(Metric):
    def __init__(self):
        super().__init__()
        self.add_state("custom_state", default=torch.tensor(0.0))

    def update_state(self, data_obj: object):
        self.custom_state: Tensor = data_obj.extract_tensor()  # some custom calculation


class MinCustomMetric(CustomMetric, MinMetric):
    def update(self, data_obj: object):
        super().update_state(data_obj)
        super().update(self.custom_state)


class MeanCustomMetric(CustomMetric, MeanMetric):
    def update(self, data_obj: object):
        super().update_state(data_obj)
        super().update(self.custom_state)


class MaxCustomMetric(CustomMetric, MaxMetric):
    def update(self, data_obj: object):
        super().update_state(data_obj)
        super().update(self.custom_state)    

Ideally, what I'd like is the following:

class CustomMetric(Metric):
    def __init__(self, aggregation: str):
        super().__init__()

        self.aggregation = aggregation
        self.add_state("custom_state", default=torch.tensor(0.0))


    def update(self, data_obj: object):
        self.custom_state: Tensor = data_obj.extract_tensor()  # some custom calculation

        # ... somehow call the aggregation functionality

Is there a nice way to do this already?

Best I've come up with is a metric factory that will inherit dynamically but I feel like it's a bit hacky and confusing.

def create_custom_metric(aggregation: str) -> Metric:
    if aggregation  == "mean":
        base_metric = MeanMetric
    # ...

    class CustomMetric(base_metric):
        def __init__(self):
            super().__init__()
            self.add_state("custom_state", default=torch.tensor(0.0))

        def update(self, data_obj: object):
            self.custom_state: Tensor = data_obj.extract_tensor()  # some custom calculation
            super().update(self.custom_state)

Thanks! Hope I clearly stated the problem.