mirror of
https://github.com/k2-fsa/icefall.git
synced 2025-08-09 01:52:41 +00:00
71 lines
2.2 KiB
Python
Executable File
71 lines
2.2 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
from lstmp import LSTMP
|
|
|
|
|
|
def test():
|
|
input_size = torch.randint(low=10, high=1024, size=(1,)).item()
|
|
hidden_size = torch.randint(low=10, high=1024, size=(1,)).item()
|
|
proj_size = hidden_size - 1
|
|
lstm = nn.LSTM(
|
|
input_size=input_size,
|
|
hidden_size=hidden_size,
|
|
num_layers=1,
|
|
bias=True,
|
|
proj_size=proj_size,
|
|
)
|
|
lstmp = LSTMP(lstm)
|
|
|
|
N = torch.randint(low=1, high=10, size=(1,)).item()
|
|
T = torch.randint(low=1, high=20, size=(1,)).item()
|
|
x = torch.rand(T, N, input_size)
|
|
h0 = torch.rand(1, N, proj_size)
|
|
c0 = torch.rand(1, N, hidden_size)
|
|
|
|
y1, (h1, c1) = lstm(x, (h0, c0))
|
|
y2, (h2, c2) = lstmp(x, (h0, c0))
|
|
|
|
assert torch.allclose(y1, y2, atol=1e-5), (y1 - y2).abs().max()
|
|
assert torch.allclose(h1, h2, atol=1e-5), (h1 - h2).abs().max()
|
|
assert torch.allclose(c1, c2, atol=1e-5), (c1 - c2).abs().max()
|
|
|
|
# lstm_script = torch.jit.script(lstm) # pytorch does not support it
|
|
lstm_script = lstm
|
|
lstmp_script = torch.jit.script(lstmp)
|
|
|
|
y3, (h3, c3) = lstm_script(x, (h0, c0))
|
|
y4, (h4, c4) = lstmp_script(x, (h0, c0))
|
|
|
|
assert torch.allclose(y3, y4, atol=1e-5), (y3 - y4).abs().max()
|
|
assert torch.allclose(h3, h4, atol=1e-5), (h3 - h4).abs().max()
|
|
assert torch.allclose(c3, c4, atol=1e-5), (c3 - c4).abs().max()
|
|
|
|
assert torch.allclose(y3, y1, atol=1e-5), (y3 - y1).abs().max()
|
|
assert torch.allclose(h3, h1, atol=1e-5), (h3 - h1).abs().max()
|
|
assert torch.allclose(c3, c1, atol=1e-5), (c3 - c1).abs().max()
|
|
|
|
lstm_trace = torch.jit.trace(lstm, (x, (h0, c0)))
|
|
lstmp_trace = torch.jit.trace(lstmp, (x, (h0, c0)))
|
|
|
|
y5, (h5, c5) = lstm_trace(x, (h0, c0))
|
|
y6, (h6, c6) = lstmp_trace(x, (h0, c0))
|
|
|
|
assert torch.allclose(y5, y6, atol=1e-5), (y5 - y6).abs().max()
|
|
assert torch.allclose(h5, h6, atol=1e-5), (h5 - h6).abs().max()
|
|
assert torch.allclose(c5, c6, atol=1e-5), (c5 - c6).abs().max()
|
|
|
|
assert torch.allclose(y5, y1, atol=1e-5), (y5 - y1).abs().max()
|
|
assert torch.allclose(h5, h1, atol=1e-5), (h5 - h1).abs().max()
|
|
assert torch.allclose(c5, c1, atol=1e-5), (c5 - c1).abs().max()
|
|
|
|
|
|
@torch.no_grad()
|
|
def main():
|
|
test()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|