EinopsError: Error while processing rearrange-reduction pattern "(b1 b2) h w c -> (b1 h) (b2 w) c"

Question

I'm learning the basics of einops to incorporate in my code.

process = transforms.Compose([
    transforms.Resize(225),
    transforms.ToTensor()
])

cat = Image.open('cat.jpeg').convert('RGB')

cat = process(cat)

rearrange(cat, '(b1 b2) h w c -> (b1 h) (b2 w) c', b1=2, b2=2)

Raises the error:

EinopsError:  Error while processing rearrange-reduction pattern "(b1 b2) h w c -> (b1 h) (b2 w) c".
 Input tensor shape: torch.Size([3, 337, 225]). Additional info: {'b1': 2, 'b2': 2}.
 Expected 4 dimensions, got 3

The error message seems pretty obvious, since I'm specifying 4 patches the output should be of the dimensions (patches, c, h , w). However, I'm not sure where am I supposed to specify that. I went over the tutorials by einops but I still didn't really find what is wrong here.

Answer 1

Einops message is exactly right in this case: in your pattern input has 4 dimensions: '(b1 b2) h w c', but you provide a tensor with only three dimensions.

You should process a batch of 4 images and stack them to get a 4 dim tensor. Your current transformation gets 3-dim tensor (single image) and returns 3-dim tensor.

Answer 2

I think the issue is that einops expects the batch dimension to be the first dimension by default. So, your input tensor shape (3, 337, 225) is not of the expected shape (4, 3, 337, 225). You can either rearrange your input tensor so that the batch dimension is the first dimension or you can specify the batch_dim argument to rearrange.