add GEM

Daikon-Sun · Daikon-Sun · commit cac7e33d7d5c · 2019-11-12T13:11:40.000+08:00
diff --git a/README.md b/README.md
@@ -50,7 +50,7 @@ In order to unify the format of all the dataset, we first ran the code in https:
 3. Make a copy of `env.example` and save it as `env`. In `env`, set the value of DATA_DIR as `data directory` and set the value of  MODEL_ROOT_DIR as `model directory`.
 4. Before training or testing, load DATA_DIR and MODEL_ROOT_DIR variables into shell environment by the following command:
    ```bash 
-   source env
+   source ./env
    ```
 
 ## Training and Testing
@@ -115,4 +115,5 @@ After running testing program, the metrics: `metrics.json` will be dumped in the
 ## Acknowledgements:
 - We use the language model offered by [transformers](https://github.com/huggingface/transformers), a great state-of-the-art natural language processing models library by Thomas Wolf et al.
 - The implementation of MAS refer to [MAS-Memory-Aware-Synapses](https://github.com/rahafaljundi/MAS-Memory-Aware-Synapses), a great Memory Aware Synapses method implementation code by Aljundi R. et al.
+- The implementation of GEM refer to [GradientEpisodicMemory](https://github.com/facebookresearch/GradientEpisodicMemory), a great Gradient Episodic Memory method implementation code by Lopez-Paz, David et al.
 - Data format conversion refer to [decaNLP](https://github.com/salesforce/decaNLP), a great The Natural Language Decathlon: Multitask Learning as Question Answering implementation code by Bryan McCann et al.
diff --git a/settings.py b/settings.py
@@ -20,6 +20,7 @@
     "lll": 0.35,
     "ewc": 0.30,
     "mas": 0.18,
+    "gem": 0.50,
 }
 TURING_ARCHS = {'Tesla V100', '2080 Ti'}
 MODEL_CLASSES = {
@@ -53,12 +54,13 @@ def parse_args():
     parser.add_argument("--model_name", type=str, default="gpt2", choices=["gpt2", "openai-gpt"])
     parser.add_argument("--n_gpus", type=int, default=1)
     parser.add_argument("--n_train_epochs", type=int, default=3)
+    parser.add_argument("--dynamic_epochs", action="store_true")
     parser.add_argument("--n_warmup_ratio", type=float, default=0.005)
     parser.add_argument("--n_workers", type=int, default=4)
     parser.add_argument("--use_sep", action="store_true")
     parser.add_argument("--reg_lambda", type=float, default=1.)
     parser.add_argument("--seed", type=int, default=42)
-    parser.add_argument("--seq_train_type", type=str, default="lll", choices=["lll","finetune","multitask","mas","ewc"])
+    parser.add_argument("--seq_train_type", type=str, default="lll", choices=["lll","finetune","multitask","mas","ewc","gem"])
     parser.add_argument("--tasks", nargs='+', default=["squad2"])
     parser.add_argument("--skip_tasks", nargs='+')
     parser.add_argument("--temperature_lm", type=float, default=1.0)
@@ -71,6 +73,7 @@ def parse_args():
     parser.add_argument("--top_p_qa", type=float, default=0.)
     parser.add_argument("--train_batch_size", type=int, default=0)
     parser.add_argument("--weight_decay", type=float, default=0.01)
+    parser.add_argument("--qp_margin", type=float, default=0.5)
     args = parser.parse_args()
 
     if args.debug:
@@ -141,9 +144,14 @@ def parse_args():
     elif args.unbound:
         pass
     else:
-        data_sizes = {task: data_attrs[task]["train"]["data_size"] for task in args.tasks}
-        max_total_data_size = max(data_sizes.values()) * args.n_train_epochs
-        args.n_train_epochs = {d[0]: min(args.max_n_epochs, max_total_data_size//d[1]) for d in data_sizes.items()}
+        if "gem" in args.seq_train_type:
+            args.memory_data = []
+        if args.dynamic_epochs:
+            data_sizes = {task: data_attrs[task]["train"]["data_size"] for task in args.tasks}
+            max_total_data_size = max(data_sizes.values()) * args.n_train_epochs
+            args.n_train_epochs = {d[0]: min(args.max_n_epochs, max_total_data_size//d[1]) for d in data_sizes.items()}
+        else:
+            args.n_train_epochs = {task: args.n_train_epochs for task in args.tasks}
 
     return args, model_config, model_class, tokenizer, config_class, special_token_ids, special_tokens, data_attrs, tokens_weight
 
diff --git a/train.py b/train.py
@@ -31,6 +31,10 @@ def train(task_ids, model):
         prev_task = args.tasks[task_ids[0]-1]
         with torch.no_grad():
             create_extra_data(tasks[0], prev_task, model, train_extra_data)
+    elif "gem" in args.seq_train_type and task_ids[0] > 0: 
+        get_real_data(tasks[0], train_extra_data, accum=False, encode=True)
+        args.memory_data.append(train_extra_data)
+        train_extra_data = []
     logger.info('extra training data size: {}'.format(len(train_extra_data)))
 
     if not model:
@@ -90,7 +94,8 @@ def train(task_ids, model):
     max_train_batch_size = max(len(train_qadata) // args.min_n_steps, args.min_batch_size)
     train_dataloader = create_dataloader(train_qadata, "train", max_train_batch_size)
     if not args.unbound and args.seq_train_type != "multitask":
-        n_train_epochs = TASK_DICT[tasks[0]]["n_train_epochs"]
+        #n_train_epochs = TASK_DICT[tasks[0]]["n_train_epochs"]
+        n_train_epochs = args.n_train_epochs[tasks[0]]
     else:
         n_train_epochs = args.n_train_epochs['_'.join(tasks)]
     n_train_optimization_steps = len(train_qadata) * n_train_epochs
@@ -104,6 +109,16 @@ def train(task_ids, model):
         {'params': [p for n, p in param_optimizer if any(nd in n for nd in no_decay)], 'weight_decay': 0.0}
     ]
 
+    if "gem" in args.seq_train_type:
+        model.task_id = task_ids[0]
+        if not hasattr(model, "grad_dims"):
+            model.grad_dims = []
+            for param in model.parameters():
+                model.grad_dims.append(param.data.numel())
+        if not hasattr(model, "grads"):
+            model.grads = torch.zeros(sum(model.grad_dims),len(args.tasks))
+            model.grads = model.grads.cuda()
+
     if args.seq_train_type in REG_TYPE_KEYS:
         optimizer = Weight_Regularized_AdamW(optimizer_grouped_parameters, lr=args.learning_rate, eps=args.adam_epsilon)
     else:
@@ -124,6 +139,8 @@ def train(task_ids, model):
 
     tot_n_steps = 0
     train_once = TrainStep(model, optimizer, scheduler)
+    if "gem" in args.seq_train_type and task_ids[0] != 0:
+        gem_step = GEMStep(model, parallel_model, train_loss_fct, optimizer)
     model.train()
     for ep in range(n_train_epochs):
         cum_loss, cum_qa_loss, cum_lm_loss, cur_n_inputs = 0, 0, 0, 0
@@ -139,6 +156,8 @@ def train(task_ids, model):
 
             losses = get_losses(parallel_model, cqa, Y, gen_X, gen_Y, train_loss_fct)
             loss = sum(losses)
+            if "gem" in args.seq_train_type and task_ids[0] != 0:
+                gem_step(task_ids[0])
             train_once(loss, n_inputs)
 
             qa_loss = losses[0].item() * n_inputs
diff --git a/utils.py b/utils.py
@@ -17,6 +17,8 @@
 from settings import TOKENIZER, LEN_FACTOR, DATA_ATTRS, MEMORY_FACTOR, MODEL_CONFIG, MODEL_CLASS
 from multiprocessing import Pool
 import sys
+import time
+import quadprog
 import io
 sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding="UTF-8")
 logger = logging.getLogger(__name__)
@@ -164,6 +166,8 @@ def __init__(self, data_paths, data_type, gen_token, extra_data=[]):
 
         data = []
         for data_path in data_paths:
+            if not data_path:
+                continue
             with open(data_path, "r") as f:
                 raw_ds = json.load(f)
             raw_ds = map(lambda x: x["paragraphs"], raw_ds["data"])
@@ -174,7 +178,7 @@ def __init__(self, data_paths, data_type, gen_token, extra_data=[]):
         
         self.data = []
         self.max_a_len = 0
-        if len(data_paths)==1 and ('wiki' in data_paths[0] or 'woz' in data_paths[0]):
+        if len(data_paths)==1 and data_paths[0] is not None and ('wiki' in data_paths[0] or 'woz' in data_paths[0]):
             #data = self._sort_by_index(data)
             #args.n_workers = 1
             if 'wiki' in data_paths[0]:
@@ -183,7 +187,8 @@ def __init__(self, data_paths, data_type, gen_token, extra_data=[]):
                 answers_file = "woz.en_answers.json" 
             with open(os.path.join(args.data_dir,answers_file),"r") as f:
                 self.answers = json.load(f)
-        self.data_tokenization(data)
+        if len(data) > 0:
+            self.data_tokenization(data)
 
         if len(extra_data) > 0:
             extra_data = map(lambda x: self.etl_single_extra_data(x), extra_data)
@@ -345,11 +350,26 @@ def __call__(self, loss, scheduler_steps):
             self.optimizer.backward(loss, update_master_grads=False)
         else:
             loss.backward()
+
         if not args.fp32:
             self.optimizer.update_master_grads()
             self.optimizer.clip_master_grads(args.max_grad_norm)
         else:
             torch.nn.utils.clip_grad_norm_(self.model.parameters(), args.max_grad_norm)
+
+        if "gem" in args.seq_train_type and self.model.task_id >0: 
+            store_grad(self.model.parameters, self.model.grads, self.model.grad_dims,self.model.task_id)
+            indx = torch.cuda.LongTensor([i for i in range(self.model.task_id)])
+            dotp = torch.mm(self.model.grads[:, self.model.task_id].unsqueeze(0),
+                            self.model.grads.index_select(1, indx))
+            if (dotp < 0).sum() != 0:
+                project2cone2(self.model.grads[:, self.model.task_id].unsqueeze(1),
+                              self.model.grads.index_select(1, indx), args.qp_margin)
+                # copy gradients back
+                overwrite_grad(self.model.parameters,
+                               self.model.grads[:, self.model.task_id],
+                               self.model.grad_dims)
+            
         if args.seq_train_type in args.REG_TYPE_KEYS:
             self.optimizer.step(self.model.reg_params)
         else:
@@ -360,6 +380,58 @@ def __call__(self, loss, scheduler_steps):
         self.optimizer.zero_grad()
 
 
+class GEMStep:
+    def __init__(self, model, parallel_model, train_loss_fct, optimizer):
+        self.model = model
+        self.parallel_model = parallel_model
+        self.train_loss_fct = train_loss_fct
+        self.optimizer = optimizer
+
+    def __call__(self,current_task_id):
+        for past_task_id, md in enumerate(args.memory_data):
+            # Not saving current task's grads.
+            if past_task_id >= current_task_id: return
+            qadata = QADataset(None, "test", "gen", md)[:90]
+            dataloader = create_dataloader(qadata, "test")
+            grads_tmp = torch.zeros(sum(self.model.grad_dims),).cuda()
+            if not args.fp32:
+                grads_tmp = grads_tmp.half() 
+            for _, _, cqa, _, Y, gen_X, gen_Y in dataloader:
+                #CHECK
+                n_inputs = sum(_cqa.shape[0] for _cqa in cqa)
+                self.optimizer.zero_grad()
+                for i in range(len(cqa)):
+                    cqa[i] = (cqa[i].to(args.device_ids[i]),)
+                    Y[i] = Y[i].to(args.device_ids[i])
+                    gen_X[i] = (gen_X[i].to(args.device_ids[i]),)
+                    gen_Y[i] = gen_Y[i].to(args.device_ids[i])
+
+                losses = get_losses(self.parallel_model, cqa, Y, gen_X, gen_Y, self.train_loss_fct)
+                loss = sum(losses)
+                if not args.fp32:
+                    self.optimizer.backward(loss, update_master_grads=False)
+                else:
+                    loss.backward()
+
+                if not args.fp32:
+                    #copy fp16 grads to fp32 grads  
+                    self.optimizer.update_master_grads()
+                    self.optimizer.clip_master_grads(args.max_grad_norm)
+                else:
+                    torch.nn.utils.clip_grad_norm_(self.model.parameters(), args.max_grad_norm)
+                i = 0
+                for param in self.model.parameters():
+                    if param.grad is not None:
+                        beg = 0 if i == 0 else sum(self.model.grad_dims[:i])
+                        end = sum(self.model.grad_dims[:i+1])
+                        grads_tmp[beg: end] += param.grad.data.view(-1)*n_inputs
+                    i += 1
+
+            grads_tmp /= len(qadata)
+            self.model.grads[:, past_task_id].copy_(grads_tmp)
+            self.optimizer.zero_grad()
+
+
 class DynamicBatchSampler(Sampler):
     def __init__(self, dataset, data_type, max_batch_size):
         self.dataset = dataset
@@ -523,11 +595,15 @@ def parse_single_real_data(data,task):
     return data
 
 
-def get_real_data(task, train_extra_data):
+def get_real_data(task, train_extra_data, accum=True, encode=True):
     task_idx = args.tasks.index(task)
-    prev_tasks = args.tasks[:task_idx]
     gen_size = DATA_ATTRS[task]["train"]["data_size"]
-    gen_size = int(np.ceil(gen_size * args.gen_lm_sample_percentage))//len(prev_tasks)
+    if accum:
+        prev_tasks = args.tasks[:task_idx]
+        gen_size = int(np.ceil(gen_size * args.gen_lm_sample_percentage))//len(prev_tasks)
+    else:
+        prev_tasks = [args.tasks[task_idx-1]]
+        gen_size = int(gen_size * args.gen_lm_sample_percentage)
 
     datum = []
     for prev_task in prev_tasks:
@@ -537,11 +613,13 @@ def get_real_data(task, train_extra_data):
         for i in indices:
             d = parse_single_real_data(data[i],prev_task)
             datum.append(d)
-            train_extra_data.append(TOKENIZER.encode(d))
+            if encode:
+                train_extra_data.append(TOKENIZER.encode(d))
         
     model_dir = get_model_dir([prev_task])
     dump_path = os.path.join(model_dir,"real.csv")
     write_extra_data(dump_path, datum)
+    return dump_path
 
 
 def read_extra_data(gen_path, train_extra_data):
@@ -728,3 +806,39 @@ def get_split_indices(data_sizes,chunk_sizes):
                 chunk_sizes.pop(0)
                 i+=1
     return records
+
+
+def store_grad(get_ps, grads, grad_dims, task_id): 
+    i = 0
+    for param in get_ps():
+        if param.grad is not None:
+            beg = 0 if i == 0 else sum(grad_dims[:i])
+            end = sum(grad_dims[:i+1])
+            grads[beg: end, task_id].copy_(param.grad.data.view(-1))
+        i += 1
+
+
+def overwrite_grad(pp, newgrad, grad_dims):
+    cnt = 0
+    for param in pp():
+        if param.grad is not None:
+            beg = 0 if cnt == 0 else sum(grad_dims[:cnt])
+            en = sum(grad_dims[:cnt + 1])
+            this_grad = newgrad[beg: en].contiguous().view(
+                param.grad.data.size())
+            param.grad.data.copy_(this_grad)
+        cnt += 1
+
+
+def project2cone2(gradient, memories, margin=0.5, eps=1e-3):
+    memories_np = memories.cpu().t().double().numpy()
+    gradient_np = gradient.cpu().contiguous().view(-1).double().numpy()
+    t = memories_np.shape[0]
+    P = np.dot(memories_np, memories_np.transpose())
+    P = 0.5 * (P + P.transpose()) + np.eye(t) * eps
+    q = np.dot(memories_np, gradient_np) * -1
+    G = np.eye(t)
+    h = np.zeros(t) + margin
+    v = quadprog.solve_qp(P, q, G, h)[0]
+    x = np.dot(v, memories_np) + gradient_np
+    gradient.copy_(torch.Tensor(x).view(-1, 1))
