Shifted indices for LazyTensors

[oliver-batchelor]

Hi!

Lets say I have a Vector of length say N.

I can easily create a LazyTensor which stacks this vector M times to make a matrix sized (M, N)

Is it possible to create a 2D LazyTensor from this vector which is shifted for each position, so I still create a matrix sized (M, N) but each time the vector is shifted over one or more indexes (and padded).

I'm interested in doing this for efficient stereo matching, where instead of creating a large 3D cost volume, it is lazily created and reduced without storing any large 3D volume.

Thanks!
Oliver

[oliver-batchelor]

I see this is subsumed by #66

[jeanfeydy]

Hi @saulzar ,

Thanks for your interesting question! First of all, all my apologies for this late reply: we have been extremely busy over the last two months and are just starting to catch up on late issues.

As far as I can tell, your question is more closely related to e.g. #24 than #66 . If you're still interested in the problem, would you mind telling us a little bit more on your problem? Depending on the dimension of the vectors and padding required, the optimal way of answering your question may vary quite a lot.

Best regards,
Jean

[oliver-batchelor]

Thanks very much for getting back to me! Yes #24 does seem a lot more related! Though it seems to me to make good use of something like ke-ops for this purpose might require being able to re-shape that diagonal - this is why I was looking at #66 as he was also discussing remapping dimensions. Here's some pytorch code of the operation I am using, at the end (two versions, one the dumb version, second the weird one by changing up strides and adding padding to add an extra dimension).

It also might be possible to do this with a batch matrix multiplication, but I had quite a bit of trouble making it work with the slightly odd strided setup.

Hopefully this first version is relatively clear - it takes as input two images of B C H W format (same sizes), and outputs a 3D volume B D H W.

It's iterating from 0..D (max_disparity) and shifting the left image across the image width dimension (3), and computing a correlation along the channel feature dimension (1). The slices are stacked to make a 3D volume B D H W.

This dimension D will be reduced (using a SumSoftMaxWeight in ke-ops terminology) to give back a B C H W image, this can probably be considered as a second step though - but it would be very interesting if it can all be done as one step!

def correlation_volume(left, right, max_disparity):
    width = left.size(3)

    def disparity_slice(i):
        slice_l = F.pad(left[..., i:width], (i, 0, 0, 0))
        slice_r = F.pad(right[..., :width-i], (i, 0, 0, 0))
        return (slice_l * slice_r).sum(dim=1) 

    slices = [disparity_slice(max_disparity - i - 1) for i in range(max_disparity)]
    return torch.stack(slices, dim=1)

def correlation_volume2(left, right, max_disparity):
    b, c, h, w = right.shape
    padded = F.pad(right, (max_disparity - 1, 0, 0, 0))
    
    strides = padded.stride()
    right_vol = padded.as_strided((b, c, h, max_disparity, w), [*strides[:-1], 1, 1], storage_offset=0)
    left_vol = left.unsqueeze(3).expand([b, c, h, max_disparity, w])

    volume = (left_vol * right_vol).sum(1)
    return volume.permute(0, 2, 1, 3)