I’m trying to use einsum in my implementation of Losses (eg. Dice Loss) for network training. So far I trained on 2D images, and it worked nicely, but now I tried to extend the code to work also on 3D data.
As I want the same code to work on both, 2 and 3D, I went for example from using
torch.einsum(‘bkwh,bkwh->bk’)
to
torch.einsum('bk…,‘bk…->bk’),
basically using ellipsis for the unknown number of dimensions. The desired output is still the same though - elementwise multiplication and then summation over all but first two dimensions.
But so, with this syntax, I want to do the summation over the ellipsis, which according to the documentation should work ( torch.einsum — PyTorch 1.8.1 documentation :
See note: " torch.einsum handles ellipsis (‘…’) differently from NumPy in that it allows dimensions covered by the ellipsis to be summed over, that is, ellipsis are not required to be part of the output. " ).
But for some reason it doesn’t work for me; when I use this particular operation on two tensors, both of size [32, 7, 24, 24, 24] for example, it throws an error: " RuntimeError: shape ‘[32, 7]’ is invalid for input of size 3096576 ". So it would seem like it doesn’t know how to sum over ellipsis after all?
Or am I somehow using it wrong?
