Upload to Main

deepcore/methods/grand.py

@@ -0,0 +1,108 @@
+from .earlytrain import EarlyTrain
+import torch, time
+import numpy as np
+from ..nets.nets_utils import MyDataParallel
+from tqdm import tqdm
+
+class GraNd(EarlyTrain):
+    def __init__(self, dst_train, args, fraction=0.5, random_seed=None, epochs=200, repeat=1,
+                 specific_model=None, balance=False, **kwargs):
+        super().__init__(dst_train, args, fraction, random_seed, epochs, specific_model,**kwargs)
+        self.epochs = epochs
+        self.n_train = len(self.dst_train)
+        self.coreset_size = round(self.n_train * fraction)
+        self.specific_model = specific_model
+        self.repeat = repeat
+
+        self.balance = balance
+
+    # def while_update(self, outputs, loss, targets, epoch, batch_idx, batch_size):
+    #     if batch_idx % self.args.print_freq == 0:
+    #         print('| Epoch [%3d/%3d] Iter[%3d/%3d]\t\tLoss: %.4f' % (
+    #             epoch, self.epochs, batch_idx + 1, (self.n_train // batch_size) + 1, loss.item()))
+
+    def before_run(self):
+        if isinstance(self.model, MyDataParallel):
+            self.model = self.model.module
+
+    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(tqdm(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')
+        self.model.train()
+        return gradients
+
+    def finish_run(self):
+        # self.model.embedding_recorder.record_embedding = True  # recording embedding vector
+
+        gradients = self.calc_gradient()
+        self.norm_matrix[:,0] = np.linalg.norm(gradients,axis=1)
+
+
+
+        # embedding_dim = self.model.get_last_layer().in_features
+        # data_loader = self.select_dm(self.dst_train, None, is_train=False)
+        # sample_num = self.n_train
+        #
+        # for i, batch in enumerate(data_loader):
+        #     self.optim.zero_grad()
+        #     image, target,batch_inds = batch['img'].cuda(), batch['label'].cuda(), batch['index'].cuda()
+        #
+        #     outputs = self.model(image)
+        #     loss = self.criterion(outputs.requires_grad_(True),
+        #                           targets.to(self.args.device)).sum()
+        #     batch_num = targets.shape[0]
+        #     with torch.no_grad():
+        #         bias_parameters_grads = torch.autograd.grad(loss, outputs)[0]
+        #         self.norm_matrix[i * self.args.selection_batch:min((i + 1) * self.args.selection_batch, sample_num),
+        #         self.cur_repeat] = torch.norm(torch.cat([bias_parameters_grads, (
+        #                 self.model.embedding_recorder.embedding.view(batch_num, 1, embedding_dim).repeat(1,
+        #                                      self.args.num_classes, 1) * bias_parameters_grads.view(
+        #                                      batch_num, self.args.num_classes, 1).repeat(1, 1, embedding_dim)).
+        #                                      view(batch_num, -1)], dim=1), dim=1, p=2)
+        #
+        # self.model.train()
+
+
+    def select(self, **kwargs):
+        # Initialize a matrix to save norms of each sample on idependent runs
+        self.norm_matrix = np.zeros([self.n_train, self.repeat])
+
+        # for self.cur_repeat in range(self.repeat):
+        self.run()
+            # self.random_seed = self.random_seed + 5
+
+        self.norm_mean = np.mean(self.norm_matrix, axis=1)
+        if not self.balance:
+            top_examples = self.train_indx[np.argsort(self.norm_mean)][::-1][:self.coreset_size]
+        else:
+            top_examples = np.array([], dtype=np.int64)
+            for c in tqdm(range(self.num_classes)):
+                c_indx = self.train_indx[self.dst_train_label == c]
+                budget = round(self.fraction * len(c_indx))
+                top_examples = np.append(top_examples, c_indx[np.argsort(self.norm_mean[c_indx])[::-1][:budget]])
+
+        return {"indices": top_examples, "scores": self.norm_mean}