Upload New File

NLU/trainer.py

+from optim import Adam, NoamOpt
+import torch
+import os
+import torch.nn as nn
+import torch.distributed
+# import torch._tensor
+from dataset import PadBatchSeq
+from torch.utils.data import DataLoader
+from torch.utils.tensorboard import SummaryWriter
+from tqdm import tqdm
+
+
+class Trainer:
+    def __init__(self, args, model, tokz, train_dataset, valid_dataset,
+                 log_dir, logger, device=torch.device('cuda'), valid_writer=None, distributed=False):
+        self.config = args
+        self.device = device
+        self.logger = logger
+        self.log_dir = log_dir
+        self.tokz = tokz
+        self.rank = torch.distributed.get_rank() if distributed else -1
+        self.train_writer = SummaryWriter(os.path.join(log_dir, 'train'))
+        if valid_writer is None:
+            self.valid_writer = SummaryWriter(os.path.join(log_dir, 'valid'))
+        else:
+            self.valid_writer = valid_writer
+        self.model = model.to(device, non_blocking=True)
+        self.criterion = nn.CrossEntropyLoss(ignore_index=tokz.pad_token_id, reduction='none').to(device)
+
+        base_optimizer = Adam(self.model.parameters(), lr=self.config.lr, weight_decay=0.01)
+        if hasattr(self.model, 'config'):
+            self.optimizer = NoamOpt(self.model.config.hidden_size, 0.1, self.config.lr_warmup, base_optimizer)
+        else:
+            self.optimizer = NoamOpt(self.model.module.config.hidden_size, 0.1, self.config.lr_warmup, base_optimizer)
+
+        self.train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset) if distributed else torch.utils.data.RandomSampler(train_dataset)
+        self.valid_sampler = torch.utils.data.distributed.DistributedSampler(valid_dataset) if distributed else None
+
+        self.train_dataloader = DataLoader(
+            train_dataset, sampler=self.train_sampler, batch_size=self.config.bs, num_workers=self.config.n_jobs, pin_memory=True,
+            collate_fn=PadBatchSeq(self.tokz.pad_token_id))
+
+        self.valid_dataloader = DataLoader(
+            valid_dataset, sampler=self.valid_sampler, batch_size=self.config.bs, num_workers=self.config.n_jobs, pin_memory=True,
+            collate_fn=PadBatchSeq(self.tokz.pad_token_id))
+
+    def state_dict(self):
+        return self.model.state_dict()
+        
+    def load_state_dict(self, state_dict):
+        self.model.load_state_dict(state_dict)
+
+    def _eval_train(self, epoch):
+        self.model.train()
+
+        intent_loss, slot_loss, intent_acc, slot_acc, step_count = 0, 0, 0, 0, 0
+        total = len(self.train_dataloader)
+        if self.rank in [-1, 0]:
+            TQDM = tqdm(enumerate(self.train_dataloader), desc='Train (epoch #{})'.format(epoch),
+                        dynamic_ncols=True, total=total)
+        else:
+            TQDM = enumerate(self.train_dataloader)
+
+        for i, data in TQDM:
+            #######################################################
+            # TODO: Complete the following function.
+            #       The following code should preform the training of the model
+            #       You can implement this function with the following steps:
+            #       1. Pass the input to GPU by calling data.to(self.device)
+            #       2. Forward the input to the model
+            #       3. Compute the loss (remember to divide the loss with self.config.batch_split to enable gradient accumulation)
+            #       4. Backward the loss
+            #       5. Update the parameters
+            #       6. Evaluate the model (by calling _eval_test) every `self.config.eval_steps` steps
+            #######################################################
+            # 1. Pass the input to GPU by calling data.to(self.device)
+            text = data['utt'].to(self.device, non_blocking=True)
+            intent_labels = data['intent'].to(self.device, non_blocking=True)
+            slot_labels = data['slot'].to(self.device, non_blocking=True)
+            mask = data['mask'].to(self.device, non_blocking=True)
+            token_type = data['token_type'].to(self.device, non_blocking=True)
+
+            # 2. Forward the input to the model
+            intent_logits, slot_logits = self.model(input_ids=text,
+                                                    attention_mask=mask,
+                                                    token_type_ids=token_type)
+
+            # 3. Compute the loss (remember to divide the loss with self.config.batch_split to enable gradient accumulation)
+            batch_intent_loss = self.criterion(intent_logits, intent_labels).mean()
+            batch_slot_loss = self.criterion(slot_logits.view(-1, slot_logits.shape[-1]), slot_labels.view(-1)).mean()
+            slot_mask = 1 - slot_labels.eq(self.tokz.pad_token_id).float()
+            batch_slot_loss = (batch_slot_loss * slot_mask.view(-1)).sum() / slot_mask.sum()
+
+            batch_loss = batch_intent_loss + batch_slot_loss
+            batch_intent_acc = (torch.argmax(intent_logits, dim=-1) == intent_labels).float().mean()
+            batch_slot_acc = (torch.argmax(slot_logits, dim=-1) == slot_labels)
+            batch_slot_acc = torch.sum(batch_slot_acc * slot_mask) / torch.sum(slot_mask)
+
+            # 4. Backward the loss
+            full_loss = batch_loss / self.config.batch_split
+            full_loss.backward()
+
+            intent_loss += batch_intent_loss.item()
+            slot_loss += batch_slot_loss.item()
+            intent_acc += batch_intent_acc.item()
+            slot_acc += batch_slot_acc.item()
+            step_count += 1
+
+            # 5. Update the parameters
+            curr_step = self.optimizer.curr_step()
+            lr = self.optimizer.param_groups[0]['lr']
+            if (i + 1) % self.config.batch_split == 0:
+                self.optimizer.step()
+                self.optimizer.zero_grad()
+
+                intent_loss /= step_count
+                slot_loss /= step_count
+                intent_acc /= step_count
+                slot_acc /= step_count
+
+                if self.rank in [-1, 0]:
+                    self.train_writer.add_scalar('loss/intent_loss', intent_loss, curr_step)
+                    self.train_writer.add_scalar('loss/slot_loss', slot_loss, curr_step)
+                    self.train_writer.add_scalar('acc/intent_acc', intent_acc, curr_step)
+                    self.train_writer.add_scalar('acc/slot_acc', slot_acc, curr_step)
+                    TQDM.set_postfix({'intent_loss': intent_loss,
+                                      'intent_acc': intent_acc,
+                                      'slot_loss': slot_loss,
+                                      'slot_acc': slot_acc})
+                    intent_loss, slot_loss, intent_acc, slot_acc, step_count = 0, 0, 0, 0, 0
+
+                # 6. Evaluate the model (by calling _eval_test) every `self.config.eval_steps` steps
+                if curr_step % self.config.eval_steps == 0:
+                    self._eval_test(epoch=epoch, step=curr_step)
+
+    def _eval_test(self, epoch, step):
+        self.model.eval()
+        with torch.no_grad():
+            dev_intent_loss = torch.tensor(0.0, dtype=torch.float32, device=self.device)
+            dev_slot_loss = torch.tensor(0.0, dtype=torch.float32, device=self.device)
+            dev_intent_acc = torch.tensor(0.0, dtype=torch.float32, device=self.device)
+            dev_slot_acc = torch.tensor(0.0, dtype=torch.float32, device=self.device)
+            count = torch.tensor(0.0, dtype=torch.float32, device=self.device)
+
+            for data in self.valid_dataloader:
+                text = data['utt'].to(self.device, non_blocking=True)
+                intent_labels = data['intent'].to(self.device, non_blocking=True)
+                slot_labels = data['slot'].to(self.device, non_blocking=True)
+                mask = data['mask'].to(self.device, non_blocking=True)
+                token_type = data['token_type'].to(self.device, non_blocking=True)
+
+                intent_logits, slot_logits = self.model(input_ids=text, attention_mask=mask, token_type_ids=token_type)
+                
+                batch_intent_loss = self.criterion(intent_logits, intent_labels)
+                batch_slot_loss = self.criterion(slot_logits.view(-1, slot_logits.shape[-1]), slot_labels.view(-1))
+                slot_mask = 1 - slot_labels.eq(self.tokz.pad_token_id).float()
+                batch_slot_loss = (batch_slot_loss * slot_mask.view(-1)).view(text.shape[0], -1).sum(dim=-1) / slot_mask.sum(dim=-1)
+                
+                dev_intent_loss += batch_intent_loss.sum()
+                dev_slot_loss += batch_slot_loss.sum()
+
+                batch_intent_acc = (torch.argmax(intent_logits, dim=-1) == intent_labels).sum()
+                batch_slot_acc = (torch.argmax(slot_logits, dim=-1) == slot_labels)
+                batch_slot_acc = torch.sum(batch_slot_acc * slot_mask, dim=-1) / torch.sum(slot_mask, dim=-1)
+
+                dev_intent_acc += batch_intent_acc
+                dev_slot_acc += batch_slot_acc.sum()
+                count += text.shape[0]
+
+            if self.rank != -1:
+                torch.distributed.all_reduce(dev_intent_loss, op=torch.distributed.reduce_op.SUM)
+                torch.distributed.all_reduce(dev_slot_loss, op=torch.distributed.reduce_op.SUM)
+                torch.distributed.all_reduce(dev_intent_acc, op=torch.distributed.reduce_op.SUM)
+                torch.distributed.all_reduce(dev_slot_acc, op=torch.distributed.reduce_op.SUM)
+                torch.distributed.all_reduce(count, op=torch.distributed.reduce_op.SUM)
+
+            dev_intent_loss /= count
+            dev_slot_loss /= count
+            dev_intent_acc /= count
+            dev_slot_acc /= count
+
+            if self.rank in [-1, 0]:
+                self.valid_writer.add_scalar('loss/intent_loss', dev_intent_loss, step)
+                self.valid_writer.add_scalar('loss/slot_loss', dev_slot_loss, step)
+                self.valid_writer.add_scalar('acc/intent_acc', dev_intent_acc, step)
+                self.valid_writer.add_scalar('acc/slot_acc', dev_slot_acc, step)
+                log_str = 'epoch {:>3}, step {}'.format(epoch, step)
+                log_str += ', dev_intent_loss {:>4.4f}'.format(dev_intent_loss)
+                log_str += ', dev_slot_loss {:>4.4f}'.format(dev_slot_loss)
+                log_str += ', dev_intent_acc {:>4.4f}'.format(dev_intent_acc)
+                log_str += ', dev_slot_acc {:>4.4f}'.format(dev_slot_acc)
+                self.logger.info(log_str)
+
+        self.model.train()
+
+    def train(self, start_epoch, epochs, after_epoch_funcs=[], after_step_funcs=[]):
+        for epoch in range(start_epoch + 1, epochs):
+            self.logger.info('Training on epoch'.format(epoch))
+            if hasattr(self.train_sampler, 'set_epoch'):
+                self.train_sampler.set_epoch(epoch)
+            self._eval_train(epoch)
+            for func in after_epoch_funcs:
+                func(epoch, self.device)
\ No newline at end of file