| 32 | |
| 33 | |
| 34 | class FocalModulation(nn.Module): |
| 35 | def __init__(self, dim, focal_window, focal_level, focal_factor=2, bias=True, proj_drop=0., |
| 36 | use_postln_in_modulation=False, normalize_modulator=False): |
| 37 | super().__init__() |
| 38 | |
| 39 | self.dim = dim |
| 40 | self.focal_window = focal_window |
| 41 | self.focal_level = focal_level |
| 42 | self.focal_factor = focal_factor |
| 43 | self.use_postln_in_modulation = use_postln_in_modulation |
| 44 | self.normalize_modulator = normalize_modulator |
| 45 | |
| 46 | self.f = nn.Linear(dim, 2 * dim + (self.focal_level + 1), bias=bias) |
| 47 | self.h = nn.Conv2d(dim, dim, kernel_size=1, stride=1, bias=bias) |
| 48 | |
| 49 | self.act = nn.GELU() |
| 50 | self.proj = nn.Linear(dim, dim) |
| 51 | self.proj_drop = nn.Dropout(proj_drop) |
| 52 | self.focal_layers = nn.ModuleList() |
| 53 | |
| 54 | self.kernel_sizes = [] |
| 55 | for k in range(self.focal_level): |
| 56 | kernel_size = self.focal_factor * k + self.focal_window |
| 57 | self.focal_layers.append( |
| 58 | nn.Sequential( |
| 59 | nn.Conv2d(dim, dim, kernel_size=kernel_size, stride=1, |
| 60 | padding=kernel_size // 2, bias=False), |
| 61 | nn.GELU(), |
| 62 | ) |
| 63 | ) |
| 64 | self.kernel_sizes.append(kernel_size) |
| 65 | if self.use_postln_in_modulation: |
| 66 | self.ln = nn.LayerNorm(dim) |
| 67 | |
| 68 | def forward(self, x): |
| 69 | """ |
| 70 | Args: |
| 71 | x: input features with shape of (B, H, W, C) |
| 72 | """ |
| 73 | C = x.shape[-1] |
| 74 | |
| 75 | # pre linear projection |
| 76 | x = self.f(x).permute(0, 3, 1, 2).contiguous() |
| 77 | q, ctx, self.gates = torch.split(x, (C, C, self.focal_level + 1), 1) |
| 78 | |
| 79 | # context aggreation |
| 80 | ctx_all = 0 |
| 81 | for l in range(self.focal_level): |
| 82 | ctx = self.focal_layers[l](ctx) |
| 83 | ctx_all = ctx_all + ctx * self.gates[:, l:l + 1] |
| 84 | ctx_global = self.act(ctx.mean(2, keepdim=True).mean(3, keepdim=True)) |
| 85 | ctx_all = ctx_all + ctx_global * self.gates[:, self.focal_level:] |
| 86 | |
| 87 | # normalize context |
| 88 | if self.normalize_modulator: |
| 89 | ctx_all = ctx_all / (self.focal_level + 1) |
| 90 | |
| 91 | # focal modulation |