torch转onnx torch.jit.script模式 crf解码问题

    torch模型pth转onnx格式时，默认为trace方式，它不支持动态for循环。

    而在crf解码的过程中，需要传入seq_length，然后循环解码，这里就需要用到script模式。

    注意：

     1、转onnx不支持reversed，需要进行改写

     2、seq_length要在torch.jit.script里转为int类型

     3、最后返回的best_tags是list[tensor]，要转换为tensor类型，因为onnx只支持tensor的输入和输出。

     4、下面这个例子，最后得到的best_tags还要reversed一下。

@torch.jit.script
def viterbi_decode(start_transitions, transitions, end_transitions, emissions, seq_length):
    seq_length = int(seq_length)
    score = start_transitions + emissions[0]
    history = []

    for i in range(1, seq_length):
        broadcast_score = score.unsqueeze(1)
        broadcast_emission = emissions[i].unsqueeze(0)
        next_score = broadcast_score + transitions + broadcast_emission
        next_score, indices = next_score.max(dim=0)
        score = next_score
        history.append(indices)

    score += end_transitions
    _, best_last_tag = score.max(dim=0)
    best_tags = [best_last_tag]
    for i in range(len(history) - 1, -1, -1):
        hist = history[i]
        best_last_tag = hist[best_tags[-1]]
        best_tags.append(best_last_tag)

    best_tags = torch.stack(best_tags, dim=0)
    return best_tags, score