solve zipformer streaming gpu inference

[whaozl]

No description provided.

[yaozengwei]

The script jit_trace_export.py exports model with torch.jit.trace. Why replace it with torch.jit.script?

[yfyeung]

Please explain the reason for this modification.

[whaozl]

@yaozengwei because it has a error RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu! when using the torch.jit.trace on encode model.

see k2-fsa/sherpa#346

so, only need encode to torch.jit.script.
the decode and joiner keep the same as torch.jit.trace.

OK

[yaozengwei]

Could you try to convert the model to cuda device instead of cpu when doing the jit.trace exporting (See

icefall/egs/librispeech/ASR/pruned_transducer_stateless7_streaming/jit_trace_export.py

Line 290 in d74822d

model.to("cpu")

)?
We also need to create the inputs on cuda device in this case. (See

icefall/egs/librispeech/ASR/pruned_transducer_stateless7_streaming/jit_trace_export.py

Line 130 in d74822d

x = torch.zeros(1, T, 80, dtype=torch.float32)

)
I wonder if we need to export the model on cuda device when we want to run the model on cuda device.

See https://pytorch.org/docs/stable/jit.html#frequently-asked-questions

[whaozl]

@yaozengwei

there are two scenes:
1、convert the model to cuda device instead of cpu when doing the jit.trace exporting
model.to("cuda:0")
it has a error RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!
or
torch.jit._trace.TracingCheckError: Tracing failed sanity checks!
ERROR: Graphs differed across invocations!
because https://github.com/k2-fsa/icefall/blob/master/egs/librispeech/ASR/pruned_transducer_stateless7_streaming/zipformer.py#L2280-L2281
,but it modify:

rows = torch.arange(start=time1 - 1, end=-1, step=-1)

rows = torch.arange(start=time1 - 1, end=-1, step=-1).cuda()

2、when use cpu to export. the sherpa online[https://github.com/k2-fsa/sherpa/blob/master/sherpa/cpp_api/bin/online-recognizer.cc] use gpu inference, it had a error
Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!

so, the best method is to modify the encode export style. using torch.jit.script.

[yaozengwei]

Could you successfully export the model if you do the change rows = torch.arange(start=time1 - 1, end=-1, step=-1).cuda()?

The reason why we export with jit.trace instead of jit.script is some inference frameworks need that.

[whaozl]

when use rows = torch.arange(start=time1 - 1, end=-1, step=-1).cuda(), it failed.

so I try to use torch.jit.script for the encode model. then, use sherpa online. it can run successfullly when use_gpu.

[yaozengwei]

when use rows = torch.arange(start=time1 - 1, end=-1, step=-1).cuda(), it failed.

so I try to use torch.jit.script for the encode model. then, use sherpa online. it can run successfullly when use_gpu.

Ok. The exported encoder that you are running on cuda device is jit.script version.