| 855 | self.num_classes = num_classes |
| 856 | |
| 857 | def update(self, minibatches, unlabeled=None): |
| 858 | device = "cuda" if minibatches[0][0].is_cuda else "cpu" |
| 859 | |
| 860 | # inputs |
| 861 | all_x = torch.cat([x for x, y in minibatches]) |
| 862 | # labels |
| 863 | all_y = torch.cat([y for _, y in minibatches]) |
| 864 | # one-hot labels |
| 865 | all_o = torch.nn.functional.one_hot(all_y, self.num_classes) |
| 866 | # features |
| 867 | all_f = self.featurizer(all_x) |
| 868 | # predictions |
| 869 | all_p = self.classifier(all_f) |
| 870 | |
| 871 | # Equation (1): compute gradients with respect to representation |
| 872 | all_g = autograd.grad((all_p * all_o).sum(), all_f)[0] |
| 873 | |
| 874 | # Equation (2): compute top-gradient-percentile mask |
| 875 | percentiles = np.percentile(all_g.cpu(), self.drop_f, axis=1) |
| 876 | percentiles = torch.Tensor(percentiles) |
| 877 | percentiles = percentiles.unsqueeze(1).repeat(1, all_g.size(1)) |
| 878 | mask_f = all_g.lt(percentiles.to(device)).float() |
| 879 | |
| 880 | # Equation (3): mute top-gradient-percentile activations |
| 881 | all_f_muted = all_f * mask_f |
| 882 | |
| 883 | # Equation (4): compute muted predictions |
| 884 | all_p_muted = self.classifier(all_f_muted) |
| 885 | |
| 886 | # Section 3.3: Batch Percentage |
| 887 | all_s = F.softmax(all_p, dim=1) |
| 888 | all_s_muted = F.softmax(all_p_muted, dim=1) |
| 889 | changes = (all_s * all_o).sum(1) - (all_s_muted * all_o).sum(1) |
| 890 | percentile = np.percentile(changes.detach().cpu(), self.drop_b) |
| 891 | mask_b = changes.lt(percentile).float().view(-1, 1) |
| 892 | mask = torch.logical_or(mask_f, mask_b).float() |
| 893 | |
| 894 | # Equations (3) and (4) again, this time mutting over examples |
| 895 | all_p_muted_again = self.classifier(all_f * mask) |
| 896 | |
| 897 | # Equation (5): update |
| 898 | loss = F.cross_entropy(all_p_muted_again, all_y) |
| 899 | self.optimizer.zero_grad() |
| 900 | loss.backward() |
| 901 | self.optimizer.step() |
| 902 | |
| 903 | return {'loss': loss.item()} |
| 904 | |
| 905 | |
| 906 | class SD(ERM): |