Let’s assume that I have a 3D tensor x = torch.rand(4, 17, 128), from which I want to average a selected span of elements (of variable length) for each row:
spans = [[[1, 3], [6,10], [11, 15]],
[[1, 10]],
[[1, 5]],
[[1,7], [8,16]]]
For example, in this case, I want to average elements 1 to 3, 6 to 10, and 11 to 15 from row one, then elements 1 to 10 from row two, etc. Each of the averaged tensors, in this case, will be of shape (128,). Then, I want to insert the span average to another tensor, defined as:
# 4
num_elements = len(spans)
# 3
max_spans = max([len(span) for span in spans])
# [128]
remaing_dims = x.size()[2:]
# [4, 3, 128]
output = torch.zeros(num_elements, max_spans, *remaing_dims)
Currently, in order to do what I want I do the following:
for i, element_spans in enumerate(spans):
for j, span in enumerate(element_spans):
output[i, j] = x[i, span[0]:span[1]].mean(0)
which is a really brute-force kind of approach. Therefore, I am wondering how can I speed-up the whole select-average-insert.
Thanks in advance!