| 9 | |
| 10 | |
| 11 | class Leaf(Node): |
| 12 | |
| 13 | def __init__(self, |
| 14 | index: int, |
| 15 | num_classes: int, |
| 16 | args: argparse.Namespace |
| 17 | ): |
| 18 | super().__init__(index) |
| 19 | |
| 20 | # Initialize the distribution parameters |
| 21 | if args.disable_derivative_free_leaf_optim: |
| 22 | self._dist_params = nn.Parameter(torch.randn(num_classes), requires_grad=True) |
| 23 | elif args.kontschieder_normalization: |
| 24 | self._dist_params = nn.Parameter(torch.ones(num_classes), requires_grad=False) |
| 25 | else: |
| 26 | self._dist_params = nn.Parameter(torch.zeros(num_classes), requires_grad=False) |
| 27 | |
| 28 | # Flag that indicates whether probabilities or log probabilities are computed |
| 29 | self._log_probabilities = args.log_probabilities |
| 30 | |
| 31 | self._kontschieder_normalization = args.kontschieder_normalization |
| 32 | |
| 33 | def forward(self, xs: torch.Tensor, **kwargs): |
| 34 | |
| 35 | # Get the batch size |
| 36 | batch_size = xs.size(0) |
| 37 | |
| 38 | # Keep a dict to assign attributes to nodes. Create one if not already existent |
| 39 | node_attr = kwargs.setdefault('attr', dict()) |
| 40 | # In this dict, store the probability of arriving at this node. |
| 41 | # It is assumed that when a parent node calls forward on this node it passes its node_attr object with the call |
| 42 | # and that it sets the path probability of arriving at its child |
| 43 | # Therefore, if this attribute is not present this node is assumed to not have a parent. |
| 44 | # The probability of arriving at this node should thus be set to 1 (as this would be the root in this case) |
| 45 | # The path probability is tracked for all x in the batch |
| 46 | if not self._log_probabilities: |
| 47 | node_attr.setdefault((self, 'pa'), torch.ones(batch_size, device=xs.device)) |
| 48 | else: |
| 49 | node_attr.setdefault((self, 'pa'), torch.zeros(batch_size, device=xs.device)) |
| 50 | |
| 51 | # Obtain the leaf distribution |
| 52 | dist = self.distribution() # shape: (k,) |
| 53 | # Reshape the distribution to a matrix with one single row |
| 54 | dist = dist.view(1, -1) # shape: (1, k) |
| 55 | # Duplicate the row for all x in xs |
| 56 | dists = torch.cat((dist,) * batch_size, dim=0) # shape: (bs, k) |
| 57 | |
| 58 | # Store leaf distributions as node property |
| 59 | node_attr[self, 'ds'] = dists |
| 60 | |
| 61 | # Return both the result of the forward pass as well as the node properties |
| 62 | return dists, node_attr |
| 63 | |
| 64 | def distribution(self) -> torch.Tensor: |
| 65 | if not self._kontschieder_normalization: |
| 66 | if self._log_probabilities: |
| 67 | return F.log_softmax(self._dist_params, dim=0) |
| 68 | else: |
no outgoing calls
no test coverage detected