-
Notifications
You must be signed in to change notification settings - Fork 42
Torchscriptable Conformer + high-level "simple" object for decoding, alignments, posteriors, plotting them, etc. #206
Torchscriptable Conformer + high-level "simple" object for decoding, alignments, posteriors, plotting them, etc. #206
Conversation
I don't think I want to go in this direction of having things all bound together in classes -- at least for the time being. |
That makes sense to me TBH because I am not even sure how to make a class like this generic enough to handle the different types of models, topologies, techniques, etc. that we're going to introduce. Anyway, I find it helpful to work with the "current best" model in other projects. I'll test it a bit more to make sure everything works as intended and I'll merge then. |
Thanks!
…On Thu, Jun 3, 2021 at 9:22 PM Piotr Żelasko ***@***.***> wrote:
That makes sense to me TBH because I am not even sure how to make a class
like this generic enough to handle the different types of models,
topologies, techniques, etc. that we're going to introduce. Anyway, I find
it helpful to work with the "current best" model in other projects. I'll
test it a bit more to make sure everything works as intended and I'll merge
then.
—
You are receiving this because you commented.
Reply to this email directly, view it on GitHub
<#206 (comment)>, or
unsubscribe
<https://github.com/notifications/unsubscribe-auth/AAZFLO24OIVGOZEAL4N4ICLTQ56YHANCNFSM457S46ZQ>
.
|
I would propose to let the code support any PyTorch model, at least for those supported by Torch Script. The user only needs to provide a I just write a small demo (see below) to show that the idea is feasible. You can see that with I believe the following small utilities would be helpful: (1)
(2)
(3)
#!/usr/bin/env python3
import torch
import torch.nn as nn
class Model(nn.Module):
def __init__(self, in_dim: int, out_dim: int):
super().__init__()
self.linear = nn.Linear(in_dim, out_dim)
def forward(self, x):
return self.linear(x)
@torch.no_grad()
def main():
m = Model(in_dim=2, out_dim=3)
x = torch.tensor([1, 2.0])
y = m(x)
script_module = torch.jit.script(m)
script_module.save('demo.pt')
new_m = torch.jit.load('demo.pt')
new_y = new_m(x)
print(y)
print(new_y)
if __name__ == '__main__':
main() |
Perhaps someone can test whether our current models are supported by TorchScript, or at least whether it would be possible to make them supported? |
I was just able to convert the Conformer to torchscript with some changes, I'll make it a part of this PR. |
Great!
…On Tue, Jun 8, 2021 at 12:13 AM Piotr Żelasko ***@***.***> wrote:
I was just able to convert the Conformer to torchscript with some changes,
I'll make it a part of this PR.
—
You are receiving this because you commented.
Reply to this email directly, view it on GitHub
<#206 (comment)>, or
unsubscribe
<https://github.com/notifications/unsubscribe-auth/AAZFLO3GE2QSYSSQQYNEP5TTRTVZLANCNFSM457S46ZQ>
.
|
…n now be used even in training but the CUDA usage is extremely low (seems like a bug but I'm not sure where?)
OK, a summary of this PR:
Things that still don't work:
I'm not sure if I have the capacity to work on this further for now -- in any case, Conformer is now scriptable which should open the way for others. |
Saving models to TorchScript works during training with It is OK to merge from my side -- please review and merge if it's adequate. |
One last remark which I forgot to mention -- I did only very naive benchmarking by running in Jupyter the following snippet: %%timeit
with torch.no_grad():
model(features, supervisions) The improvement from normal to TorchScripted model was small -- 140ms vs 130ms. It used a V100 GPU with ~30 cuts in the batch. So it's only the training time when I noticed the slowdown. |
FYI: Though originally TorchScript was advertised for performance, it is now mainly solving the problem of deployment. The resulting object is deployable to C++/iOS/Android. The performance improvement efforts were moved to AI compiler, so in general we can't expect performance improvement. I say it's lucky if you get performance improvement if you get any. This might not be relevant to your application, but when scripting a model, it is possible to perform irreversible operation. For example, recently I added TorchScript-able wav2vec2 to torchaudio. For the sake of supporting quantization, I added a hook for scripting that removes weight normalization forward hook. My rational was that the model was mainly intended for inference so removing a hook is fine. However if a model is scripted during the training, then the wav2vec2 model from torch audio is not compatible with snowfall in kind of unexpected way. Since torch script object file is only architecture and pareters, it feels to me that creating a tool that makes scripted model from training checkpoint file is simpler. But however I do not know the design principles of snowfall or the context of this work, if that's desired, I think it's okay. |
If you are looking for a way to speed up training, quantization aware training is one approach. I heard there was a case where it both improved training time AND accuracy at the same time. |
Thanks @mthrok, that makes a lot of sense. I wondered if people use TorchScript to speed up the training but now it's clear it's not the case. BTW could you elaborate on the AI compiler? Is it Glow, or sth else? The |
Yeah I think one of them. (but I do not know much so please take it with a grain of salt 😥) There is also |
There’s also an ONNX exporter here that converts TorchScript modules https://pytorch.org/docs/stable/onnx.html maybe it’s worth looking into.. |
Thanks! |
Note: this is unfinished, I put together multiple code snippets but didn't polish it in any way. It could be a part of Icefall with a bit of work and documentation.
You can use it e.g. in a jupyter notebook to plot posteriors and alignments, note the high level methods (which all accept cuts as inputs):
compute_features
compute_posteriors
decode
align
align_ctm
plot_alignments
plot_posteriors
You can get the plots from #203 with this.