| 213 | |
| 214 | |
| 215 | class SqueezeExcite(Module): |
| 216 | # global context network - attention-esque squeeze-excite variant (https://arxiv.org/abs/2012.13375) |
| 217 | |
| 218 | def __init__(self, dim, *, dim_out=None, dim_hidden_min=16, init_bias=-10): |
| 219 | super().__init__() |
| 220 | dim_out = default(dim_out, dim) |
| 221 | |
| 222 | self.to_k = nn.Conv2d(dim, 1, 1) |
| 223 | dim_hidden = max(dim_hidden_min, dim_out // 2) |
| 224 | |
| 225 | self.net = nn.Sequential( |
| 226 | nn.Conv2d(dim, dim_hidden, 1), nn.LeakyReLU(0.1), nn.Conv2d(dim_hidden, dim_out, 1), nn.Sigmoid() |
| 227 | ) |
| 228 | |
| 229 | nn.init.zeros_(self.net[-2].weight) |
| 230 | nn.init.constant_(self.net[-2].bias, init_bias) |
| 231 | |
| 232 | def forward(self, x): |
| 233 | orig_input, batch = x, x.shape[0] |
| 234 | is_video = x.ndim == 5 |
| 235 | |
| 236 | if is_video: |
| 237 | x = rearrange(x, "b c f h w -> (b f) c h w") |
| 238 | |
| 239 | context = self.to_k(x) |
| 240 | |
| 241 | context = rearrange(context, "b c h w -> b c (h w)").softmax(dim=-1) |
| 242 | spatial_flattened_input = rearrange(x, "b c h w -> b c (h w)") |
| 243 | |
| 244 | out = einsum("b i n, b c n -> b c i", context, spatial_flattened_input) |
| 245 | out = rearrange(out, "... -> ... 1") |
| 246 | gates = self.net(out) |
| 247 | |
| 248 | if is_video: |
| 249 | gates = rearrange(gates, "(b f) c h w -> b c f h w", b=batch) |
| 250 | |
| 251 | return gates * orig_input |
| 252 | |
| 253 | |
| 254 | # token shifting |