Computing Shar features in parallel (CPU and/or GPU)

[popcornell]

Hi all,

I am looking into a way to parallelize feature extraction in the sharded format.
In particular referring to

lhotse/examples/04-lhotse-shar.ipynb

Line 125 in b869488

"from lhotse.shar import ArrayTarWriter, AudioTarWriter, JsonlShardWriter\n",

Here copy-pasted for convenience.

# ArrayTarWriter detects the %06d formatting pattern and will auto-increment the shard number.
with ArrayTarWriter(
    f"{data_dir}/fbank.%06d.tar", shard_size=1000, compression="lilcom"
) as writer:
    for cut in cuts_audio:

        # `feats` is a numpy array with log Mel filter bank features.
        feats = cut.compute_features(fbank)

        # `cut` now contains a field `cut.fbank` with metadata manifest for the features,
        # and a method `cut.load_fbank()` that loads the features (respects pad/truncation).
        cut = cut.attach_tensor(
            "fbank", feats, frame_shift=fbank.frame_shift, temporal_dim=0
        )

        # We store the features under key `cut.id`, because during loading we'll check that the IDs match
        # to avoid data errors. We also store the feature manifest to have some information about this data.
        writer.write(cut.id, feats, cut.fbank)

feature_shards = writer.output_paths

I tried to do the following to parallelize over CPU:

class FeatExtractor(torch.utils.data.Dataset):
    def __init__(self, writers , data_dir, cuts, shard_size=1000):
        self.data_dir = data_dir
        self.cuts = cuts
        self.shard_size = shard_size

        self.writers = writers

        self.fbank = Fbank()
    def __len__(self):
        return len(self.cuts)


    def __getitem__(self, item):
        cut = self.cuts[item]
        feats = cut.compute_features(self.fbank)

        # `cut` now contains a field `cut.fbank` with metadata manifest for the features,
        # and a method `cut.load_fbank()` that loads the features (respects pad/truncation).
        cut = cut.attach_tensor(
            "wavlm", feats, frame_shift=fbank.frame_shift, temporal_dim=0
        )
        c_writer = self.writers[int(item / self.shard_size)]
        # We store the features under key `cut.id`, because during loading we'll check that the IDs match
        # to avoid data errors. We also store the feature manifest to have some information about this data.
        c_writer.write(cut.id, feats, cut.wavlm)
    

data_dir = "./feats_mp"
os.makedirs(data_dir, exist_ok=True)
fbank = Fbank()
n_writers = int(np.ceil(len(cuts_train) / shard_size))
writers = [ArrayTarWriter("{}/fbank.{:06d}.tar".format(data_dir, x),
                    compression="lilcom", shard_size=1) for x in range(n_writers)]
dataset = FeatExtractor(writers, data_dir, cuts_train, shard_size)

dloader = torch.utils.data.DataLoader(dataset, batch_size=1, shuffle=False, num_workers=1, collate_fn=lambda x: x)


for i_batch, elem in enumerate(tqdm.tqdm(dloader)):
            pass

[x.close for x in writers]

feature_shards = [writer.output_paths[0] for writer in writers]
cuts_audio_feat_2 = CutSet.from_shar(
    fields={
        "cuts": shards["cuts"],
        "recording": shards["recording"],
        "wavlm": feature_shards,
    }
)

print(cuts_audio_feat_2[0].to_dict())
print(cuts_audio_feat_2[0].load_wavlm())

Basically I have a writer for each shard and I dispatch different workers (using pytorch as a easy workaround) which will access the correct writer by self.writers[int(item / self.shard_size)].
The problem is that this code only works when I have only one worker:num_workers=1 in the dataloader.
Otherwise it fails with RuntimeError: Uneven number of files in the tarfile (expected to iterate pairs of binary data + JSON metadata.

Here is the full script (basic example on mini-librispeech): pre_compute_sharded.zip

[pzelasko]

There’s a CLI utility for built in feature extraction for shar that you can use or adapt:

lhotse/lhotse/bin/modes/shar.py

Line 132 in b869488

def compute_features(

[popcornell]

Thanks ! I am going to hack this then !
The idea is to use multi-GPU feature extraction for WavLM features.
I think I can close this discussion.

[popcornell]

BTW, I am trying to add WavLM as a feature extractor.
Everything good, sharded feature computation works except that I cannot pass the validation of manifests when extracting the features due to the fact that the model is inconsistent with lhotse.utils.compute_num_frames:

https://github.com/lhotse-speech/lhotse/blob/b86948886237f6acf707d4dc8af81279e0064318/lhotse/utils.py#L390C5-L390C23

WavLM has a output resolution of 0.02 seconds but the length of the features is not consistent with that formula.
What it is the best approach to bypass/overcome this if we want to add pre-trained models as features ?
One solution would be that this function becomes actually a method of FeatureExtractor and can be overridden.

[pzelasko]

Im afraid the behavior of compute num frames is too „hardcoded” in Lhotse to change or customize it now. I can think of the following solutions:

• (best) try porting the logic of snip_edges=False of this function to pad the input inside the feature extractor

  lhotse/lhotse/features/kaldi/layers.py

  Line 743 in b869488

  def _get_strided_batch(

• do nothing and truncate the wavlm output to satisfy compute num frames (but padding would be tricky if it’s too short)
• use custom fields and TemporalArray to store the features instead, I think they won’t „complain” about this mismatch (and then collate using collate_custom_field in the dataset class during training)

[popcornell]

do nothing and truncate the wavlm output to satisfy compute num frames (but padding would be tricky if it’s too short)

It is shorter sometimes that is the problem !
I will go with the snip_edges solution I don't think it will be a problem

[popcornell]

What about features that are not audio-based e.g. BERT output from text ? (lots of multi-modal text and audio stuff these days so they are used quite a lot)
For those I think the only option is custom field.

[pzelasko]

Yes just use custom fields. Either a list of int in the manifest, or save a numpy int array to an HDF file (random reads) or a tar file (sequential reads).