| 123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105 |
- """ Gather-Excite Attention Block
- Paper: `Gather-Excite: Exploiting Feature Context in CNNs` - https://arxiv.org/abs/1810.12348
- Official code here, but it's only partial impl in Caffe: https://github.com/hujie-frank/GENet
- I've tried to support all of the extent both w/ and w/o params. I don't believe I've seen another
- impl that covers all of the cases.
- NOTE: extent=0 + extra_params=False is equivalent to Squeeze-and-Excitation
- Hacked together by / Copyright 2021 Ross Wightman
- """
- from typing import Optional, Tuple, Type, Union
- import math
- from torch import nn as nn
- import torch.nn.functional as F
- from .create_act import create_act_layer, get_act_layer
- from .create_conv2d import create_conv2d
- from .helpers import make_divisible
- from .mlp import ConvMlp
- class GatherExcite(nn.Module):
- """ Gather-Excite Attention Module
- """
- def __init__(
- self,
- channels: int,
- feat_size: Optional[Tuple[int, int]] = None,
- extra_params: bool = False,
- extent: int = 0,
- use_mlp: bool = True,
- rd_ratio: float = 1./16,
- rd_channels: Optional[int] = None,
- rd_divisor: int = 1,
- add_maxpool: bool = False,
- act_layer: Type[nn.Module] = nn.ReLU,
- norm_layer: Type[nn.Module] = nn.BatchNorm2d,
- gate_layer: Union[str, Type[nn.Module]] = 'sigmoid',
- device=None,
- dtype=None,
- ):
- dd = {'device': device, 'dtype': dtype}
- super().__init__()
- self.add_maxpool = add_maxpool
- act_layer = get_act_layer(act_layer)
- self.extent = extent
- if extra_params:
- self.gather = nn.Sequential()
- if extent == 0:
- assert feat_size is not None, 'spatial feature size must be specified for global extent w/ params'
- self.gather.add_module(
- 'conv1', create_conv2d(channels, channels, kernel_size=feat_size, stride=1, depthwise=True, *dd))
- if norm_layer:
- self.gather.add_module(f'norm1', nn.BatchNorm2d(channels, *dd))
- else:
- assert extent % 2 == 0
- num_conv = int(math.log2(extent))
- for i in range(num_conv):
- self.gather.add_module(
- f'conv{i + 1}',
- create_conv2d(channels, channels, kernel_size=3, stride=2, depthwise=True, *dd))
- if norm_layer:
- self.gather.add_module(f'norm{i + 1}', nn.BatchNorm2d(channels, *dd))
- if i != num_conv - 1:
- self.gather.add_module(f'act{i + 1}', act_layer(inplace=True))
- else:
- self.gather = None
- if self.extent == 0:
- self.gk = 0
- self.gs = 0
- else:
- assert extent % 2 == 0
- self.gk = self.extent * 2 - 1
- self.gs = self.extent
- if not rd_channels:
- rd_channels = make_divisible(channels * rd_ratio, rd_divisor, round_limit=0.)
- self.mlp = ConvMlp(channels, rd_channels, act_layer=act_layer, *dd) if use_mlp else nn.Identity()
- self.gate = create_act_layer(gate_layer)
- def forward(self, x):
- size = x.shape[-2:]
- if self.gather is not None:
- x_ge = self.gather(x)
- else:
- if self.extent == 0:
- # global extent
- x_ge = x.mean(dim=(2, 3), keepdims=True)
- if self.add_maxpool:
- # experimental codepath, may remove or change
- x_ge = 0.5 * x_ge + 0.5 * x.amax((2, 3), keepdim=True)
- else:
- x_ge = F.avg_pool2d(
- x, kernel_size=self.gk, stride=self.gs, padding=self.gk // 2, count_include_pad=False)
- if self.add_maxpool:
- # experimental codepath, may remove or change
- x_ge = 0.5 * x_ge + 0.5 * F.max_pool2d(x, kernel_size=self.gk, stride=self.gs, padding=self.gk // 2)
- x_ge = self.mlp(x_ge)
- if x_ge.shape[-1] != 1 or x_ge.shape[-2] != 1:
- x_ge = F.interpolate(x_ge, size=size)
- return x * self.gate(x_ge)
|
