| 632 | return mean_diff + cova_diff |
| 633 | |
| 634 | def update(self, minibatches, unlabeled=None): |
| 635 | objective = 0 |
| 636 | penalty = 0 |
| 637 | nmb = len(minibatches) |
| 638 | |
| 639 | features = [self.featurizer(xi) for xi, _ in minibatches] |
| 640 | classifs = [self.classifier(fi) for fi in features] |
| 641 | targets = [yi for _, yi in minibatches] |
| 642 | |
| 643 | for i in range(nmb): |
| 644 | objective += F.cross_entropy(classifs[i], targets[i]) |
| 645 | for j in range(i + 1, nmb): |
| 646 | penalty += self.mmd(features[i], features[j]) |
| 647 | |
| 648 | objective /= nmb |
| 649 | if nmb > 1: |
| 650 | penalty /= (nmb * (nmb - 1) / 2) |
| 651 | |
| 652 | self.optimizer.zero_grad() |
| 653 | (objective + (self.hparams['mmd_gamma'] * penalty)).backward() |
| 654 | self.optimizer.step() |
| 655 | |
| 656 | if torch.is_tensor(penalty): |
| 657 | penalty = penalty.item() |
| 658 | |
| 659 | return {'loss': objective.item(), 'penalty': penalty} |
| 660 | |
| 661 | |
| 662 | class MMD(AbstractMMD): |