mirror of
https://github.com/modelscope/FunASR.git
synced 2026-10-12 02:19:18 +00:00
feat: implement Whisper forward() for fine-tuning support (fixes #2678)
- Implemented cross-entropy loss computation in WhisperWarp.forward() - Added fine-tuning example script and documentation - Users can now fine-tune any Whisper model on custom data via FunASR training framework
This commit is contained in:
@@ -0,0 +1,65 @@
|
||||
# Whisper Fine-tuning with FunASR
|
||||
|
||||
Fine-tune OpenAI Whisper models on your own data using FunASR's training framework.
|
||||
|
||||
## Supported Models
|
||||
|
||||
- whisper-tiny / whisper-tiny.en
|
||||
- whisper-base / whisper-base.en
|
||||
- whisper-small / whisper-small.en
|
||||
- whisper-medium / whisper-medium.en
|
||||
- whisper-large-v1 / whisper-large-v2 / whisper-large-v3 / whisper-large-v3-turbo
|
||||
|
||||
## Data Preparation
|
||||
|
||||
Prepare data in JSONL format:
|
||||
|
||||
```json
|
||||
{"key": "utt001", "source": "/path/to/audio1.wav", "target": "the transcription text"}
|
||||
{"key": "utt002", "source": "/path/to/audio2.wav", "target": "another transcription"}
|
||||
```
|
||||
|
||||
## Fine-tuning
|
||||
|
||||
```bash
|
||||
bash finetune.sh
|
||||
```
|
||||
|
||||
Or customize directly:
|
||||
|
||||
```python
|
||||
from funasr import AutoModel
|
||||
|
||||
model = AutoModel(model="Whisper-large-v3", model_conf={"hub": "openai"})
|
||||
|
||||
# Training uses the forward() method which computes cross-entropy loss
|
||||
# on (mel-spectrogram, token_ids) pairs
|
||||
```
|
||||
|
||||
## Key Parameters
|
||||
|
||||
| Parameter | Default | Description |
|
||||
|-----------|---------|-------------|
|
||||
| model | Whisper-large-v3 | Model size |
|
||||
| lr | 1e-5 | Learning rate (lower for larger models) |
|
||||
| max_epoch | 10 | Training epochs |
|
||||
| batch_size | 4 | Per-GPU batch size |
|
||||
| warmup_steps | 500 | LR warmup |
|
||||
|
||||
## Tips
|
||||
|
||||
- For Chinese fine-tuning, use `whisper-large-v3` (best multilingual base)
|
||||
- Freeze encoder for faster training: add `++train_conf.freeze_param="model.encoder"`
|
||||
- Use smaller learning rates (1e-5 ~ 5e-6) to avoid catastrophic forgetting
|
||||
- Recommended: 100+ hours of target-domain audio for meaningful improvement
|
||||
|
||||
## After Fine-tuning
|
||||
|
||||
```python
|
||||
from funasr import AutoModel
|
||||
|
||||
# Load fine-tuned model
|
||||
model = AutoModel(model="/path/to/exp/whisper_finetune")
|
||||
result = model.generate(input="test.wav")
|
||||
print(result[0]["text"])
|
||||
```
|
||||
@@ -0,0 +1,31 @@
|
||||
#!/bin/bash
|
||||
# Whisper Fine-tuning with FunASR
|
||||
#
|
||||
# This script fine-tunes OpenAI Whisper models on custom data using FunASR's training framework.
|
||||
# Supports: whisper-tiny, whisper-base, whisper-small, whisper-medium, whisper-large-v3
|
||||
#
|
||||
# Data format: JSONL with "audio" and "text" fields
|
||||
# {"key": "utt1", "source": "/path/to/audio.wav", "target": "transcription text"}
|
||||
|
||||
export CUDA_VISIBLE_DEVICES=0,1
|
||||
|
||||
model_name="Whisper-large-v3"
|
||||
train_data="data/train.jsonl"
|
||||
val_data="data/val.jsonl"
|
||||
output_dir="exp/whisper_finetune"
|
||||
|
||||
python -m funasr.bin.train \
|
||||
++model="${model_name}" \
|
||||
++model_conf.hub="openai" \
|
||||
++train_data_set_list="${train_data}" \
|
||||
++valid_data_set_list="${val_data}" \
|
||||
++dataset_conf.batch_size=4 \
|
||||
++dataset_conf.num_workers=4 \
|
||||
++train_conf.output_dir="${output_dir}" \
|
||||
++train_conf.max_epoch=10 \
|
||||
++train_conf.lr=1e-5 \
|
||||
++train_conf.warmup_steps=500 \
|
||||
++optim="adam" \
|
||||
++optim_conf.lr=1e-5 \
|
||||
++scheduler="warmuplr" \
|
||||
++scheduler_conf.warmup_steps=500
|
||||
@@ -65,9 +65,48 @@ class WhisperWarp(nn.Module):
|
||||
|
||||
def forward(
|
||||
self,
|
||||
speech: torch.Tensor = None,
|
||||
speech_lengths: torch.Tensor = None,
|
||||
text: torch.Tensor = None,
|
||||
text_lengths: torch.Tensor = None,
|
||||
**kwargs,
|
||||
):
|
||||
"""Forward pass for training."""
|
||||
pass
|
||||
"""Forward pass for training. Computes cross-entropy loss.
|
||||
|
||||
Args:
|
||||
speech: (B, T, D) mel-spectrogram features
|
||||
speech_lengths: (B,) lengths of each audio
|
||||
text: (B, U) token IDs (with SOT/EOT tokens)
|
||||
text_lengths: (B,) lengths of each text sequence
|
||||
Returns:
|
||||
dict with "loss" and optionally "stats"
|
||||
"""
|
||||
if speech is None or text is None:
|
||||
raise ValueError("forward() requires speech and text for training")
|
||||
|
||||
# Encoder
|
||||
audio_features = self.model.encoder(speech)
|
||||
|
||||
# Decoder: shift text right for teacher forcing
|
||||
# text format: [SOT, lang, task, ..., tokens, EOT]
|
||||
decoder_input = text[:, :-1]
|
||||
decoder_target = text[:, 1:]
|
||||
|
||||
# Decoder forward
|
||||
logits = self.model.decoder(decoder_input, audio_features)
|
||||
|
||||
# Cross-entropy loss (ignore padding, token_id = -1 or pad)
|
||||
loss = F.cross_entropy(
|
||||
logits.reshape(-1, logits.size(-1)),
|
||||
decoder_target.reshape(-1),
|
||||
ignore_index=-100,
|
||||
)
|
||||
|
||||
stats = {
|
||||
"loss": loss.detach().item(),
|
||||
"batch_size": speech.size(0),
|
||||
}
|
||||
return {"loss": loss, "stats": stats}
|
||||
|
||||
def inference(
|
||||
self,
|
||||
|
||||
Reference in New Issue
Block a user