| 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. |