[음성 인식 프로젝트] Whisper 파인튜닝 하기

본 포스팅에서는 음성을 텍스트로 변환해주는 OpenAI의 Whisper 모델을 구음장애 환자 데이터로 파인튜닝 해보겠습니다. 여기서 음성인식 모델은 자동음성인식(ASR, automatic speech recognition) 기술로, 인간의 음성을 텍스트로 변환하는 기술입니다.

 

Whisper 모델

Whispser는 OpenAI에서 개발한 음성 인식 모델로 Whisper large-v3의 경우 약 100만 시간 가량의 음성 데이터를 통해 사전 훈련되었습니다. 논문에 따르면, Whisper는 Transformer sequence-to-sequence 모델로 구성되어 있으며, 음성 인식/음성 번역/음성 활동 감지를 수행할 수 있습니다. 모델의 구조는 아래 그림과 같습니다. Transformer 구조로 인코더-디코더로 이루어진 것을 확인할 수 있습니다.

Radford, A., Kim, J. W., Xu, T., Brockman, G., McLeavey, C., & Sutskever, I. (2023, July). Robust speech recognition via large-scale weak supervision. In  International Conference on Machine Learning  (pp. 28492-28518). PMLR.

 

파인튜닝

전이학습, 미세조정으로 불리는 파인튜닝(fine-tuning)은 대규모 데이터를 통해 학습된 거대 모델을 특정 도메인이나 작업에 맞게 조정하는 과정입니다.
본 프로젝트에서는 한국어 구음장애 환자들을 위한 음성인식 모델을 개발하기 위해 한국어 구음장애 음성에 대해 파인튜닝을 진행하겠습니다. 데이터에 대한 설명과 전처리는 이전 포스팅을 확인해주시면 됩니다.

음성데이터 추출: https://ysg2997.tistory.com/51

[음성 인식 모델 프로젝트] 음성 데이터 시각화 및 특성 추출

음성데이터 전처리: https://ysg2997.tistory.com/52

[음성 인식 모델 프로젝트] 음성 데이터 침묵구간 - 비침묵구간 분리하기

 

학습 과정

미세 조정을 위한 전체 과정은 아래 그림과 같습니다.

 

학습 코드로 살펴보기

전체 과정을 코드와 함께 알아보겠습니다. 구글 코랩 환경에서 진행했습니다.

먼저 GPU를 할당합니다.

gpu_info = !nvidia-smi
gpu_info = '\n'.join(gpu_info)
if gpu_info.find('failed') >= 0:
  print('Not connected to a GPU')
else:
  print(gpu_info)

 

 

라이브러리 설치

!pip install datasets>=2.6.1
!pip install git+https://github.com/huggingface/transformers
!pip install librosa
!pip install evaluate>=0.30
!pip install jiwer
!pip install --upgrade pip
!pip install --upgrade git+https://github.com/huggingface/transformers.git accelerate datasets[audio]

 

라이브러리 불러오기

허깅페이스의 pipeline 클래스를 사용하여 학습을 진행하겠습니다.

import torch
from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor, pipeline
from datasets import load_dataset

# 디바이스 GPU 설정
device = "cuda:0" if torch.cuda.is_available() else "cpu"
torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32

# 모델 이름 설정
model_id = "openai/whisper-large-v3"

# 모델 불러오기
model = AutoModelForSpeechSeq2Seq.from_pretrained(
    model_id, torch_dtype=torch_dtype, low_cpu_mem_usage=True, use_safetensors=True
)
model.to(device)

processor = AutoProcessor.from_pretrained(model_id)

# pipline()으로 설정하기.
pipe = pipeline(
    "automatic-speech-recognition",
    model=model,
    tokenizer=processor.tokenizer,
    feature_extractor=processor.feature_extractor,
    max_new_tokens=128,
    chunk_length_s=30,
    batch_size=16,
    return_timestamps=True,
    torch_dtype=torch_dtype,
    device=device,
)

# 데이터세트 불러오기
dataset = load_dataset("distil-whisper/librispeech_long", "clean", split="validation")
sample = dataset[0]["audio"]

result = pipe(sample)
print(result["text"])

 

모델 불러오기

허깅페이스에 로그인 하겠습니다.

이때, 각자의 허깅페이스 계정의 토큰을 입력해야 합니다.

from huggingface_hub import notebook_login

notebook_login()

 

데이터 불러오기

파인튜닝을 위한 데이터 셋을 불러오겠습니다.

from datasets import load_dataset, DatasetDict

common_voice = DatasetDict()

common_voice["train"] = load_dataset("mozilla-foundation/common_voice_16_1", "ko", split="train+validation", use_auth_token=True)
common_voice["test"] = load_dataset("mozilla-foundation/common_voice_16_1", "ko", split="test", use_auth_token=True)

# 불필요한 데이터 삭제
common_voice = common_voice.remove_columns(["accent", "age", "client_id", "down_votes", "gender", "locale", "path", "segment", "up_votes"])
common_voice

 

불러온 데이터는 아래와 같은 DatasetDict 형태로 저장되어있습니다.

### result ###

DatasetDict({
    train: Dataset({
        features: ['audio', 'sentence', 'variant'],
        num_rows: 636
    })
    test: Dataset({
        features: ['audio', 'sentence', 'variant'],
        num_rows: 282
    })
})

 

데이터 전처리

데이터 전처리에는 아래 두가지 모듈을 사용합니다.

• Whisper Feature Extractor
• Whisper Tokenizer

 

Whisper Feature Extractor

• 음성은 시간에 따라 변하는 1차원 배열로 표현된됩니다. 주어진 시간단계에서 배열의 값은 신호 진폭을 의미합니다.
• 음성 데이터는 연속적이고 무한한 수의 진폭 값을 포함한다.
• 오디오 데이터는 샘플링을 통해 'sample/sec' 단위, 즉 Hz로 측정합니다.
  • 더 높은 sampling rate는 더 많은 신호에 접근하지만, 더 많은 값을 저장하고 있는 것입니다.
• 이때, 오디오 특성에 따라 sampling rate를 일치 시키는 것이 중요합니다.
• Whisper Feature Extractor는 sampling rate가 16kHz인 데이터를 기대합니다.

from transformers import WhisperFeatureExtractor

# feature_extractor 불러오기
feature_extractor = WhisperFeatureExtractor.from_pretrained("openai/whisper-small")

 

WhisperTokenizer

• 텍스트 데이터를 토큰화 합니다.
• 전통적으로 ASR 작업에서는 CTC(Connection Temporal Classification)을 사용합니다.
• 즉, 텍스트 데이터에 대한 CTC 토크나이저가 적용되는 것입니다.
• WhisperTokenizer는 언어와 작업을 지정하기만 하면 됩니다.

from transformers import WhisperTokenizer

tokenizer = WhisperTokenizer.from_pretrained("openai/whisper-small", language="Korean", task="transcribe")

 

위 두가지 전처리를 합친 WhisperProcessor

WhisperProcessor는 Whisper Feature Extractor + WhisperTokenizer를 한번에 매핑합니다.

from transformers import WhisperProcessor

processor = WhisperProcessor.from_pretrained("openai/whisper-small", language="Korean", task="transcribe")

 

위 데이터가 48,000sr 이므로, 16,000으로 변경해줍니다.

from datasets import Audio

common_voice = common_voice.cast_column("audio", Audio(sampling_rate=16000))

 

 

최종적으로 데이터를 불러오고 매핑해주는 코드입니다.

def prepare_dataset(batch):
    # 리샘플링 from 48 to 16kHz
    audio = batch["audio"]

    # audio array -> log-Mel spectrogram
    batch["input_features"] = feature_extractor(audio["array"], sampling_rate=audio["sampling_rate"]).input_features[0]

    # label 추가
    batch["labels"] = tokenizer(batch["sentence"]).input_ids
    return batch
    
    common_voice = common_voice.map(prepare_dataset, remove_columns=common_voice.column_names["train"], num_proc=4)

 

모델 학습

• 파인튜닝 학습에서는 간단히 Trainer()를 사용합니다.
• 평가 지표: WER(단어 오류율)입니다.
  • Compute_metrics 함수를 정의합니다.

 

• DataCollator
  • 음성 모델은 input_features와 레이블을 독립적으로 처리합니다.
  • input_features는 특징 추출기로 처리되고 레이블은 토크나이저로 처리될 것입니다.
• input_features는 이미 30초로 채워져 있고 고정 차원의 log-Mel 스펙트로그램으로 변환되어 있습니다.
• 본 DataCollator는 일괄 처리된 PyTorch 텐서로 변환하는 것뿐만 수행합니다.
• return_tensors=pt와 함께 특징 추출기의 .pad 메소드를 사용하여 이를 수행합니다.
• 반면에 레이블은 당연히 패딩되지 않습니다.
  • 먼저 토크나이저의 .pad 메소드를 사용하여 배치의 최대 길이까지 시퀀스를 채움.
  • 이후, 패딩 토큰은 -100으로 대체되어 손실을 계산할 때 이러한 토큰이 고려되지 않음.

import torch

from dataclasses import dataclass
from typing import Any, Dict, List, Union

@dataclass
class DataCollatorSpeechSeq2SeqWithPadding:
    processor: Any

    def __call__(self, features: List[Dict[str, Union[List[int], torch.Tensor]]]) -> Dict[str, torch.Tensor]:
        # split inputs and labels since they have to be of different lengths and need different padding methods
        # first treat the audio inputs by simply returning torch tensors
        input_features = [{"input_features": feature["input_features"]} for feature in features]
        batch = self.processor.feature_extractor.pad(input_features, return_tensors="pt")

        # get the tokenized label sequences
        label_features = [{"input_ids": feature["labels"]} for feature in features]
        # pad the labels to max length
        labels_batch = self.processor.tokenizer.pad(label_features, return_tensors="pt")

        # replace padding with -100 to ignore loss correctly
        labels = labels_batch["input_ids"].masked_fill(labels_batch.attention_mask.ne(1), -100)

        # if bos token is appended in previous tokenization step,
        # cut bos token here as it's append later anyways
        if (labels[:, 0] == self.processor.tokenizer.bos_token_id).all().cpu().item():
            labels = labels[:, 1:]

        batch["labels"] = labels

        return batch


# Data Collator 이니셜라이즈
data_collator = DataCollatorSpeechSeq2SeqWithPadding(processor=processor)

 

평가지표

import locale
print(locale.getpreferredencoding())

def getpreferredencoding(do_setlocale = True):
    return "UTF-8"
locale.getpreferredencoding = getpreferredencoding

 

아래 코드에서 compute_metrics 함수는 다음과 같이 작동합니다.

1. predict의 label_ids에서 -100을 pad_token_id로 바꿈.
2. 예측 ID와 레이블 ID를 문자열로 디코딩함.
3. 예측 라벨과 참조 라벨 간의 WER을 계산함.

import evaluate

metric = evaluate.load("wer")

def compute_metrics(pred):
    pred_ids = pred.predictions
    label_ids = pred.label_ids

    # replace -100 with the pad_token_id
    label_ids[label_ids == -100] = tokenizer.pad_token_id

    # we do not want to group tokens when computing the metrics
    pred_str = tokenizer.batch_decode(pred_ids, skip_special_tokens=True)
    label_str = tokenizer.batch_decode(label_ids, skip_special_tokens=True)

    wer = 100 * metric.compute(predictions=pred_str, references=label_str)

    return {"wer": wer}

 

학습 argument 설정

이제 학습을 위한 설정을 해줍니다.

• output_dir: 학습된 모델 가중치가 저장될 로컬 디렉토리.
• generation_max_length: 평가 중에 자동적으로 생성할 최대 토큰 수.
• save_steps: 훈련 중에 중간 체크포인트가 저장되고, save_steps 훈련마다 허브에 비동기적 업로드.
• eval_steps: 훈련 단계 중간 체크포인트 평가가 수행.
• report_to: 훈련 로그가 저장될 위치. ex)azure_ml / comet_ml / mlflow / neptune / tensorboard / wandb

from transformers import Seq2SeqTrainingArguments

training_args = Seq2SeqTrainingArguments(
    output_dir="./whisper-small-hi",  # change to a repo name of your choice
    per_device_train_batch_size=16,
    gradient_accumulation_steps=1,  # increase by 2x for every 2x decrease in batch size
    learning_rate=1e-5,
    warmup_steps=500,
    max_steps=4000,
    gradient_checkpointing=True,
    fp16=True,
    evaluation_strategy="steps",
    per_device_eval_batch_size=8,
    predict_with_generate=True,
    generation_max_length=225,
    save_steps=1000,
    eval_steps=1000,
    logging_steps=25,
    report_to=["tensorboard"],
    load_best_model_at_end=True,
    metric_for_best_model="wer",
    greater_is_better=False,
    push_to_hub=True,
)

 

학습 진행

from transformers import Seq2SeqTrainer

trainer = Seq2SeqTrainer(
    args=training_args,
    model=model,
    train_dataset=common_voice["train"],
    eval_dataset=common_voice["test"],
    data_collator=data_collator,
    compute_metrics=compute_metrics,
    tokenizer=processor.feature_extractor,
)

trainer.train()

### result ###

/usr/local/lib/python3.10/dist-packages/torch/utils/checkpoint.py:429: UserWarning: torch.utils.checkpoint: please pass in use_reentrant=True or use_reentrant=False explicitly. The default value of use_reentrant will be updated to be False in the future. To maintain current behavior, pass use_reentrant=True. It is recommended that you use use_reentrant=False. Refer to docs for more details on the differences between the two variants.
  warnings.warn(
`use_cache = True` is incompatible with gradient checkpointing. Setting `use_cache = False`...
 [ 4/4000 00:19 < 10:35:32, 0.10 it/s, Epoch 0.07/100]
Step	Training Loss	Validation Loss