batch inference

[ltm920716]

你好，我修改了推理代码实现了batch推理，但是rnn那个部分的推理时间和batch成线性关系，请问能帮忙解释一下这个现象么，谢谢

lenText = nT
nsteps = nT
output = torch.zeros(nB, lenText).int()
output_probs = torch.zeros(nB, lenText).float()

hidden = torch.zeros(nB, self.nchannel).type_as(C.data)
stats = torch.zeros(nB, self.nchannel).type_as(C.data)
prev_emb = self.char_embeddings.index_select(0, torch.zeros(nB).type_as(C.data).long())

for step in range(nsteps):
hidden, stats = self.rnn(torch.cat((C[step, :, :], prev_emb), dim=1), (hidden, stats))
step_result = self.generator(hidden)
step_result = F.softmax(step_result, dim=-1)
max_prob_index = torch.argmax(step_result, dim=1)
max_prob = step_result.index_select(dim=1, index=max_prob_index)
output[:, step] = max_prob_index
output_probs[:, step] = max_prob
prev_emb = self.char_embeddings.index_select(0, max_prob_index.long())

return output, output_probs