Compute metric.
(self)
| 149 | self.features.append(features) |
| 150 | |
| 151 | def compute(self) -> Tuple[Tensor, Tensor]: |
| 152 | """Compute metric.""" |
| 153 | features = dim_zero_cat(self.features) |
| 154 | # random permute the features |
| 155 | idx = torch.randperm(features.shape[0]) |
| 156 | features = features[idx] |
| 157 | |
| 158 | # calculate probs and logits |
| 159 | prob = features.softmax(dim=1) |
| 160 | log_prob = features.log_softmax(dim=1) |
| 161 | |
| 162 | # split into groups |
| 163 | prob = prob.chunk(self.splits, dim=0) |
| 164 | log_prob = log_prob.chunk(self.splits, dim=0) |
| 165 | |
| 166 | # calculate score per split |
| 167 | mean_prob = [p.mean(dim=0, keepdim=True) for p in prob] |
| 168 | kl_ = [p * (log_p - m_p.log()) for p, log_p, m_p in zip(prob, log_prob, mean_prob)] |
| 169 | kl_ = [k.sum(dim=1).mean().exp() for k in kl_] |
| 170 | kl = torch.stack(kl_) |
| 171 | |
| 172 | # return mean and std |
| 173 | return kl.mean(), kl.std() |
| 174 | |
| 175 | def plot( |
| 176 | self, val: Optional[Union[Tensor, Sequence[Tensor]]] = None, ax: Optional[_AX_TYPE] = None |