perch2-torch-embed / README.md
postpop's picture
Add embedding-only perch2-torch (torch.export .pt2) and perch2_embed.py helper
fb3c11b verified
|
Raw History Blame Contribute Delete
6.74 kB
---
license: apache-2.0
library_name: pytorch
pipeline_tag: feature-extraction
base_model:
- DAS-DeepAudioSegmenter/perch2-torch
tags:
- audio
- bioacoustics
- embeddings
- torch-export
- das
---
# PERCH 2 embeddings (PyTorch, embedding only)
The PERCH 2 embedding network from [perch2-torch](https://huggingface.co/DAS-DeepAudioSegmenter/perch2-torch) as a
single self-contained `torch.export` file. The 14,795-class label head is left out, so the file holds just the waveform
frontend and EfficientNet-B3 backbone: 10.7 M instead of 101.7 M parameters, and 47.6 MB instead of 408 MB. For class
logits, use [perch2-torch](https://huggingface.co/DAS-DeepAudioSegmenter/perch2-torch).
## Use the model file
`perch_v2_embed.pt2` takes a batch of 5 s mono windows at 32 kHz, `(batch, 160000)` with batch 1-4096, and returns
the unpooled spatial embedding `(batch, 16, 4, 1536)`: 16 time frames and 4 frequency positions per window.
```python
import torch
from huggingface_hub import hf_hub_download
from torch.export.passes import move_to_device_pass
path = hf_hub_download("DAS-DeepAudioSegmenter/perch2-torch-embed", "perch_v2_embed.pt2")
core = move_to_device_pass(torch.export.load(path), "cpu").module() # or "cuda" / "mps"
windows = torch.zeros(2, 160_000) # two 5 s windows of 32 kHz audio
spatial = core(windows) # (2, 16, 4, 1536)
embedding = spatial.mean(dim=(1, 2)) # (2, 1536), one embedding per window
```
Requires PyTorch >= 2.11. `torch.export` files load in the version they were exported with or newer; tested with 2.11
and 2.14.
## Limitations of the model file on its own
- **Fixed 5 s input.** Each window must be exactly 160,000 samples. Shorter audio has to be zero-padded, and longer
audio cut into windows by you.
- **32 kHz assumed.** The model does not resample; whatever you pass is treated as 32 kHz audio.
- **Per-window output only.** You get unpooled frames for each window. There is no song-level average or maximum, and
no continuous time axis across overlapping windows, which needs the overlap frames trimmed and averaged.
- **Memory.** All windows passed in one call (up to 4,096) are processed at once, so long recordings have to be split
into batches.
## Alternative: `perch2_embed.py` for audio of any length
`perch2_embed.py` handles all of the above: windowing, padding, batching, overlap pooling and truncation. It needs only
torch and huggingface_hub.
```bash
pip install torch huggingface_hub soundfile
hf download DAS-DeepAudioSegmenter/perch2-torch-embed perch2_embed.py --local-dir .
```
```python
import soundfile as sf
from perch2_embed import load, embed
core = load("auto") # downloads and caches perch_v2_embed.pt2
wave, sr = sf.read("song.wav", dtype="float32") # the model assumes 32 kHz: resample first for normal speed
out = embed(core, wave if wave.ndim == 1 else wave.mean(axis=1), sample_rate=sr)
out["avg"].shape, out["max"].shape, out["temporal"].shape # (1536,), (1536,), (frames, 1536)
```
### Options
`load(device="cpu", path=None, revision=None)`
- `device`: `"cpu"`, `"cuda"`, `"mps"`, or `"auto"` (CUDA, then MPS, then CPU).
- `path`: a local copy of `perch_v2_embed.pt2`; otherwise it is downloaded from this repo and cached.
- `revision`: a repo revision to pin the downloaded file.
`embed(core, wave, hop_frames=1, max_seconds=None, sample_rate=32000, batch_size=32)`
- `wave`: mono 1-D float array or tensor of any length, fed to the model unchanged (treated as 32 kHz).
- `hop_frames`: window hop in 0.32 s frames, 1-15. `1` is the model's own hop and reproduces the full perch2-torch
model; larger values compute fewer windows (see below).
- `max_seconds`: keep only the first `max_seconds` of audio. `None` embeds everything.
- `sample_rate`: the rate at which `wave` was sampled. Nothing is resampled: it only converts `max_seconds` to samples
and frame positions to seconds in `times`.
- `batch_size`: 5 s windows per forward pass, which sets memory use.
### Returns
A dict with:
- `avg` `(1536,)`: the same as the full model's `embedding` output;
- `max` `(1536,)`: the maximum over the same overlap-trimmed frames (the full model does not output this);
- `temporal` `(frames, 1536)`: stitched frames 0.32 s apart in model time, averaged over the 4 frequency positions;
- `spatial` `(frames, 4, 1536)`: the same frames with the frequency positions kept;
- `times` `(frames,)`: frame centres in seconds;
- `n_windows`: the number of 5 s windows computed.
### Notes
- The model always treats its input as 32,000 samples per second, so the sample rate of `wave` sets the playback
speed the model hears. Audio at 32 kHz is heard at normal speed. Audio at a higher rate fed unchanged is slowed
down and lowered in pitch by `sample_rate / 32000`. For example, a recording resampled to 64 kHz (e.g.
`librosa.load(path, sr=64000)`) and fed unchanged is heard at half speed and half pitch, and a 44.1 kHz file read
as is plays at 0.73x speed. Windows, hops and frames are fixed in samples: one 5 s window covers 160,000 / sample_rate
seconds of your audio and frames are 10,240 / sample_rate seconds apart (0.32 s at 32 kHz, 0.16 s at 64 kHz). Pass the
true `sample_rate` so that `max_seconds` and `times` are in seconds of the original audio.
- `hop_frames > 1` computes fewer windows, but the embeddings change. On 75 bird songs of 8 s or longer (audio loaded at
64 kHz), cosine similarity to the default hop was:
| hop | windows computed | avg: median (worst) | max: median (worst) |
|---|---|---|---|
| 2 (0.64 s) | 52% | 0.993 (0.889) | 0.978 (0.930) |
| 4 (1.28 s) | 28% | 0.989 (0.888) | 0.982 (0.937) |
| 8 (2.56 s) | 15% | 0.977 (0.879) | 0.983 (0.938) |
| 15 (4.8 s) | 10% | 0.955 (0.840) | 0.969 (0.862) |
Use the same hop for everything you compare.
## Verification
Against perch2-torch (revision `d5d3bb6a`):
- the core matches the full model's frontend and backbone: CPU difference <= 4e-6, MPS relative difference 2.4e-6;
- `embed()` at the default hop matches the full model's windowing and pooling for inputs of 1 to 361 windows (avg
within 3e-7, max and temporal within 1e-5), and `avg` equals the full model's `embedding` within 3e-7;
- truncation equals embedding the pre-cut audio;
- PyTorch 2.14 loads the file with timm unavailable and gives identical output.
## Attribution and license
Exported from [perch2-torch](https://huggingface.co/DAS-DeepAudioSegmenter/perch2-torch), which is derived from the
Apache-2.0 [PERCH ONNX model](https://huggingface.co/justinchuby/Perch-onnx) and the Apache-2.0
[wrice EfficientNet-B3 checkpoint](https://huggingface.co/wrice/perch-v2-efficientnet-b3). Apache-2.0.