I'm trying to understand how to use your attention module based on the figure above and the code below.

From what I understand from the non-local paper, if I have an input feature of m_batchsize, channels, height, width = input_.size(),

then n = m_batchsize*height*width and d = channels.

So in the code below, I should use channels = in_channels, key_channels, value_channels.

But what should the head_countbe? Should it be divisible by the number of channels?

import torch
from torch import nn
from torch.nn import functional as f

class EfficientAttention(nn.Module):
    
    def __init__(self, in_channels, key_channels, head_count, value_channels):
        super().__init__()
        self.in_channels = in_channels
        self.key_channels = key_channels
        self.head_count = head_count
        self.value_channels = value_channels

        self.keys = nn.Conv2d(in_channels, key_channels, 1)
        self.queries = nn.Conv2d(in_channels, key_channels, 1)
        self.values = nn.Conv2d(in_channels, value_channels, 1)
        self.reprojection = nn.Conv2d(value_channels, in_channels, 1)

    def forward(self, input_):
        n, _, h, w = input_.size()
        keys = self.keys(input_).reshape((n, self.key_channels, h * w))
        queries = self.queries(input_).reshape(n, self.key_channels, h * w)
        values = self.values(input_).reshape((n, self.value_channels, h * w))
        head_key_channels = self.key_channels // self.head_count
        head_value_channels = self.value_channels // self.head_count
        
        attended_values = []
        for i in range(self.head_count):
            key = f.softmax(keys[
                :,
                i * head_key_channels: (i + 1) * head_key_channels,
                :
            ], dim=2)
            query = f.softmax(queries[
                :,
                i * head_key_channels: (i + 1) * head_key_channels,
                :
            ], dim=1)
            value = values[
                :,
                i * head_value_channels: (i + 1) * head_value_channels,
                :
            ]
            context = key @ value.transpose(1, 2)
            attended_value = (
                context.transpose(1, 2) @ query
            ).reshape(n, head_value_channels, h, w)
            attended_values.append(attended_value)

        aggregated_values = torch.cat(attended_values, dim=1)
        reprojected_value = self.reprojection(aggregated_values)
        attention = reprojected_value + input_

        return attention

Hi, it's a great job. I'm a rookie for DL.

The EfficientAttention have four parameters: in_channels, key_channels, head_count, value_channels. I know key_channels = d_k, value_channels = d_v. What are head_count means? What is a common setting for four parameters?

After I ran the code you provided, I compared it with the traditional self-attention mechanism. I found the results are quite different from traditional ones which are truly normalized, but the efficient one is not. Can you offer some help?

i am trying to use the efficient attention. I am confused about the params. I change the key_channels and the values, but the output dimension keeps the same, why?

Hi,

Would you please tell me how to use efficient attention to be similar to the scaled dot-prodcut attention?

I notice that you apply softmax for both query and key. So is it right to set temperature = (d_k) ** 0.25 and apply key = f.softmax(keys[ :, i * head_key_channels: (i + 1) * head_key_channels, : ] / temperature, dim=2) and query = f.softmax(queries[ :, i * head_key_channels: (i + 1) * head_key_channels, : ] / temperature, dim=1) for both query and key in your code to make it similar to the scaled dot-product attention?

Thank you!

Hi, I found the attention map computed using the script is C x C. Shouldn't it be (H x W) x (H x W) if we want spatial attention?

Thank you for any information that you can provide.

In Seq2Seq models, it is common to apply mask to remove paddings, or to remove future inputs of a causal model. Is it possible to do so in Efficient Attention, as it does not have a key seq to query seq mapping?

Hello, Thanks for the excellent work. I would like to show different attention maps produced using different query positions (reference bouding boxes). For example when there are three bouding boxes on the image I would like to calculate the attention w.r.t. them and show the difference. Is it possible to do this?

Hi,

I'm slightly confused as to what the difference between these two in figure 7 is in the paper? PSMNet (original) 1.09 0 PSMNet (baseline) 0.51 0

Thanks, Oliver