Issue
I am having some problem trying to convert the following function into a function only manipulating tensors.
def valid_sequence_output(sequence_output, valid_mask):
bs, max_len, feat_dim = sequence_output.shape
valid_output = torch.zeros(bs, max_len, feat_dim, dtype=torch.float32)
for i in range(bs):
jj = -1
for j in range(max_len):
if valid_mask[i][j].item() == 1:
jj += 1
valid_output[i][jj] = sequence_output[i][j]
return valid_output
where the input tensors can be created as follow:
size = ((2,5,2))
sequence_output = torch.randint(0, 250, size=size)
valid_mask = torch.randint(0, 2, size=size[:2])
I basically aim at “bubbling up” the non null rows of sequence_output
. As an example having sequence_output
equal to:
tensor([[[ 0, 0],
[ 15, 47],
[124, 230],
[ 0, 0],
[ 65, 31]],
[[ 0, 0],
[ 0, 0],
[ 0, 0],
[139, 228],
[224, 205]]])
I am trying to obtain the following tensor:
tensor([[[ 15, 47],
[124, 230],
[ 65, 31],
[ 0, 0],
[ 0, 0]],
[[139, 228],
[224, 205],
[ 0, 0],
[ 0, 0],
[ 0, 0]]])
If somebody has a suggestion on how to do this, I would really appreciate :D
Solution
I managed to come up with a quite nasty solution (which he's probably suboptimal) by constructing a matrix A
of zeros and ones that will swap rows of X
when performing a matrix multiplication...
def vso(seq_out, valid_mask):
X = torch.where(valid_mask.unsqueeze(-1) == 1, seq_out, torch.zeros_like(seq_out))
bs, max_len, _ = X.shape
tu = torch.unique(torch.nonzero(X)[:, :2], dim=0)
batch_axis = tu[:, 0]
rows_axis = tu[:, 1]
a = torch.arange(bs).repeat(batch_axis.shape).reshape(batch_axis.shape[0], -1).T
T = torch.cumsum(batch_axis == a, dim=1) - 1
cols_axis = T[batch_axis, torch.arange(batch_axis.shape[0])]
A = torch.zeros((bs, max_len, max_len))
A[(batch_axis, cols_axis, rows_axis)] = 1
valid_output = torch.matmul(A, X)
return valid_output
Still looking for better answers though !
Answered By - Jules
0 comments:
Post a Comment
Note: Only a member of this blog may post a comment.