Dynamic batch the data until the total frames in batch reach `max_frames_in_batch` Args: data: Iterable[{key, feat, label}] max_frames_in_batch: max_frames in one batch Returns: Iterable[List[{key, feat, label}]]
(data, max_frames_in_batch=12000, mode='train')
| 397 | |
| 398 | |
| 399 | def dynamic_batch(data, max_frames_in_batch=12000, mode='train'): |
| 400 | """ Dynamic batch the data until the total frames in batch |
| 401 | reach `max_frames_in_batch` |
| 402 | |
| 403 | Args: |
| 404 | data: Iterable[{key, feat, label}] |
| 405 | max_frames_in_batch: max_frames in one batch |
| 406 | |
| 407 | Returns: |
| 408 | Iterable[List[{key, feat, label}]] |
| 409 | """ |
| 410 | buf = [] |
| 411 | longest_frames = 0 |
| 412 | for sample in data: |
| 413 | assert 'acoustic_token' in sample |
| 414 | assert isinstance(sample['acoustic_token'], torch.Tensor) |
| 415 | |
| 416 | if 'semantic_token' in sample: |
| 417 | new_sample_frames = sample['semantic_token'][0].shape[0] |
| 418 | else: |
| 419 | new_sample_frames = sample['semantic_token'] |
| 420 | |
| 421 | if "text_token" in sample: |
| 422 | new_sample_frames += len(sample['text_token']) |
| 423 | |
| 424 | longest_frames = max(longest_frames, new_sample_frames) |
| 425 | frames_after_padding = longest_frames * (len(buf) + 1) |
| 426 | |
| 427 | if frames_after_padding > max_frames_in_batch: |
| 428 | if len(buf) > 0: |
| 429 | yield buf |
| 430 | buf = [sample] |
| 431 | longest_frames = new_sample_frames |
| 432 | else: |
| 433 | buf.append(sample) |
| 434 | if len(buf) > 0: |
| 435 | yield buf |
| 436 | |
| 437 | |
| 438 | def batch(data, batch_type='static', batch_size=16, max_frames_in_batch=12000, |