[수정] build_configuration.py 수정, [추가] train.py 학습 추가, pytorch_lightnin…

build_configuration.py

@@ -9,11 +9,12 @@
 def main(args):
 
     tokenizer = BartTokenizer.from_pretrained(args.tokenizer_path)
-    bart_config = BartConfig()
+    bart_config = BartConfig.from_pretrained('facebook/bart-base')
     bart_config.vocab_size = len(tokenizer)
     bart_config.eos_token_id = tokenizer.eos_token_id
     bart_config.bos_token_id = tokenizer.bos_token_id
     bart_config.pad_token_id = tokenizer.pad_token_id
+    bart_config.mask_token_id = tokenizer.mask_token_id
 
     bart_config.save_pretrained(args.config_path)
 

src/data.py

@@ -95,7 +95,8 @@ def train_dataloader(self):
                 dataset=dataset,
                 batch_size=self.batch_size,
                 num_workers=self.num_workers,
-                collate_fn=self.collate_fn)
+                collate_fn=self.collate_fn,
+                drop_last=True)
 
         return dataloader
 
@@ -138,3 +139,4 @@ def main(args):
 
     args = parser.parse_args()
     main(args)
+

src/module.py

@@ -0,0 +1,76 @@
+import pytorch_lightning as pl
+import torch
+import torchvision
+
+from transformers import BartForConditionalGeneration
+from transformers import BartConfig
+
+from pl_bolts.optimizers.lr_scheduler import LinearWarmupCosineAnnealingLR
+from pl_bolts.optimizers.lars_scheduling import LARSWrapper
+
+class BartModule(pl.LightningModule):
+
+    def __init__(self, config, learning_rate,
+                 weight_decay, max_epochs, warmup_epochs):
+        super(BartModule, self).__init__()
+
+        self.save_hyperparameters()
+        self.config = BartConfig.from_pretrained(config)
+        self.model = BartForConditionalGeneration(
+                config=self.config)
+
+        self.train_step = 0
+
+    def forward(self, input_ids, attention_mask=None, labels=None):
+        output = self.model(
+                input_ids=input_ids,
+                attention_mask=attention_mask,
+                labels=labels,
+                return_dict=True, output_hidden_states=True)
+        return output
+
+    def shared_step(self, batch):
+
+        input_ids, attention_mask, labels = batch
+        output = self.forward(
+                input_ids=input_ids,
+                attention_mask=attention_mask,
+                labels=labels)
+        loss = output.loss
+
+        pred = torch.argmax(output.logits, axis=2)
+        pred = pred.eq(labels).view(-1).to(dtype=torch.float)
+        pred = pred.mean()
+        return loss, pred
+
+    def training_step(self, batch, batch_idx):
+
+        self.train_step += 1
+
+        loss, pred = self.shared_step(batch)
+
+        self.logger.experiment.add_scalar('data/train_loss', loss, self.train_step)
+        self.logger.experiment.add_scalar('data/train_pred', pred, self.train_step)
+        self.logger.experiment.add_scalar('data/lr', self.optimizers[0].param_groups[0]['lr'], self.train_step)
+
+        return loss
+
+    def configure_optimizers(self):
+
+        optimizer = torch.optim.Adam(
+                self.parameters(),
+                lr=self.hparams.learning_rate,
+                weight_decay=self.hparams.weight_decay)
+
+        optimizer = LARSWrapper(optimizer)
+
+        scheduler = LinearWarmupCosineAnnealingLR(
+            optimizer,
+            warmup_epochs=self.hparams.warmup_epochs,
+            max_epochs=self.hparams.max_epochs
+        )
+
+        self.optimizers = [optimizer]
+
+        return self.optimizers, [scheduler]
+

train.py

@@ -3,8 +3,72 @@
     To implement code for training your model.
 """
 import pytorch_lightning
+import argparse
 import torch
 
+import pytorch_lightning as pl
+
+from src.module import BartModule
+from src.data import BartDataModule
+from pathlib import Path
+from transformers import BartTokenizer
+
+from pytorch_lightning.callbacks import ModelCheckpoint
+from pytorch_lightning import loggers
+
 pytorch_lightning.seed_everything(777)
 torch.backends.cudnn.deterministic = True
 torch.backends.cudnn.benchmark = False
+
+def main():
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--tokenizer_path', default='tokenizers', type=str)
+    parser.add_argument('--corpus', default='test.txt', type=str)
+    parser.add_argument('--mask_path', default='dataset.json', type=str)
+    parser.add_argument('--config_path', default='kobart', type=str)
+    parser.add_argument('--logger', default='kobart', type=str)
+    parser.add_argument('--learning_rate', default=5e-4, type=float)
+    parser.add_argument('--weight_decay', default=5e-4, type=float)
+    parser.add_argument('--max_epochs', default=100, type=int)
+    parser.add_argument('--warmup_epochs', default=3, type=int)
+    args = parser.parse_args()
+
+    proj_dir = Path()
+    corpus_dir = proj_dir / "corpus"
+    comment_dir = corpus_dir / "comment"
+    source_path = comment_dir / args.corpus
+    mask_path = comment_dir / args.mask_path
+
+    tokenizer = BartTokenizer.from_pretrained(args.tokenizer_path)
+    dm = BartDataModule(
+            source_path=source_path,
+            mask_path=mask_path,
+            tokenizer=tokenizer,
+            batch_size=2,
+            num_workers=1)
+    dm.setup()
+    train_dataloader = dm.train_dataloader()    
+
+    checkpoint_callback = ModelCheckpoint(
+            save_top_k=-1, verbose=True)
+    logger = loggers.TensorBoardLogger(args.logger)
+    model = BartModule(
+            config=args.config_path,
+            learning_rate=args.learning_rate,
+            weight_decay=args.weight_decay,
+            max_epochs=args.max_epochs,
+            warmup_epochs=args.warmup_epochs)
+
+    device_count = torch.cuda.device_count()
+    trainer = pl.Trainer(
+            # gpus=device_count,
+            # distributed_backend='ddp',
+            max_epochs=args.max_epochs,
+            checkpoint_callback=checkpoint_callback,
+            logger=logger)
+
+    trainer.fit(model, train_dataloader)
+
+if __name__ == '__main__':
+    main()