# coding=utf-8
# Copyright 2022 The IDEA Authors. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ------------------------------------------------------------------------------------------------
# Modified from:
# https://github.com/facebookresearch/ConvNeXt/blob/main/models/convnext.py
# ------------------------------------------------------------------------------------------------
import torch
import torch.nn as nn
import torch.nn.functional as F
[docs]class LayerNorm(nn.Module):
r"""LayerNorm which supports both channel_last (default) and channel_first data format.
The inputs data format should be as follows:
- channel_last: (bs, h, w, channels)
- channel_first: (bs, channels, h, w)
Args:
normalized_shape (tuple): The size of the input feature dim.
eps (float): A value added to the denominator for
numerical stability. Default: True.
channel_last (bool): Set True for `channel_last` input data
format. Default: True.
"""
def __init__(self, normalized_shape, eps=1e-6, channel_last=True):
super().__init__()
self.weight = nn.Parameter(torch.ones(normalized_shape))
self.bias = nn.Parameter(torch.zeros(normalized_shape))
self.eps = eps
self.channel_last = channel_last
self.normalized_shape = (normalized_shape,)
[docs] def forward(self, x):
"""Forward function for `LayerNorm`"""
if self.channel_last:
return F.layer_norm(x, self.normalized_shape, self.weight, self.bias, self.eps)
else:
u = x.mean(1, keepdim=True)
s = (x - u).pow(2).mean(1, keepdim=True)
x = (x - u) / torch.sqrt(s + self.eps)
x = self.weight[:, None, None] * x + self.bias[:, None, None]
return x