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