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