pytorch中不定长序列补齐方法

pytorch中不定长序列补齐方法

第二种方法通常是在load一个batch数据时, 在collate_fn中进行补齐的. collate_fn使用方法

以下给出两种思路:

第一种思路是比较容易想到的, 就是对一个batch的样本进行遍历, 然后使用np.pad对每一个样本进行补齐.

for unit in data:
        mask = np.zeros(max_length)
        s_len = len(unit[0])    # calculate the length of sequence in each unit
        mask[: s_len] = 1
        unit[0] = np.pad(unit[0], (0, max_length - s_len), 'constant', constant_values=(0, 0))
        mask_batch.append(mask)

但是这种方法在batch size很大的情况下会很慢, 因为使用for循环进行了遍历. 我在实际用的时候, 当batch_size=128时, 一个batch的加载时间甚至是一个batch训练时间的几倍!

因此, 我想到如何并行地对序列进行补齐. 第二种方法的思路就是使用torch中自带的pad_sequence来并行补齐.

        batch_sequence = list(map(lambda x: torch.tensor(x[findex]), x_data))
        batch_data[feat] = torch.nn.utils.rnn.pad_sequence(batch_sequence).T

可以看到这里使用pad_sequence一次性对整个batch进行补齐. 下面对这个函数进行详细说明.

pad_sequence详解

from torch.utils.rnn import pad_sequence

a = torch.ones(10)
b = torch.ones(6)
c = torch.ones(20)

abc = pad_sequence([a,b,c])		# shape(20, 3)

注意这个函数接收的是一个元素为tensor的列表, 而不是tensor.

最终, 这个函数会将所有tensor转换为tensor矩阵#shape(max_length, batch_size). 因此, 在使用完后通常还需要转置一下.