Merge pull request #98 from dlipatov/fix-chunksize

Fix chunksize in utils.py when batch is less than ncpu value
This commit is contained in:
dscripka 2024-01-18 20:18:59 -05:00 committed by GitHub
commit dc5a234218
No known key found for this signature in database
GPG key ID: B5690EEEBB952194

View file

@ -269,8 +269,9 @@ class AudioFeatures():
result = self._get_melspectrogram(batch)
elif pool:
chunksize = batch.shape[0]//ncpu if batch.shape[0] >= ncpu else 1
result = np.array(pool.map(self._get_melspectrogram,
batch, chunksize=batch.shape[0]//ncpu))
batch, chunksize=chunksize))
melspecs[i:i+batch_size, :, :] = result.squeeze()
@ -330,8 +331,9 @@ class AudioFeatures():
result = self.embedding_model_predict(batch)
elif pool:
chunksize = batch.shape[0]//ncpu if batch.shape[0] >= ncpu else 1
result = np.array(pool.map(self._get_embeddings_from_melspec,
batch, chunksize=batch.shape[0]//ncpu))
batch, chunksize=chunksize))
for j, ndx2 in zip(range(0, result.shape[0], n_frames), ndcs):
embeddings[ndx2, :, :] = result[j:j+n_frames]