[Pytorch] 내가 쓰는 config 코드

이제 실험할 때 쓰는 구조가 거의 굳어졌다. 

내가 기본적으로 사용하는 config 코드이다. 허깅페이스 예제 코드와 거의 비슷하고 출처도 동일하다.

 

---

• data: 데이터 폴더
  • preprocessing.py
• eda: 각종 분석용 코드 및 시각화 코드용 폴더
• src
  • functions: 3가지 코드는 필수적으로 사용한다. Dataset과 Collator는 data에서 만들어서 사용
    • data.py
    • utils.py
    • metric.py
  • model: 모델에 대한 함수와 train,eval 함수가 들어있는 폴더
    • model.py
    • main_functions.py
• config.py: 모델에 대한 하이퍼파라미터와 경로 등을 저장해서 사용한다. 쉽게 바꿀 수 있다.
• run.py: 실행 파일

---

"""config.py"""
import argparse
import os
from pathlib import Path
import logging

logger = logging.getLogger(__name__)


class Options():
    def __init__(self):
        self.parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
        self.initialize_parser()

    def initialize_parser(self):

        self.parser.add_argument('--do_train', type=bool, default=True, help='Whether training or evaluation')
        self.parser.add_argument("--initialization", type=bool, default=True)
        self.parser.add_argument('--name', type=str, default='',    
                                 help='name of the experiment')
        self.parser.add_argument('--checkpoint_dir', type=str, default='./checkpoint/',
                                 help='models are saved here')
        self.parser.add_argument('--checkpoint_step', type=str, default='', help='checkpoint step')
        self.parser.add_argument('--device', type=str, default="cuda")
        self.parser.add_argument('--seed', type=int, default=0, help="random seed for initialization")

        # pretrained model parameters
        self.parser.add_argument('--pretrained_tokenizer', type=str, default="bert-base",
                                 help="Load pre-trained tokenizer")
        self.parser.add_argument('--pretrained_model', type=str, default="bert-base",
                                 help="Load pre-trained language model")

        # training parameters
        self.parser.add_argument("--train_batch_size", default=32, type=int,
                                 help="Batch size per GPU/CPU for training.")
        self.parser.add_argument("--eval_batch_size", default=16, type=int,
                                 help="Batch size per GPU/CPU for evaluation.")
        self.parser.add_argument('--loss_print_freq', type=int, default=100,
                                 help='print train loss <loss_print_freq> steps during training')
        self.parser.add_argument('--eval_freq', type=int, default=500,
                                 help='evaluate model every <eval_freq> steps during training')
        self.parser.add_argument('--save_freq', type=int, default=500,
                                 help='save model every <save_freq> steps during training')

    def add_optim_options(self):
        self.parser.add_argument('--total_steps', type=int, default=20000)
        self.parser.add_argument('--scheduler_steps', type=int, default=None,
                        help='total number of step for the scheduler, if None then scheduler_total_step = total_step')
        self.parser.add_argument('--accumulation_steps', type=int, default=2)
        self.parser.add_argument('--dropout_rate', type=int, default=0.3)
        self.parser.add_argument('--warmup_steps', type=int, default=0)
        self.parser.add_argument('--lr', type=float, default=5e-5, help='learning rate')
        self.parser.add_argument('--clip', type=float, default=1., help='gradient clipping')
        self.parser.add_argument('--optim', type=str, default='adamw')
        self.parser.add_argument('--adam_epsilon', type=int, default=1e-8)
        self.parser.add_argument('--scheduler', type=str, default='linear')
        self.parser.add_argument('--weight_decay', type=float, default=0.0)

    def add_argument_options(self):
        self.parser.add_argument('--data_dir', type=str,default= './data')

        self.parser.add_argument('--train_data', type=str, default='train.json')
        self.parser.add_argument('--test_data', type=str, default='test.json')     
        self.parser.add_argument('--vocab_data', type=str, default='')
        self.parser.add_argument('--add_vocab', type=bool, default=True)

    def parse(self):
        opt = self.parser.parse_args()
        return opt

사용시에는 다음과 같이 사용한다.

"""run.py"""
import os
import logging
from pathlib import Path
from torchsummaryX import summary

from config import Options     # config 파일 사용

from transformers import AutoConfig, AutoTokenizer, AutoModel

os.environ['CUDA_LAUNCH_BLOCKING'] = "1"
os.environ["CUDA_VISIBLE_DEVICES"] = "0"

if __name__ == "__main__":

    # setting
    options = Options()
    options.add_argument_options()
    options.add_optim_options()
    opt = options.parse()

    # Path & logging
    output_path = Path(opt.checkpoint_dir) / opt.name
    output_path.mkdir(parents=True, exist_ok=True)

    cache_path = output_path / 'cache'
    cache_path.mkdir(parents=True, exist_ok=True)

    logger = logging.getLogger(__name__)
    logger = init_logger(output_path / 'run.log')
    set_seed(opt)

    # Load Pretrained, Dataset
    config = AutoConfig.from_pretrained(opt.pretrained_model)
    pre_trained_model = AutoModel.from_pretrained(opt.pretrained_model, config=config)    
    tokenizer = AutoTokenizer.from_pretrained(opt.pretrained_tokenizer)