Output the attention weights from MILModel

[DocandBean]

Hi MONAI Universe,

Is there a method to output the attention weights from the MIL Model?

Given input data of the form [B, N, C, H, W], I would expect an N x 1 output array with the weights ranging 0-1. See for instance Figure 4 of Attention-based Deep Multiple Instance Learning.

Any guidance would be much appreciated!

[wyli]

Hello, there's a flag no_head to skip the bag level fusing, perhaps you can set it to False to get BxNxfeatures results

MONAI/monai/networks/nets/milmodel.py

Line 233 in ca90628

def forward(self, x: torch.Tensor, no_head: bool = False) -> torch.Tensor:

then you can compute the output scores accordingly as implemented here

MONAI/monai/networks/nets/milmodel.py

Line 174 in ca90628

def calc_head(self, x: torch.Tensor) -> torch.Tensor:

[DocandBean]

Thanks for the suggestion. Worked perfectly.

[relyativist]

@DocandBean @wyli The proposed method in answer will not output you attention weights. Setting the no_head=True will just output you the weights after backbone CNN. If you want attention matrix, you have to change the TransformerEncoderLayer, because the torch implementation of transformer can't output you the attention matrix. For instance you can implement self-attention:

class MultiheadSelfAttentionBlock(nn.Module):
    """
    Multihead self attention block
    params:
        dim: input dimension
        heads: number of heads
        dim_head: dimension of each head
        dropout: dropout rate
    """
    def __init__(self,
                 embedding_dim : int,
                 heads : int = 8,
                 dropout : float = 0.):
        super().__init__()
        self.layer_norm = nn.LayerNorm(normalized_shape=embedding_dim)
        self.mhsa = nn.MultiheadAttention(embedding_dim, heads, dropout)
    
    def forward(self, x):
        x = self.layer_norm(x)
        attn_output, attn_weights = self.mhsa(x, x, x, need_weights=True, average_attn_weights=False)
        return attn_output, attn_weights

where attn_output, attn_weights = self.mhsa(x, x, x, need_weights=True, average_attn_weights=False) will output you the attention matrices from each head.