| 79 | |
| 80 | # from https://github.com/facebookresearch/ConvNeXt-V2/blob/3608f67cc1dae164790c5d0aead7bf2d73d9719b/models/utils.py#L105 |
| 81 | class GlobalResponseNorm(nn.Module): |
| 82 | def __init__(self, dim): |
| 83 | super().__init__() |
| 84 | self.gamma = nn.Parameter(torch.zeros(1, 1, 1, dim)) |
| 85 | self.beta = nn.Parameter(torch.zeros(1, 1, 1, dim)) |
| 86 | |
| 87 | def forward(self, x): |
| 88 | agg_norm = torch.norm(x, p=2, dim=(1, 2), keepdim=True) |
| 89 | stand_div_norm = agg_norm / (agg_norm.mean(dim=-1, keepdim=True) + 1e-6) |
| 90 | return self.gamma * (x * stand_div_norm) + self.beta + x |
| 91 | |
| 92 | |
| 93 | class SDCascadeAttnBlock(nn.Module): |