Add mask_tensor util #35
Conversation
Utility to create a boolean tensor mask on the tensors device based on sequence lengths.
JackTemaki
requested review from
albertz,
curufinwe,
christophmluscher,
mmz33,
Marvin84,
Judyxujj,
SimBe195,
vieting,
ZijYang and
Atticus1806
| :param seq_len: [B] | ||
| :return: Mask of [B,T] | ||
| """ | ||
| seq_len = seq_len.to(device=tensor.device) |
There was a problem hiding this comment.
I am not sure if this case should be made within this function.
There was a problem hiding this comment.
The device handling? So completely removing tensor and having just a seq_len to mask function?
There was a problem hiding this comment.
you still need the tensor for the shape right? But I feel like this casting of the seq_len should be done beforehand (and at least in my setups usually already is done with the model call. So this might be redundant imho.
There was a problem hiding this comment.
You could do it without the tensor (using max(seq_len)), I think we can add this as a second function
There was a problem hiding this comment.
If seq_len is already on the device, this should have not more overhead than one integer comparison
There was a problem hiding this comment.
is this max then exportable?
In general if the overhead is low I don't mind and we can leave it like this.
| import torch | ||
|
|
||
|
|
||
| def mask_tensor(tensor: torch.Tensor, seq_len: torch.Tensor) -> torch.Tensor: |
There was a problem hiding this comment.
I'm not sure about the name. I find it a bit too generic. And also ambiguous. (Does it mask a tensor? No, actually, it returns a mask tensor. For what? Based on what?)
In TF, the corresponding function is called tf.sequence_mask. Why not just use the same name here? So all people who have used TF at some point in the past (almost everyone) knows what this is about.
There was a problem hiding this comment.
tf.sequence_mask has a different signature, and we can also add such a function. This is just not what I wanted to have here.
There was a problem hiding this comment.
For the name, yes this can be more verbose if needed.
| import torch | ||
|
|
||
|
|
||
| def mask_tensor(tensor: torch.Tensor, seq_len: torch.Tensor) -> torch.Tensor: |
There was a problem hiding this comment.
You should not use tensor here as an argument, as this is not really required. What you actually want is an argument maxlen, where you can just pass tensor.shape[1] or whatever. That argument could also be optional, and then it would just take the max from seq_len.
There was a problem hiding this comment.
What you actually want
What I want is to have it with a tensor and make sure the mask is on the same device as the tensor. But maybe the majority disagrees, and I keep this for myself.
|
Closing, as this specific signature seems not to be beneficial for people other than me. |
Utility to create a boolean tensor mask on the tensors device based on sequence lengths.