[Pytorch Lightning] BART를 이용해 Text를 요약해보자 - 4 (완)

• 1편 : 프로젝트 요약 및 LightningModule 설계
• 2편 : Dataset 구축 및 DataLoader 구축
• 3편 : logger 작성, wandb 연동 및 확인
• 4편 : 모델 로딩 및 실행결과(wandb) 확인

이번에는 작성한 코드를 이용해서 모델 로딩 및 실행결과(wandb) 확인을 해보겠습니다.

환경설정 및 초기화

사용할 모델 및 토크나이저 선택

1편을 보시면 알겠지만, 저는 한국어로 입력했을 때, 영어로 출력하도록 Task를 만들었습니다. 그래서 저는 한국어-영어 쌍으로 훈련된 모델을 허깅페이스에서 가져와서 사용했습니다. 궁금하신분은 링크를 클릭해보시면 됩니다. 간단히 말하면 한국어를 입력했을 때, 영어로 출력하는 KoBART 모델입니다.

model_id = 'chunwoolee0/circulus-kobart-en-to-ko'

model = AutoModelForSeq2SeqLM.from_pretrained(args['model_path'], cache_dir='/home/hgjeong/hdd1/hub')
tokenizer = AutoTokenizer.from_pretrained(args['tokenizer'])

 

하이퍼 파라미터 설정

args = dict()
args['output_dir'] ='20240512'
args['model_path'] = 'chunwoolee0/circulus-kobart-en-to-ko'
args['tokenizer'] = 'chunwoolee0/circulus-kobart-en-to-ko'

args['gpus'] = '1'
args['epochs'] = 100
args['max_learning_rate'] = 2e-5
args['min_learning_rate'] = 1e-6
args['warmup_rate'] = 0.1
args['max_seq_len'] = 128
args['batch_size_train'] =64
args['batch_size_valid'] =8
args['logging_interval'] =100
args['evaluate_interval'] =1.0
args['seed'] =93
args['wandb_project'] ='blog'
args['accumulate_grad_batches'] = 1

• output_dir: 모델과 출력 파일들을 저장할 디렉토리 경로 ('20240512')
• model_path: 사용할 모델의 경로
• tokenizer: 사용할 토크나이저의 경로
• gpus: 훈련에 사용할 GPU 수
• epochs: 훈련할 총 에폭 수 (100)
• max_learning_rate: 최대 학습률 (2e-5)
• min_learning_rate: 최소 학습률 (1e-6)
• warmup_rate: 학습률 웜업 비율 (0.1), 학습 초기에 학습률을 점진적으로 증가시키는 비율
• max_seq_len: 입력 시퀀스의 최대 길이 (128), 너무 긴 입력은 잘라내기 처리
• batch_size_train: 훈련 데이터의 배치 크기 (64)
• batch_size_valid: 검증 데이터의 배치 크기 (8)
• logging_interval: 로깅 간격 (100), 이 간격으로 훈련 로그를 기록
• evaluate_interval: 평가 간격 (1.0), 각 에폭마다 모델 성능을 평가
• seed: 랜덤 시드 (93)
• wandb_project: Weights & Biases에서 사용할 프로젝트 이름 ('blog'), 실험 관리 및 모니터링을 위해 설정
• accumulate_grad_batches: 그래디언트 누적 배치 수

 

데이터 로딩

train_dataloader, valid_dataloader = OrderDataLoader("../test.jsonl", tokenizer, args['batch_size_train'], args['batch_size_valid'], args['max_seq_len'])

 

훈련 준비(pl.LightningModule 및 pl.Trainer 초기화)

model = AutoModelForSeq2SeqLM.from_pretrained(args['model_path'])

lightning_module = StoryModule(
    model,
    args['output_dir'],
    total_steps,
    args['max_learning_rate'],
    args['min_learning_rate'],
    args['warmup_rate'],
)

하이퍼 파라미터를 Module에 입력한다.

trainer = pl.Trainer(
    strategy="auto",
    accelerator="gpu",
    logger=train_loggers,
    max_epochs=args['epochs'],
    log_every_n_steps=args['logging_interval'],
    val_check_interval=args['evaluate_interval'],
    accumulate_grad_batches=args['accumulate_grad_batches'],
    callbacks=[LearningRateMonitor(logging_interval="step")],
    devices=args['gpus'],
)

하이퍼 파라미터를 Trainer에 입력한다.

 

훈련 시작

trainer.fit(lightning_module, train_dataloader, valid_dataloader)

train_dataloader와 valid_dataloader를 통해 훈련하고 검증을 진행한다.

 

실행 및 결과확인(wandb)

wandb를 통해 loss(train)와 loss(valid)가 둘 다 epoch가 늘어나면서 줄어드는 것을 확인할 수 있다.

전체 코드

https://github.com/Capstone-Amigo/OrderIt-AI/tree/main/ai

이곳에서 확인할 수 있다. BART 훈련에 도움이 되었으면 하는 마음에 다음과 같이 블로그 글을 적어보았다. 다들 이 글을 통해 시간을 아꼈으면 좋겠다. 감사합니다!