Focal Modulation Args: dim (int): Number of input channels. proj_drop (float, optional): Dropout ratio of output. Default: 0.0 focal_level (int): Number of focal levels focal_window (int): Focal window size at focal level 1 focal_factor (int, default=2):
| 42 | return x |
| 43 | |
| 44 | class FocalModulation(nn.Module): |
| 45 | """ Focal Modulation |
| 46 | |
| 47 | Args: |
| 48 | dim (int): Number of input channels. |
| 49 | proj_drop (float, optional): Dropout ratio of output. Default: 0.0 |
| 50 | focal_level (int): Number of focal levels |
| 51 | focal_window (int): Focal window size at focal level 1 |
| 52 | focal_factor (int, default=2): Step to increase the focal window |
| 53 | use_postln (bool, default=False): Whether use post-modulation layernorm |
| 54 | """ |
| 55 | |
| 56 | def __init__(self, dim, proj_drop=0., focal_level=2, focal_window=7, focal_factor=2, use_postln=False, use_postln_in_modulation=False, scaling_modulator=False): |
| 57 | |
| 58 | super().__init__() |
| 59 | self.dim = dim |
| 60 | |
| 61 | # specific args for focalv3 |
| 62 | self.focal_level = focal_level |
| 63 | self.focal_window = focal_window |
| 64 | self.focal_factor = focal_factor |
| 65 | self.use_postln_in_modulation = use_postln_in_modulation |
| 66 | self.scaling_modulator = scaling_modulator |
| 67 | |
| 68 | self.f = nn.Linear(dim, 2*dim+(self.focal_level+1), bias=True) |
| 69 | self.h = nn.Conv2d(dim, dim, kernel_size=1, stride=1, padding=0, groups=1, bias=True) |
| 70 | |
| 71 | self.act = nn.GELU() |
| 72 | self.proj = nn.Linear(dim, dim) |
| 73 | self.proj_drop = nn.Dropout(proj_drop) |
| 74 | self.focal_layers = nn.ModuleList() |
| 75 | |
| 76 | if self.use_postln_in_modulation: |
| 77 | self.ln = nn.LayerNorm(dim) |
| 78 | |
| 79 | for k in range(self.focal_level): |
| 80 | kernel_size = self.focal_factor*k + self.focal_window |
| 81 | self.focal_layers.append( |
| 82 | nn.Sequential( |
| 83 | nn.Conv2d(dim, dim, kernel_size=kernel_size, stride=1, groups=dim, |
| 84 | padding=kernel_size//2, bias=False), |
| 85 | nn.GELU(), |
| 86 | ) |
| 87 | ) |
| 88 | |
| 89 | def forward(self, x): |
| 90 | """ Forward function. |
| 91 | |
| 92 | Args: |
| 93 | x: input features with shape of (B, H, W, C) |
| 94 | """ |
| 95 | B, nH, nW, C = x.shape |
| 96 | x = self.f(x) |
| 97 | x = x.permute(0, 3, 1, 2).contiguous() |
| 98 | q, ctx, gates = torch.split(x, (C, C, self.focal_level+1), 1) |
| 99 | |
| 100 | ctx_all = 0 |
| 101 | for l in range(self.focal_level): |