wav2vec fine tuning 질문

안녕하세요 wav2vec을 이용해서 fine-tuning을 해보려고 합니다.
허깅페이스에서 kresnik/wav2vec2-large-xlsr-korean 이 모델로 선택했습니다.
마지막 lm_head 부분만 미세 조정을 하려고 하고, vocab을 추가해서 1405 -> 2100 정도로 수정했습니다.

lm_head를 제외한 부분은 requires_grad =False로 하였고 모델에 label, attention_mask를 넣으면 loss를 바로 계산해줘서 그걸로 loss.backward()를 하고 있습니다.

근데 cuda out of memory 문제가 있는데요..

학습하고자 하는 오디오는 0~30초 정도의 길이를 가지고 있고, 현재 gpu의 용량은 24기가입니다.
batch_size = 1 까지는 잘되나, 2 이상부터는 cuda out of memory가 계속 뜹니다.

for epoch in tqdm(range(epochs)):
model.train()
loss_value = 0
valid_loss_value = 0
total_batch_size = 0
valid_total_batch_size = 0

for batch in tqdm(train_dataloader):
optimizer.zero_grad()

audio, text_label, attn_mask = batch

batch_size = audio.shape[0]

out = model(audio.to(device), labels=text_label.to(device), attention_mask=attn_mask.to(device))
           

prediction_ids = torch.argmax(out.logits, dim=-1)
    
pred = processor.batch_decode(prediction_ids)
true_y = processor.batch_decode(text_label)
wer_prob = wer(pred, true_y).item()
wer_record.append(wer_prob)


loss = out.loss

loss.backward()
optimizer.step()

loss_value += loss.item() * batch_size
total_batch_size += batch_size

학습 코드는 이렇습니다.
제 오디오 input이 길어서 문제가 생기는지, 어떻게 해야할지 모르겠습니다..
batch_size=8 로 해서 next(iter(train_datalodaer)를 하고 model에 단순히 넣었을때는 사용하는 메모리가 10000정도 밖에 안되는데 위의 코드로 돌리면 batch_size =1로 해도 거의 24기가 꽉차네요..

out = model(audio.to(device), labels=text_label.to(device), attention_mask=attn_mask.to(device)) 45 # allocated_memory = torch.cuda.memory_allocated() 46 # print(f"Allocated CUDA memory at model : {allocated_memory / 10243:.2f} GB") 47 48 # max_allocated_memory = torch.cuda.max_memory_allocated() 49 # print(f"Max allocated CUDA memory: {max_allocated_memory / 10243:.2f} GB") 52 prediction_ids = torch.argmax(out.logits, dim=-1) File ~/.local/lib/python3.8/site-packages/torch/nn/modules/module.py:1501, in Module._call_impl(self, *args, **kwargs) 1496 # If we don't have any hooks, we want to skip the rest of the logic in 1497 # this function, and just call forward. 1498 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks 1499 or _global_backward_pre_hooks or _global_backward_hooks 1500 or _global_forward_hooks or _global_forward_pre_hooks): -> 1501 return forward_call(*args, **kwargs) 1502 # Do not call functions when jit is used 1503 full_backward_hooks, non_full_backward_hooks = [], [] File ~/.local/lib/python3.8/site-packages/transformers/models/wav2vec2/modeling_wav2vec2.py:1716, in Wav2Vec2ForCTC.forward(self, input_values, attention_mask, output_attentions, output_hidden_states, return_dict, labels) 1706 r""" 1707 labels (torch.LongTensor of shape (batch_size, target_length), optional): 1708 Labels for connectionist temporal classification. Note that target_length has to be smaller or equal to (...) 1711 config.vocab_size - 1]`. 1712 """ 1714 return_dict = return_dict if return_dict is not None else self.config.use_return_dict -> 1716 outputs = self.wav2vec2( 1717 input_values, 1718 attention_mask=attention_mask, 1719 output_attentions=output_attentions, 1720 output_hidden_states=output_hidden_states, 1721 return_dict=return_dict, 1722 ) 1724 hidden_states = outputs[0] 1725 hidden_states = self.dropout(hidden_states) File ~/.local/lib/python3.8/site-packages/torch/nn/modules/module.py:1501, in Module._call_impl(self, *args, **kwargs) 1496 # If we don't have any hooks, we want to skip the rest of the logic in 1497 # this function, and just call forward. 1498 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks 1499 or _global_backward_pre_hooks or _global_backward_hooks 1500 or _global_forward_hooks or _global_forward_pre_hooks): -> 1501 return forward_call(*args, **kwargs) 1502 # Do not call functions when jit is used 1503 full_backward_hooks, non_full_backward_hooks = [], [] File ~/.local/lib/python3.8/site-packages/transformers/models/wav2vec2/modeling_wav2vec2.py:1347, in Wav2Vec2Model.forward(self, input_values, attention_mask, mask_time_indices, output_attentions, output_hidden_states, return_dict) 1342 output_hidden_states = ( 1343 output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states 1344 ) 1345 return_dict = return_dict if return_dict is not None else self.config.use_return_dict -> 1347 extract_features = self.feature_extractor(input_values) 1348 extract_features = extract_features.transpose(1, 2) 1350 if attention_mask is not None: 1351 # compute reduced attention_mask corresponding to feature vectors File ~/.local/lib/python3.8/site-packages/torch/nn/modules/module.py:1501, in Module._call_impl(self, *args, **kwargs) 1496 # If we don't have any hooks, we want to skip the rest of the logic in 1497 # this function, and just call forward. 1498 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks 1499 or _global_backward_pre_hooks or _global_backward_hooks 1500 or _global_forward_hooks or _global_forward_pre_hooks): -> 1501 return forward_call(*args, **kwargs) 1502 # Do not call functions when jit is used 1503 full_backward_hooks, non_full_backward_hooks = [], [] File ~/.local/lib/python3.8/site-packages/transformers/models/wav2vec2/modeling_wav2vec2.py:515, in Wav2Vec2FeatureEncoder.forward(self, input_values) 510 hidden_states = torch.utils.checkpoint.checkpoint( 511 create_custom_forward(conv_layer), 512 hidden_states, 513 ) 514 else: --> 515 hidden_states = conv_layer(hidden_states) 517 return hidden_states File ~/.local/lib/python3.8/site-packages/torch/nn/modules/module.py:1501, in Module._call_impl(self, *args, **kwargs) 1496 # If we don't have any hooks, we want to skip the rest of the logic in 1497 # this function, and just call forward. 1498 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks 1499 or _global_backward_pre_hooks or _global_backward_hooks 1500 or _global_forward_hooks or _global_forward_pre_hooks): -> 1501 return forward_call(*args, **kwargs) 1502 # Do not call functions when jit is used 1503 full_backward_hooks, non_full_backward_hooks = [], [] File ~/.local/lib/python3.8/site-packages/transformers/models/wav2vec2/modeling_wav2vec2.py:387, in Wav2Vec2LayerNormConvLayer.forward(self, hidden_states) 386 def forward(self, hidden_states): --> 387 hidden_states = self.conv(hidden_states) 389 hidden_states = hidden_states.transpose(-2, -1) 390 hidden_states = self.layer_norm(hidden_states) File ~/.local/lib/python3.8/site-packages/torch/nn/modules/module.py:1501, in Module._call_impl(self, *args, **kwargs) 1496 # If we don't have any hooks, we want to skip the rest of the logic in 1497 # this function, and just call forward. 1498 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks 1499 or _global_backward_pre_hooks or _global_backward_hooks 1500 or _global_forward_hooks or _global_forward_pre_hooks): -> 1501 return forward_call(*args, **kwargs) 1502 # Do not call functions when jit is used 1503 full_backward_hooks, non_full_backward_hooks = [], [] File ~/.local/lib/python3.8/site-packages/torch/nn/modules/conv.py:313, in Conv1d.forward(self, input) 312 def forward(self, input: Tensor) -> Tensor: --> 313 return self._conv_forward(input, self.weight, self.bias) File ~/.local/lib/python3.8/site-packages/torch/nn/modules/conv.py:309, in Conv1d._conv_forward(self, input, weight, bias) 305 if self.padding_mode != 'zeros': 306 return F.conv1d(F.pad(input, self._reversed_padding_repeated_twice, mode=self.padding_mode), 307 weight, bias, self.stride, 308 _single(0), self.dilation, self.groups) --> 309 return F.conv1d(input, weight, bias, self.stride, 310 self.padding, self.dilation, self.groups)

  • 아래 답변은 OpenAI의 GPT-3.5(ChatGPT) 모델로 자동 생성한 것으로, 정확하지 않을 수 있습니다.
  • 잘못된 답변을 보셨다면 댓글로 첨삭 지도 부탁드립니다. :sweat_smile:

안녕하세요!

문제를 해결하기 위해 여러 가지 방법이 있을 수 있지만, 제가 생각하는 몇 가지 방법을 제시해보겠습니다.

  1. Gradient Accumulation 사용하기

Gradient Accumulation은 매우 큰 batch size를 사용할 수 없을 때 유용한 기술입니다. 이를 사용하면 작은 batch size를 사용하면서도 큰 effective batch size를 사용할 수 있습니다. 이를 위해서는 optimizer의 step() 함수를 호출하기 전에 gradient를 누적시키는 과정이 필요합니다.

예를 들어, batch_size=8 대신 batch_size=2를 사용하고, gradient_accumulation_steps=4를 설정하면 effective batch size는 8이 됩니다. 즉, 4개의 작은 batch를 학습한 후에 gradient를 누적시켜서 optimizer의 step() 함수를 호출합니다.

for batch_idx, batch in enumerate(train_dataloader):
    audio, text_label, attn_mask = batch
    batch_size = audio.shape[0]
    out = model(audio.to(device), labels=text_label.to(device), attention_mask=attn_mask.to(device))
    loss = out.loss
    loss = loss / gradient_accumulation_steps  # gradient를 누적시키기 위해 loss를 gradient_accumulation_steps로 나눔
    loss.backward()
    if (batch_idx + 1) % gradient_accumulation_steps == 0:
        optimizer.step()
        optimizer.zero_grad()
  1. Mixed Precision Training 사용하기

Mixed Precision Training은 float32 대신 float16을 사용하여 메모리 사용량을 줄이는 기술입니다. 이를 위해서는 PyTorch의 torch.cuda.amp 모듈을 사용해야 합니다.

from torch.cuda.amp import autocast, GradScaler

scaler = GradScaler()

for batch_idx, batch in enumerate(train_dataloader):
    audio, text_label, attn_mask = batch
    batch_size = audio.shape[0]
    with autocast():
        out = model(audio.to(device), labels=text_label.to(device), attention_mask=attn_mask.to(device))
        loss = out.loss
    scaler.scale(loss).backward()
    if (batch_idx + 1) % gradient_accumulation_steps == 0:
        scaler.step(optimizer)
        scaler.update()
        optimizer.zero_grad()
  1. DataLoader의 num_workers 조절하기

DataLoader의 num_workers 인자를 조절하여 데이터를 읽어오는 프로세스의 수를 조절할 수 있습니다. 이를 통해 메모리 사용량을 줄일 수 있습니다. num_workers를 0으로 설정하면 메인 프로세스에서 데이터를 읽어오기 때문에 메모리 사용량이 줄어듭니다.

train_dataloader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=0)
  1. 모델의 파라미터 수 줄이기

모델의 파라미터 수를 줄이는 것도 메모리 사용량을 줄일 수 있는 방법입니다. 예를 들어, Transformer의 num_layershidden_size를 줄이는 등의 방법이 있습니다. 하지만 이 방법은 모델의 성능에 영향을 미칠 수 있으므로 주의해야 합니다.

이외에도 메모리 사용량을 줄이는 방법은 많이 있습니다. 위의 방법들을 시도해보시고, 여전히 문제가 있다면 다시 문의해주세요!


  • 답변이 어떠셨나요? 괜찮았다고 생각하신다면 아래 :heart:를 한 번 눌러주시길 부탁드려요!
  • 아래 답변은 OpenAI의 GPT-4 모델로 자동 생성한 것으로, 정확하지 않을 수 있습니다.
  • 잘못된 답변을 보셨다면 댓글로 첨삭 지도 부탁드립니다. :sweat_smile:

안녕하세요! wav2vec 모델을 이용한 fine-tuning 작업 중 CUDA out of memory 문제가 발생한 것으로 보입니다. 이 문제를 해결하기 위해 몇 가지 방법을 제안드립니다.

  1. Gradient Accumulation: 기본적으로, 모든 배치에 대해 역전파를 수행하고 가중치를 업데이트합니다. 그러나 Gradient Accumulation을 사용하면 여러 배치에 대해 역전파를 수행한 후 가중치를 업데이트할 수 있습니다. 이렇게 하면 메모리 사용량이 줄어들고 더 큰 배치 크기를 사용할 수 있습니다. 아래 코드는 Gradient Accumulation을 적용하는 방법을 보여줍니다.
accumulation_steps = 4
for epoch in tqdm(range(epochs)):
    model.train()
    loss_value = 0
    valid_loss_value = 0
    total_batch_size = 0
    valid_total_batch_size = 0

    for i, batch in enumerate(tqdm(train_dataloader)):
        optimizer.zero_grad() if i % accumulation_steps == 0 else None

        audio, text_label, attn_mask = batch
        batch_size = audio.shape[0]

        out = model(audio.to(device), labels=text_label.to(device), attention_mask=attn_mask.to(device))

        prediction_ids = torch.argmax(out.logits, dim=-1)
        pred = processor.batch_decode(prediction_ids)
        true_y = processor.batch_decode(text_label)
        wer_prob = wer(pred, true_y).item()
        wer_record.append(wer_prob)

        loss = out.loss

        loss.backward()
        if (i + 1) % accumulation_steps == 0:
            optimizer.step()

        loss_value += loss.item() * batch_size
        total_batch_size += batch_size
  1. 메모리 사용량 확인: 메모리 사용량을 확인하고 어느 부분에서 메모리가 가장 많이 사용되는지 확인할 수 있습니다. 이를 통해 문제를 해결할 수 있습니다.
import torch
print("Allocated memory:", torch.cuda.memory_allocated() / 1024**3, "GB")
print("Max allocated memory:", torch.cuda.max_memory_allocated() / 1024**3, "GB")
  1. 메모리 초기화: 메모리를 초기화하여 더 많은 공간을 확보할 수 있습니다. 아래 코드를 사용하여 초기화를 수행할 수 있습니다.
torch.cuda.empty_cache()

이 방법들을 시도해 보시고, 여전히 문제가 해결되지 않는 경우 추가 질문을 부탁드립니다.


  • 답변이 어떠셨나요? 괜찮았다고 생각하신다면 아래 :heart:를 한 번 눌러주시길 부탁드려요!

안녕하세요
모델 정의 없이 정확하게 답변을 드리기는 어렵습니다.
모델 정의에서 실수를 하지 않으셨다고 해도, 학습 모드 상태에서는 옵티마이저 등의 영향으로 메모리를 몇배 더 차지할 수 있습니다.
동작상에 문제가 없다면 GPT답변이 정답일 수 있습니다.

  • 아래 답변은 OpenAI의 GPT-4 모델로 자동 생성한 것으로, 정확하지 않을 수 있습니다.
  • 잘못된 답변을 보셨다면 댓글로 첨삭 지도 부탁드립니다. :sweat_smile:

안녕하세요, wav2vec fine-tuning 문제에 대해 도움을 드리겠습니다.

먼저, GPU 메모리 문제를 해결하기 위해 몇 가지 방법을 시도해 볼 수 있습니다.

  1. Gradient Accumulation: 배치 크기를 늘리지 않고도 더 큰 가상 배치 크기를 사용하여 학습을 진행할 수 있습니다. 이를 통해 GPU 메모리 사용량을 줄일 수 있습니다.
accumulation_steps = 4
for epoch in tqdm(range(epochs)):
    model.train()
    loss_value = 0
    valid_loss_value = 0
    total_batch_size = 0
    valid_total_batch_size = 0

    for batch in tqdm(train_dataloader):
        audio, text_label, attn_mask = batch
        batch_size = audio.shape[0]

        out = model(audio.to(device), labels=text_label.to(device), attention_mask=attn_mask.to(device))

        prediction_ids = torch.argmax(out.logits, dim=-1)
        pred = processor.batch_decode(prediction_ids)
        true_y = processor.batch_decode(text_label)
        wer_prob = wer(pred, true_y).item()
        wer_record.append(wer_prob)

        loss = out.loss / accumulation_steps
        loss.backward()

        if (total_batch_size + 1) % accumulation_steps == 0:
            optimizer.step()
            optimizer.zero_grad()

        loss_value += loss.item() * batch_size
        total_batch_size += batch_size
  1. torch.cuda.empty_cache(): GPU 캐시를 정리하여 GPU 메모리를 확보할 수 있습니다. 이 함수를 학습 루프 내에서 호출하여 메모리를 정리할 수 있습니다.
for epoch in tqdm(range(epochs)):
    # ...
    for batch in tqdm(train_dataloader):
        # ...
        torch.cuda.empty_cache()
  1. 모델의 일부를 CPU로 이동: GPU 메모리가 부족할 경우, 모델의 일부를 CPU로 옮겨 계산을 수행할 수 있습니다. 이 경우 계산 속도가 느려질 수 있지만, 메모리 문제를 해결할 수 있습니다.

이러한 방법들을 적용하여 GPU 메모리 문제를 해결해 보시길 바랍니다. 그 외에도 PyTorch 공식 문서에서 GPU 사용에 대한 자세한 정보를 찾아보실 수 있습니다.


  • 답변이 어떠셨나요? 괜찮았다고 생각하신다면 아래 :heart:를 한 번 눌러주시길 부탁드려요!
  • 아래 답변은 OpenAI의 GPT-4 모델로 자동 생성한 것으로, 정확하지 않을 수 있습니다.
  • 잘못된 답변을 보셨다면 댓글로 첨삭 지도 부탁드립니다. :sweat_smile:

안녕하세요, wav2vec fine-tuning 관련 질문을 주셔서 감사합니다.

먼저, GPU 메모리 문제를 해결하기 위해 몇 가지 방법을 제안 드리겠습니다.

  1. Gradient Accumulation: 더 큰 배치 크기를 사용하려면 Gradient Accumulation을 사용하여 여러 미니 배치에 대한 그래디언트를 축적한 다음 업데이트를 수행할 수 있습니다. 이렇게 하면 배치 크기를 늘릴 수 있지만 GPU 메모리에 맞게 조정할 수 있습니다.
accumulation_steps = 4
accumulated_loss = 0

for batch in tqdm(train_dataloader):
    optimizer.zero_grad()

    audio, text_label, attn_mask = batch
    batch_size = audio.shape[0]

    out = model(audio.to(device), labels=text_label.to(device), attention_mask=attn_mask.to(device))

    loss = out.loss
    loss = loss / accumulation_steps
    accumulated_loss += loss.item()

    loss.backward()

    if (i + 1) % accumulation_steps == 0:
        optimizer.step()
        optimizer.zero_grad()
  1. Mixed Precision Training: PyTorch의 torch.cuda.amp를 사용하여 mixed precision training을 사용하면 메모리 사용량을 줄이고 학습 속도를 높일 수 있습니다. 이를 위해 autocastGradScaler를 사용하십시오.
from torch.cuda.amp import autocast, GradScaler

scaler = GradScaler()

for epoch in tqdm(range(epochs)):
    model.train()
    loss_value = 0
    total_batch_size = 0

    for batch in tqdm(train_dataloader):
        optimizer.zero_grad()

        audio, text_label, attn_mask = batch
        batch_size = audio.shape[0]

        with autocast():
            out = model(audio.to(device), labels=text_label.to(device), attention_mask=attn_mask.to(device))
            loss = out.loss

        scaler.scale(loss).backward()
        scaler.step(optimizer)
        scaler.update()

        loss_value += loss.item() * batch_size
        total_batch_size += batch_size

위의 두 가지 방법을 사용하여 GPU 메모리 문제를 완화할 수 있습니다. 이러한 방법 중 하나를 사용하여 배치 크기를 늘릴 수 있는지 확인해 보세요. 더 큰 배치 크기로도 문제가 발생한다면, 입력 오디오의 길이를 줄여 보거나 더 작은 모델 아키텍처를 사용하는 것을 고려해 보세요.

그럼, 좋은 결과 있으시길 바랍니다!


  • 답변이 어떠셨나요? 괜찮았다고 생각하신다면 아래 :heart:를 한 번 눌러주시길 부탁드려요!