Defining your own custom criterion for dynamics evaluation

Defining your own custom criterion can be very useful as the same instance you can either use it for analysing training dynamics or you can use it as IB-based loss objective for HSIC networks.

Importing glow modules

# importing Estimator class from glow
from glow.information_bottleneck import Estimator

Custom criterion class

Following is the general structure which your custom criterion class should follow in order for the PyGlow to detect your criterion function. NOTE: In most of the cases you won’t require implementation for eval_dynamics_segment so only criterion function is sufficient for PyGlow to accept your instance. You need eval_dynamics_segment implementation when you totally have to change the way you want to process a single dynamics segment.

class CustomCriterion(Estimator):
    This is a custom class for implementing your
    own criterion function (measure of dependence
    between random variables)

        function (callable, optional): this is an example argument
        which denotes any function which your criterion demands
        gpu (bool): if true then all the computation is done on
        `GPU` otherwise on `CPU`
        **kwargs: The keyword arguments store all the parameters
        that are essential for the criterion function

    def __init__(self, function, gpu, **kwargs):
        super().__init__(gpu, **kwargs)
        self.func = function  # attributes other than parameters

    def criterion(self, x, y):
        Custom criterion function.

            x (torch.Tensor): contains random variable one
            y (torch.Tensor): contains random variable two

            (torch.Tensor): calculated criterion of the two random variables 'x' and 'y'

        # Define your logic for calculating criterion here
        # TODO

    # implement this only if you want to totally change the
    # logic of evaluating a single dyanmics segment else
    # there is inbuilt function that handles it.
    def eval_dynamics_segment(self, dynamics_segment):
        Evaluates criterion for a dynamics segment using the
        criterion defined in the class itself.

            dynamics_segment(iterable): this is a list of :class:`torch.Tensor`
            type objects whose first entry is input batch, last entry is labels
            batch and intermediate entries are hidden layers outputs batch

            (iterable): list of calculated coordinates according to the criterion
            with length equal to 'len(dynamics_segment)-2'

        segment_size = len(dynamics_segment)
        # NOTE: dynamics_segment is a list where
        # input = dynamics_segment[0]
        # label = dynamics_segment[segment_size - 1]
        # and remaining tensors are hidden layers output
        output_segment = []  # output list which contains calculated coordinates
        for idx in range(1, segment_size-1):
            h = dynamics_segment[idx]
            # Define your logic of using h, input and label here
            # to apply your custom criterion on them.
            output_segment.append([self.criterion(h, x), self.criterion(h, y)])

        return output_segment