File size: 2,572 Bytes
d80b3db |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 |
# Copyright (c) Meta Platforms, Inc. and affiliates.
# Copyright 2025 AI for Oncology Research Group. 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.
#
# References:
# https://github.com/facebookresearch/dino/blob/master/vision_transformer.py
# https://github.com/rwightman/pytorch-image-models/tree/master/timm/models/vision_transformer.py
from typing import Union
import torch
from torch import nn
class LayerScale(nn.Module):
"""Layer scale module for scaling the output of a layer.
Parameters
----------
dim : int
Dimension of the layer scale.
init_values : float or torch.Tensor, optional
Initial values for the layer scale, by default 1e-5. If a tensor is provided, it should have shape (dim,).
inplace : bool, optional
Whether to perform the operation in-place, by default False.
"""
def __init__(
self,
dim: int,
init_values: Union[float, torch.Tensor] = 1e-5,
inplace: bool = False,
) -> None:
"""Inits :class:`LayerScale
Parameters
----------
dim : int
Dimension of the layer scale.
init_values : float or torch.Tensor, optional
Initial values for the layer scale, by default 1e-5. If a tensor is provided, it should have shape (dim,).
inplace : bool, optional
Whether to perform the operation in-place, by default False.
"""
super().__init__()
self.inplace = inplace
self.gamma = nn.Parameter(init_values * torch.ones(dim))
def forward(self, x: torch.Tensor) -> torch.Tensor:
"""Forward pass of :class:`LayerScale`.
Parameters
----------
x : torch.Tensor
Input tensor of shape (B, N, C) where B is the batch size, N is the sequence length, and C is
the feature dimension.
Returns
-------
torch.Tensor
Scaled output tensor of shape (B, N, C).
"""
return x.mul_(self.gamma) if self.inplace else x * self.gamma
|