ptame.models.components.attention.ttame_attention

  1from typing import List
  2
  3import torch
  4from torch import nn
  5from torch.nn import functional as F
  6
  7from ptame.utils.masking_utils import minmax_4d
  8
  9
 10class TtameAttention(nn.Module):
 11    """TtameAttention is a class that represents the attention mechanism used
 12    in TTAME."""
 13
 14    def __init__(self, ft_size: List[torch.Size], num_classes=1000):
 15        super().__init__()
 16        feat_height = ft_size[0][2] if ft_size[0][2] <= 56 else 56
 17        self.resolution = feat_height
 18        self.interpolate = lambda inp: F.interpolate(
 19            inp,
 20            size=(feat_height, feat_height),
 21            mode="bilinear",
 22            align_corners=False,
 23        )
 24        in_channels_list = [o[1] for o in ft_size]
 25        self.channels = in_channels_list
 26        # noinspection PyTypeChecker
 27        self.convs = nn.ModuleList(
 28            [
 29                nn.Conv2d(
 30                    in_channels=in_channels,
 31                    out_channels=in_channels,
 32                    kernel_size=1,
 33                    padding=0,
 34                    bias=True,
 35                )
 36                for in_channels in in_channels_list
 37            ]
 38        )
 39        self.bn_channels = in_channels_list
 40        self.bns = nn.ModuleList(
 41            [nn.BatchNorm2d(channels) for channels in self.bn_channels]
 42        )
 43        self.relu = nn.ReLU()
 44        # for each extra layer we need 1000 more channels to input to the fuse
 45        # convolution
 46        fuse_channels = sum(in_channels_list)
 47        # noinspection PyTypeChecker
 48        self.fuser = nn.Conv2d(
 49            in_channels=fuse_channels,
 50            out_channels=num_classes,
 51            kernel_size=1,
 52            padding=0,
 53            bias=True,
 54        )
 55        self.num_classes = num_classes
 56
 57    def forward(self, features, **kwargs):
 58        """Forward pass of the attention mechanism."""
 59        # Fusion Strategy
 60        feature_maps = features
 61        if kwargs.get("opticam", False):
 62            return features[0]
 63        # Now all feature map sets are of the same HxW
 64        # conv
 65        class_maps = [
 66            op(feature) for op, feature in zip(self.convs, feature_maps)
 67        ]
 68        # batch norm
 69        class_maps = [op(feature) for op, feature in zip(self.bns, class_maps)]
 70        # add (skip connection)
 71        class_maps = [
 72            class_map + feature_map
 73            for class_map, feature_map in zip(class_maps, feature_maps)
 74        ]
 75        # activation
 76        class_maps = [self.relu(class_map) for class_map in class_maps]
 77        # upscale
 78        class_maps = [self.interpolate(feature) for feature in class_maps]
 79        # concat
 80        class_maps = torch.cat(class_maps, 1)
 81        # fuse into num_classes channels
 82        c = self.fuser(class_maps)  # batch_size x1xWxH
 83        if not self.training:
 84            return minmax_4d(c)
 85        else:
 86            return torch.sigmoid(c)
 87
 88    @torch.no_grad()
 89    def get_contributions(self):
 90        """Calculate the contributions of the attention mechanism."""
 91        for name, param in self.fuser.named_parameters():
 92            if "weight" in name:
 93                weights = param.squeeze()
 94        channels = self.channels.copy()
 95        contribs = torch.stack(
 96            [
 97                a.sum(dim=1)
 98                for a in weights.softmax(dim=1).split(channels, dim=1)
 99            ],
100            dim=1,
101        )
102        return [f"{i}" for i in range(len(self.channels))], contribs