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:
游雁
2026-05-28 03:39:07 +08:00
parent c83e666df3
commit a8d23eecbb
3 changed files with 137 additions and 2 deletions
@@ -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"])
```
+31
View File
@@ -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
+41 -2
View File
@@ -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,