pytorch 中的批量光束搜索答案

【问题标题】：Batch-wise beam search in pytorchpytorch 中的批量光束搜索
【发布时间】：2021-01-29 01:47:48
【问题描述】：

我正在尝试在文本生成模型中实现波束搜索解码策略。这是我用来解码输出概率的函数。

def beam_search_decoder(data, k):
    sequences = [[list(), 0.0]]
    # walk over each step in sequence
    for row in data:
        all_candidates = list()
        for i in range(len(sequences)):
            seq, score = sequences[i]
            for j in range(len(row)):
                candidate = [seq + [j], score - torch.log(row[j])]
                all_candidates.append(candidate)
        # sort candidates by score
        ordered = sorted(all_candidates, key=lambda tup:tup[1])
        sequences = ordered[:k]
    return sequences

现在你可以看到这个函数是在考虑到 batch_size 1 的情况下实现的。为批量大小添加另一个循环将使算法O(n^4)。和现在一样慢。有什么办法可以提高这个功能的速度。我的模型输出通常大小为(32, 150, 9907)，格式为(batch_size, max_len, vocab_size)

【问题讨论】：

您应该在 Pytorch 中通过快速 google 搜索找到光束搜索实现。请注意，波束解码并非易事，有几个因素相关。因此，建议使用其中一种可用的实现。
束搜索策略在测试期间是有意义的。你不能维护一个batch_size=1 并并行处理测试示例吗？
您还可以查看beam search implementation 和代码in this repo 使用修改后的Transformer 进行图像字幕。该实现使用 PyTorch 的 register_buffer 来缓存前一个时间步的输入，以便在当前时间步中只提供新的输入，并且速度相当快。

标签： python deep-learning nlp pytorch beam-search

【解决方案1】：

下面是我的实现，可能比for循环的实现要快一点。

import torch


def beam_search_decoder(post, k):
    """Beam Search Decoder

    Parameters:

        post(Tensor) – the posterior of network.
        k(int) – beam size of decoder.

    Outputs:

        indices(Tensor) – a beam of index sequence.
        log_prob(Tensor) – a beam of log likelihood of sequence.

    Shape:

        post: (batch_size, seq_length, vocab_size).
        indices: (batch_size, beam_size, seq_length).
        log_prob: (batch_size, beam_size).

    Examples:

        >>> post = torch.softmax(torch.randn([32, 20, 1000]), -1)
        >>> indices, log_prob = beam_search_decoder(post, 3)

    """

    batch_size, seq_length, _ = post.shape
    log_post = post.log()
    log_prob, indices = log_post[:, 0, :].topk(k, sorted=True)
    indices = indices.unsqueeze(-1)
    for i in range(1, seq_length):
        log_prob = log_prob.unsqueeze(-1) + log_post[:, i, :].unsqueeze(1).repeat(1, k, 1)
        log_prob, index = log_prob.view(batch_size, -1).topk(k, sorted=True)
        indices = torch.cat([indices, index.unsqueeze(-1)], dim=-1)
    return indices, log_prob

【讨论】：

【解决方案2】：

你可以使用这个库

https://pypi.org/project/pytorch-beam-search/

它为 PyTorch 序列模型实现了 Beam Search、Greedy Search 和采样。

下面的 sn-p 实现了一个 Transformer seq2seq 模型并使用它来生成预测。

#pip install pytorch-beam-search
from pytorch_beam_search import seq2seq

# Create vocabularies
# Tokenize the way you need
source = [list("abcdefghijkl"), list("mnopqrstwxyz")]
target = [list("ABCDEFGHIJKL"), list("MNOPQRSTWXYZ")]
# An Index object represents a mapping from the vocabulary to
# to integers (indices) to feed into the models
source_index = seq2seq.Index(source)
target_index = seq2seq.Index(target)

# Create tensors
X = source_index.text2tensor(source)
Y = target_index.text2tensor(target)
# X.shape == (n_source_examples, len_source_examples) == (2, 11)
# Y.shape == (n_target_examples, len_target_examples) == (2, 12)

# Create and train the model
model = seq2seq.Transformer(source_index, target_index)    # just a PyTorch model
model.fit(X, Y, epochs = 100)    # basic method included

# Generate new predictions
new_source = [list("new first in"), list("new second in")]
new_target = [list("new first out"), list("new second out")]
X_new = source_index.text2tensor(new_source)
Y_new = target_index.text2tensor(new_target)
loss, error_rate = model.evaluate(X_new, Y_new)    # basic method included
predictions, log_probabilities = seq2seq.beam_search(model, X_new) 
output = [target_index.tensor2text(p) for p in predictions]
output

【讨论】：