algos/models/algorithm_base.py
def process(
            self,
            batch,
            EPSILON=0,
            enforced_mask=None,
            extract_formulas=False,
            apply_formulas=False,
            fit_decision_tree=False,
            apply_decision_tree=False,
            compute_losses=True,
            hardcode_concepts=None,
            hardcode_outputs=False,
            debug=False):
        '''
        Method that takes a batch, does all iterations of every graph inside it
        and accumulates all metrics/losses.
        '''

        SIZE, STEPS_SIZE = self.prepare_constants(batch)
        self.hardcode_concepts = hardcode_concepts
        self.hardcode_outputs = hardcode_outputs
        self.apply_decision_tree = apply_decision_tree

        # Pytorch Geometric batches along the node dimension, but we execute
        # along the temporal (step) dimension, hence we need to transpose
        # a few tensors. Done by `prepare_batch`.
        batch = utils.prepare_batch(batch)
        if self.drop_last_concept:
            batch.concepts = batch.concepts[:, :, :-1]
            batch.last_concepts_real = batch.last_concepts_real[:, :-1]
        # When we want to calculate last step metrics/accuracies
        # we need to take into account again different termination per graph
        # hence we save last step tensors (e.g. outputs/concepts) into their
        # corresponding tensor. The function below prepares these tensors
        # (all set to zeros, except masking for computation, which are ones)
        self.set_initial_last_states(batch)
        # Prepare masking tensors (each graph does at least 1 iteration of the algo)
        self.prepare_initial_masks(batch)
        # A flag if we had a wrong graph in the batch. Used for visualisation
        # of what went wrong
        self.wrong_flag = False
        assert self.mask_cp.all(), self.mask_cp
        to_process = torch.ones([batch.num_graphs], device=_DEVICE)

        while self.loop_condition(
                batch.termination[self.steps-1] if self.training or fit_decision_tree or not self.concepts_frozen() else self.mask_cp,
                STEPS_SIZE):

            # take inputs/target values/concepts
            x_curr, y_curr = batch.x[self.steps], batch.y[self.steps]
            concepts_real = batch.concepts[self.steps]
            concepts_continue_real = batch.concepts[min(self.steps+1, len(batch.concepts)-1)]
            if not self.training:
                assert (self.last_continue_logits > 0).any() or True

            # Some algorithms, e.g. parallel colouring outputs fewer values than it takes
            # (e.g. priorities for colouring are unchanged one every step)
            # so if we reuse our last step outputs, they need to be fed back in.
            # NOTE self.get_step_input always takes x_curr, if we train
            inp = self.get_step_input(x_curr, batch)

            true_termination = batch.termination[self.steps] if self.steps < STEPS_SIZE else torch.zeros_like(batch.termination[self.steps-1])

            # Does one iteration of the algo and accumulates statistics
            self.loop_body(batch,
                           inp,
                           y_curr,
                           concepts_real,
                           true_termination,
                           compute_losses,
                           extract_formulas=extract_formulas,
                           apply_formulas=apply_formulas,
                           fit_decision_tree=fit_decision_tree,
                           apply_decision_tree=apply_decision_tree,
                           concepts_continue_real=concepts_continue_real)
            # And calculate what graphs would execute on the next step.
            self.mask, self.mask_cp = type(self).get_masks(self.training, batch, true_termination if  fit_decision_tree or not self.concepts_frozen() else self.last_continue_logits, enforced_mask)
        if not self.training:
            # After we are done with the execution, accumulate statistics related
            # to last step accuracies.

            self.aggregate_last_step(
                batch.batch,
                self.last_output_logits,
                batch.y[-1],
                self.last_concepts_logits > 0,
                batch.last_concepts_real, self.last_preds_with_formula)

        Method that takes a batch, does all iterations of every graph inside it
        and accumulates all metrics/losses.
        
Method that takes a batch, does all iterations of every graph inside it
and accumulates all metrics/losses.
Method that takes a batch does all iterations of every graph inside it and accumulates all metrics / losses .
def process(
            self,
            batch,
            EPSILON=0,
            enforced_mask=None,
            extract_formulas=False,
            apply_formulas=False,
            fit_decision_tree=False,
            apply_decision_tree=False,
            compute_losses=True,
            hardcode_concepts=None,
            hardcode_outputs=False,
            debug=False):
        SIZE, STEPS_SIZE = self.prepare_constants(batch)
        self.hardcode_concepts = hardcode_concepts
        self.hardcode_outputs = hardcode_outputs
        self.apply_decision_tree = apply_decision_tree
        batch = utils.prepare_batch(batch)
        if self.drop_last_concept:
            batch.concepts = batch.concepts[:, :, :-1]
            batch.last_concepts_real = batch.last_concepts_real[:, :-1]
        self.set_initial_last_states(batch)
        self.prepare_initial_masks(batch)
        self.wrong_flag = False
        assert self.mask_cp.all(), self.mask_cp
        to_process = torch.ones([batch.num_graphs], device=_DEVICE)
        while self.loop_condition(
                batch.termination[self.steps-1] if self.training or fit_decision_tree or not self.concepts_frozen() else self.mask_cp,
                STEPS_SIZE):
            x_curr, y_curr = batch.x[self.steps], batch.y[self.steps]
            concepts_real = batch.concepts[self.steps]
            concepts_continue_real = batch.concepts[min(self.steps+1, len(batch.concepts)-1)]
            if not self.training:
                assert (self.last_continue_logits > 0).any() or True
            inp = self.get_step_input(x_curr, batch)
            true_termination = batch.termination[self.steps] if self.steps < STEPS_SIZE else torch.zeros_like(batch.termination[self.steps-1])
            self.loop_body(batch,
                           inp,
                           y_curr,
                           concepts_real,
                           true_termination,
                           compute_losses,
                           extract_formulas=extract_formulas,
                           apply_formulas=apply_formulas,
                           fit_decision_tree=fit_decision_tree,
                           apply_decision_tree=apply_decision_tree,
                           concepts_continue_real=concepts_continue_real)
            self.mask, self.mask_cp = type(self).get_masks(self.training, batch, true_termination if  fit_decision_tree or not self.concepts_frozen() else self.last_continue_logits, enforced_mask)
        if not self.training:
            self.aggregate_last_step(
                batch.batch,
                self.last_output_logits,
                batch.y[-1],
                self.last_concepts_logits > 0,
                batch.last_concepts_real, self.last_preds_with_formula)
def process ( self , batch , EPSILON = 0 , enforced_mask = None , extract_formulas = False , apply_formulas = False , fit_decision_tree = False , apply_decision_tree = False , compute_losses = True , hardcode_concepts = None , hardcode_outputs = False , debug = False ) : SIZE , STEPS_SIZE = self . prepare_constants ( batch ) self . hardcode_concepts = hardcode_concepts self . hardcode_outputs = hardcode_outputs self . apply_decision_tree = apply_decision_tree batch = utils . prepare_batch ( batch ) if self . drop_last_concept : batch . concepts = batch . concepts [ : , : , : - 1 ] batch . last_concepts_real = batch . last_concepts_real [ : , : - 1 ] self . set_initial_last_states ( batch ) self . prepare_initial_masks ( batch ) self . wrong_flag = False assert self . mask_cp . all ( ) , self . mask_cp to_process = torch . ones ( [ batch . num_graphs ] , device = _DEVICE ) while self . loop_condition ( batch . termination [ self . steps - 1 ] if self . training or fit_decision_tree or not self . concepts_frozen ( ) else self . mask_cp , STEPS_SIZE ) : x_curr , y_curr = batch . x [ self . steps ] , batch . y [ self . steps ] concepts_real = batch . concepts [ self . steps ] concepts_continue_real = batch . concepts [ min ( self . steps + 1 , len ( batch . concepts ) - 1 ) ] if not self . training : assert ( self . last_continue_logits > 0 ) . any ( ) or True inp = self . get_step_input ( x_curr , batch ) true_termination = batch . termination [ self . steps ] if self . steps < STEPS_SIZE else torch . zeros_like ( batch . termination [ self . steps - 1 ] ) self . loop_body ( batch , inp , y_curr , concepts_real , true_termination , compute_losses , extract_formulas = extract_formulas , apply_formulas = apply_formulas , fit_decision_tree = fit_decision_tree , apply_decision_tree = apply_decision_tree , concepts_continue_real = concepts_continue_real ) self . mask , self . mask_cp = type ( self ) . get_masks ( self . training , batch , true_termination if fit_decision_tree or not self . concepts_frozen ( ) else self . last_continue_logits , enforced_mask ) if not self . training : self . aggregate_last_step ( batch . batch , self . last_output_logits , batch . y [ - 1 ] , self . last_concepts_logits > 0 , batch . last_concepts_real , self . last_preds_with_formula )
Method that takes a batch, does all iterations of every graph inside it
and accumulates all metrics/losses.
Method that takes a batch does all iterations of every graph inside it and accumulates all metrics / losses .
'''
        Method that takes a batch, does all iterations of every graph inside it
        and accumulates all metrics/losses.
        ''' # Pytorch Geometric batches along the node dimension, but we execute
 # along the temporal (step) dimension, hence we need to transpose
 # a few tensors. Done by `prepare_batch`.
 # When we want to calculate last step metrics/accuracies
 # we need to take into account again different termination per graph
 # hence we save last step tensors (e.g. outputs/concepts) into their
 # corresponding tensor. The function below prepares these tensors
 # (all set to zeros, except masking for computation, which are ones)
 # Prepare masking tensors (each graph does at least 1 iteration of the algo)
 # A flag if we had a wrong graph in the batch. Used for visualisation
 # of what went wrong
 # take inputs/target values/concepts
 # Some algorithms, e.g. parallel colouring outputs fewer values than it takes
 # (e.g. priorities for colouring are unchanged one every step)
 # so if we reuse our last step outputs, they need to be fed back in.
 # NOTE self.get_step_input always takes x_curr, if we train
 # Does one iteration of the algo and accumulates statistics
 # And calculate what graphs would execute on the next step.
 # After we are done with the execution, accumulate statistics related
 # to last step accuracies.

self
batch
EPSILON
enforced_mask
extract_formulas
apply_formulas
fit_decision_tree
apply_decision_tree
compute_losses
hardcode_concepts
hardcode_outputs
debug