Fix encoder causality in TransformerASR (see issue #2182) - #2262
Conversation
|
I will take a look when I have some time. I'll try to use some of my streaming check tools to see if there isn't any further bug in causal mode. |
|
We started writing streaming-related things a while back and ended up stopping in the middle because we did not have the resources ... so more things might be missing. Maybe @s1zhang can help a bit as well if he sees weird things while exploring streaming with SB. |
asumagic
left a comment
There was a problem hiding this comment.
Using the following code, with the streaming test functions taken from the streaming PR #2140:
import torch
from speechbrain.lobes.models.transformer.TransformerASR import TransformerASR
from speechbrain.utils.streaming import infer_dependency_matrix, plot_dependency_matrix
from matplotlib import pyplot as plt
model = TransformerASR(
tgt_vocab=100,
input_size=32,
d_model=64,
nhead=4,
num_encoder_layers=2,
num_decoder_layers=2,
d_ffn=64*4,
causal=True
)
model.eval()
input_shape = (4, 50, 32) # (batch_size, seq_len, input_size)
with torch.no_grad():
dep_matrix = infer_dependency_matrix(model.encode, seq_shape=input_shape)
plot_dependency_matrix(dep_matrix)
plt.show()I can confirm that TransformerASR.encode looks to be truly causal now. This is not the case if I do not use the causal fixes from this PR.
|
For reference, using that repro code with the PR: Without the PR: A red cell means that changing a specific input had an effect on a specific output. So at least with this test code, no output can see beyond its corresponding input as expected. |
|
Thank you @lucadellalib for the fix and @asumagic for the review! |
No description provided.