Hello. I was going through the attention implementation and had a question about the operation by the einsum function.
I am not very familiar with the Einstein summation convention and I believe that there are many others like myself.
I would be very grateful if explanations about the operation were added as comments.
As it is, I find it difficult to understand which parts correspond to what in the paper.
Many thanks if you could help me out.

In Equation 2 of the paper, the query and the key are fed into inner-product operation, instead of point multiplication.

So the follow line

Hi, are the following codes wrong?
"v_out_h, v_out_w = v_out.split(self.out_channels // 2, dim=1)
v_out = torch.cat((v_out_h + self.rel_h, v_out_w + self.rel_w), dim=1)"
shouldn't it be something as follows?
"k_out_h, k_out_w = k_out.split(self.out_channels // 2, dim=1)
k_out = torch.cat((k_out_h + self.rel_h, k_out_w + self.rel_w), dim=1)"
Because the relative distance embedding is supposed to added to the "key" instead of the "value", right?

after

k_out = k_out.contiguous().view(batch, self.groups, self.out_channels // self.groups, height, width, -1)

it gives error,

RuntimeError: shape '[2, 1, 16, 34, 34, -1]' is invalid for input of size 294912

this is because of use of unfold on k_out before this, so height, width is not consistent.

I believe .sum(dim=2) should be added to the below line:

for dot-product in Equation 2 in the paper.

Hello ,the v_out_h 's shape is [n,c//2,kernel_sizeh,kernel_sizew] ,but the self.rel_h 's shape is [out_channels // 2, 1, 1, kernel_size, 1 ] ,how can you add these two different shape ? hope your early reply!

Thanks for sharing the great idea.
When I read the paper, I have some issues about how to compute the relative positional embeddings, i.e. r_(a-i, b-j) , from a row offset a-i and a column offset b-j.
Is there any explicit formula for the calculating process?
Looking forward to your reply.

Hi, I am facing issues with the Resnet50 model training on CIFAR-10. Even with lr of 0.01 its throwing Nan after around 10 epochs (suddenly), so, I am not quite sure how to train the resnet50 model. Hoping for a quick reply! Thanks.

Hi,

I'm confused about embedding in steam attention

It seems that col and row emb are the same, so features are aggregated across cols and rows with the same softmax values. Isn't emb supposed to be 2 dimensional here? Like this:

        emb = emb_logit_a.unsqueeze(2) + emb_logit_b.unsqueeze(1) # [m, ks, ks] p(m, a, b)
        emb = F.softmax(emb.view(self.m, -1), dim=-1)
        emb = emb.view(self.m, 1, 1, 1, 1, self.kernel_size, self.kernel_size)
        v_out = v_out * emb

I think they use matrix multiplication, not dot product

Hello,
I am testing your Resnet50 model with stem is True and at the first training step, my loss is NaN and the accuracy is decreasing? Is that a bug?

Also I didn't see this problem when I train the model ResNet 26.

I am having issues trying to change the batch size. The code produces a mismatch error each time I change the batch size. Below is the error I am getting;
RuntimeError: size mismatch, m1: [4 x 8192], m2: [2048 x 5] at /pytorch/aten/src/THC/generic/THCTensorMathBlas.cu:283

When I try to load pretrained model, I kept getting this error.

RuntimeError: Error(s) in loading state_dict for Model:
Missing key(s) in state_dict: "layer1.0.conv2.0.rel_h", "layer1.0.conv2.0.rel_w", "layer1.0.conv2.0.key_conv.weight", "layer1.0.conv2.0.query_conv.weight", "layer1.0.conv2.0.value_conv.weight", "layer2.0.conv2.0.rel_h", "layer2.0.conv2.0.rel_w", "layer2.0.conv2.0.key_conv.weight", "layer2.0.conv2.0.query_conv.weight", "layer2.0.conv2.0.value_conv.weight", "layer2.1.conv2.0.rel_h", "layer2.1.conv2.0.rel_w", "layer2.1.conv2.0.key_conv.weight", "layer2.1.conv2.0.query_conv.weight", "layer2.1.conv2.0.value_conv.weight", "layer3.0.conv2.0.rel_h", "layer3.0.conv2.0.rel_w", "layer3.0.conv2.0.key_conv.weight", "layer3.0.conv2.0.query_conv.weight", "layer3.0.conv2.0.value_conv.weight", "layer3.1.conv2.0.rel_h", "layer3.1.conv2.0.rel_w", "layer3.1.conv2.0.key_conv.weight", "layer3.1.conv2.0.query_conv.weight", "layer3.1.conv2.0.value_conv.weight", "layer3.2.conv2.0.rel_h", "layer3.2.conv2.0.rel_w", "layer3.2.conv2.0.key_conv.weight", "layer3.2.conv2.0.query_conv.weight", "layer3.2.conv2.0.value_conv.weight", "layer3.3.conv2.0.rel_h", "layer3.3.conv2.0.rel_w", "layer3.3.conv2.0.key_conv.weight", "layer3.3.conv2.0.query_conv.weight", "layer3.3.conv2.0.value_conv.weight", "layer4.0.conv2.0.rel_h", "layer4.0.conv2.0.rel_w", "layer4.0.conv2.0.key_conv.weight", "layer4.0.conv2.0.query_conv.weight", "layer4.0.conv2.0.value_conv.weight".
Unexpected key(s) in state_dict: "layer1.0.conv2.0.weight", "layer1.0.conv2.0.bias", "layer2.0.conv2.0.weight", "layer2.0.conv2.0.bias", "layer2.1.conv2.0.weight", "layer2.1.conv2.0.bias", "layer3.0.conv2.0.weight", "layer3.0.conv2.0.bias", "layer3.1.conv2.0.weight", "layer3.1.conv2.0.bias", "layer3.2.conv2.0.weight", "layer3.2.conv2.0.bias", "layer3.3.conv2.0.weight", "layer3.3.conv2.0.bias", "layer4.0.conv2.0.weight", "layer4.0.conv2.0.bias".

my code for loading pretrained model is below:

    file_path = 'path to file_ckpt.tar'
    checkpoint = torch.load(file_path)
    
    model.load_state_dict(checkpoint['state_dict'])
    model = nn.DataParallel(model, dim=0)
    model.cuda()
    start_epoch = checkpoint['epoch']
    best_acc = checkpoint['best_acc']

    model_parameters = checkpoint['parameters']

ResNet-26 on CIFAR-10 (8.3M params) using your implementation without any changes fails to fit in memory in HAL and even on Colab - P100, T4 and P4. What was your system configuration for doing the benchmarks?

Hi, thanks for your nice work.
On a large dataset like ImageNet, proposed self-attention mechanism consumes large amounts of memory because of unfold operation (im2col).
One would want to share point-wise feature vector among sliding windows.
Do you have any idea?

in attention.py, class AttentionConv

replacing out = torch.einsum('bnchwk,bnchwk -> bnchw', out, v_out)

with out = (out*v_out).sum(dim=5)

made running time more than 2x faster while training on IMAGENET (2 min vs 53s per 100 step, batchsize 25) which is still 3.5x slower than training normal ResNet on IMAGENET

(Not sure whether this model works for IMAGENET or not)

You create these parameters every time you pass through forward so theirs previous values are lost. I think these parameters should be created once in __init__.

Dear @leaderj1001

Have You success to train with Imagenet Dataset? Because i found this error when use the Imagenet:

RuntimeError: CUDA out of memory. Tried to allocate 3.66 GiB (GPU 0; 10.91 GiB total capacity; 8.04 GiB already allocated; 1.66 GiB free; 8.05 GiB reserved in total by PyTorch)

Thank You

In my case out_channels is not an even number so
v_out_h, v_out_w = v_out.split(self.out_channels // 2, dim=1)
fails, any idea how to fix it?

My task is MRI-based pseudo-CT generation.
I combine the SASA with the CycleGAN(https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix).
Considering the GPU memory usage I halved the number of feature maps.
I found that the network without SASA performs best.
When I add SASA to D, the results are really bad.
Did anyone manage to add SASA to CycleGAN?
Hoping for a reply, Thank you.

Should not the key_conv, query_conv and value_conv be defined with group parameters? I am not clear what group parameters do in the following lines.

Hello, I want to ask you a question. In relative position embeddings , why should the number of channels be divided by 2（out_channels // 2）?

self.rel_h = nn.Parameter(torch.randn(out_channels // 2, 1, 1, kernel_size, 1), requires_grad=True) self.rel_w = nn.Parameter(torch.randn(out_channels // 2, 1, 1, 1, kernel_size), requires_grad=True)

the acc for cifar10 reported in the readme is low, is it consistent with the paper.