26 lines
940 B
Python
26 lines
940 B
Python
|
import numpy as np
|
|||
|
|
import torch
|
|||
|
|
|
|||
|
|
|
|||
|
|
# 对一个batch的数据处理
|
|||
|
|
def collate_fn(batch):
|
|||
|
|
# 找出音频长度最长的
|
|||
|
|
batch = sorted(batch, key=lambda sample: sample[0].shape[0], reverse=True)
|
|||
|
|
max_audio_length = batch[0][0].shape[0]
|
|||
|
|
batch_size = len(batch)
|
|||
|
|
# 以最大的长度创建0张量
|
|||
|
|
inputs = np.zeros((batch_size, max_audio_length), dtype='float32')
|
|||
|
|
input_lens_ratio = []
|
|||
|
|
labels = []
|
|||
|
|
for x in range(batch_size):
|
|||
|
|
sample = batch[x]
|
|||
|
|
tensor = sample[0]
|
|||
|
|
labels.append(sample[1])
|
|||
|
|
seq_length = tensor.shape[0]
|
|||
|
|
# 将数据插入都0张量中，实现了padding
|
|||
|
|
inputs[x, :seq_length] = tensor[:]
|
|||
|
|
input_lens_ratio.append(seq_length/max_audio_length)
|
|||
|
|
input_lens_ratio = np.array(input_lens_ratio, dtype='float32')
|
|||
|
|
labels = np.array(labels, dtype='int64')
|
|||
|
|
return torch.tensor(inputs), torch.tensor(labels), torch.tensor(input_lens_ratio)
|