Upload to Main

deepcore/methods/submodular.py

@@ -0,0 +1,116 @@
+from .earlytrain import EarlyTrain
+import numpy as np
+import torch
+from .methods_utils import cossim_np, submodular_function, submodular_optimizer
+from ..nets.nets_utils import MyDataParallel
+
+
+class Submodular(EarlyTrain):
+    def __init__(self, dst_train, args, fraction=0.5, random_seed=None, epochs=200, specific_model=None, balance=True,
+                 function="GraphCut", greedy="LazyGreedy", metric="cossim", **kwargs):
+        super(Submodular, self).__init__(dst_train, args, fraction, random_seed, epochs, specific_model, **kwargs)
+
+        if greedy not in submodular_optimizer.optimizer_choices:
+            raise ModuleNotFoundError("Greedy optimizer not found.")
+        print(f"The Submodular Method is {function}")
+        self._greedy = greedy
+        self._metric = metric
+        self._function = function
+
+        self.balance = balance
+
+    def before_train(self):
+        pass
+
+    def after_loss(self, outputs, loss, targets, batch_inds, epoch):
+        pass
+
+    def before_epoch(self):
+        pass
+
+    def after_epoch(self):
+        pass
+
+    def before_run(self):
+        pass
+
+    def num_classes_mismatch(self):
+        raise ValueError("num_classes of pretrain dataset does not match that of the training dataset.")
+
+
+    def calc_gradient(self, index=None):
+        '''
+        Calculate gradients matrix on current network for specified training dataset.
+        '''
+        self.model.eval()
+        data_loader = self.select_dm(self.dst_train, index, is_train=False)
+        # Initialize a matrix to save gradients.
+        # (on cpu)
+        gradients = []
+
+        for i, batch in enumerate(data_loader):
+
+            self.optim.zero_grad()
+            image, label = batch['img'].cuda(), batch['label'].cuda()
+            bs_size = image.shape[0]
+            loss,visual_embedding,logit = self.model(image,label,cal_gradient=True)
+            embed_dim = visual_embedding.shape[-1]
+            with torch.no_grad():
+                bias_parameters_grads = torch.autograd.grad(loss, logit)[0]
+                weight_parameters_grads = visual_embedding.view(bs_size, 1,
+                                        -1).repeat(1, self.num_classes, 1) *\
+                                        bias_parameters_grads.view(bs_size, self.num_classes,
+                                        1).repeat(1, 1, embed_dim)
+                gradients.append(torch.cat([bias_parameters_grads, weight_parameters_grads.flatten(1)],
+                                            dim=1).cpu().numpy())
+
+        gradients = np.concatenate(gradients, axis=0,dtype=np.float32)
+        print('Finish Gradient Calculation')
+        return gradients
+
+    def finish_run(self):
+        if isinstance(self.model, MyDataParallel):
+            self.model = self.model.module
+
+        # Turn on the embedding recorder and the no_grad flag
+
+        self.model.no_grad = True
+        self.train_indx = np.arange(self.n_train)
+
+        gradients = self.calc_gradient(index=None)
+
+        if self.balance:
+            selection_result = np.array([], dtype=np.int64)
+            for c in range(self.num_classes):
+                print(f'class {c}')
+                c_indx = self.train_indx[self.dst_train_label == c]
+                # Calculate gradients into a matrix
+                c_gradients = gradients[c_indx]
+                # Instantiate a submodular function
+                submod_function = submodular_function.__dict__[self._function](index=c_indx,
+                                    similarity_kernel=lambda a, b:cossim_np(c_gradients[a], c_gradients[b]))
+                submod_optimizer = submodular_optimizer.__dict__[self._greedy](args=self.args,
+                                    index=c_indx, budget=round(self.fraction * len(c_indx)), already_selected=[])
+
+                c_selection_result = submod_optimizer.select(gain_function=submod_function.calc_gain,
+                                                             update_state=submod_function.update_state)
+                selection_result = np.append(selection_result, c_selection_result)
+        else:
+            # Calculate gradients into a matrix
+            gradients = self.calc_gradient()
+            # Instantiate a submodular function
+            submod_function = submodular_function.__dict__[self._function](index=self.train_indx,
+                                        similarity_kernel=lambda a, b: cossim_np(gradients[a], gradients[b]))
+            submod_optimizer = submodular_optimizer.__dict__[self._greedy](args=self.args, index=self.train_indx,
+                                                                              budget=self.coreset_size)
+            selection_result = submod_optimizer.select(gain_function=submod_function.calc_gain,
+                                                       update_state=submod_function.update_state)
+
+        self.model.no_grad = False
+        return {"indices": selection_result}
+
+    def select(self, **kwargs):
+        selection_result = self.run()
+        return selection_result
+
+