File size: 603 Bytes
bc90483 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 | # Copyright (c) Meta Platforms, Inc. and affiliates.
#
# This software may be used and distributed in accordance with
# the terms of the DINOv3 License Agreement.
import torch.nn as nn
class LayerNorm2D(nn.Module):
def __init__(self, normalized_shape, norm_layer=nn.LayerNorm):
super().__init__()
self.ln = norm_layer(normalized_shape) if norm_layer is not None else nn.Identity()
def forward(self, x):
"""
x: N C H W
"""
x = x.permute(0, 2, 3, 1)
x = self.ln(x)
x = x.permute(0, 3, 1, 2)
return x
|