Shortcuts

Source code for mmcv.cnn.bricks.scale

# Copyright (c) OpenMMLab. All rights reserved.
import torch
import torch.nn as nn


[docs]class Scale(nn.Module): """A learnable scale parameter. This layer scales the input by a learnable factor. It multiplies a learnable scale parameter of shape (1,) with input of any shape. Args: scale (float): Initial value of scale factor. Default: 1.0 """ def __init__(self, scale: float = 1.0): super().__init__() self.scale = nn.Parameter(torch.tensor(scale, dtype=torch.float))
[docs] def forward(self, x: torch.Tensor) -> torch.Tensor: return x * self.scale
class LayerScale(nn.Module): """LayerScale layer. Args: dim (int): Dimension of input features. inplace (bool): Whether performs operation in-place. Default: `False`. data_format (str): The input data format, could be 'channels_last' or 'channels_first', representing (B, C, H, W) and (B, N, C) format data respectively. Default: 'channels_last'. scale (float): Initial value of scale factor. Default: 1.0 """ def __init__(self, dim: int, inplace: bool = False, data_format: str = 'channels_last', scale: float = 1e-5): super().__init__() assert data_format in ('channels_last', 'channels_first'), \ "'data_format' could only be channels_last or channels_first." self.inplace = inplace self.data_format = data_format self.weight = nn.Parameter(torch.ones(dim) * scale) def forward(self, x) -> torch.Tensor: if self.data_format == 'channels_first': shape = tuple((1, -1, *(1 for _ in range(x.dim() - 2)))) else: shape = tuple((*(1 for _ in range(x.dim() - 1)), -1)) if self.inplace: return x.mul_(self.weight.view(*shape)) else: return x * self.weight.view(*shape)
Read the Docs v: latest
Versions
master
latest
2.x
1.x
v1.7.0
v1.6.2
v1.6.1
v1.6.0
v1.5.3
v1.5.2_a
v1.5.1
v1.5.0
v1.4.8
v1.4.7
v1.4.6
v1.4.5
v1.4.4
v1.4.3
v1.4.2
v1.4.1
v1.4.0
v1.3.18
v1.3.17
v1.3.16
v1.3.15
v1.3.14
v1.3.13
v1.3.12
v1.3.11
v1.3.10
v1.3.9
v1.3.8
v1.3.7
v1.3.6
v1.3.5
v1.3.4
v1.3.3
v1.3.2
v1.3.1
v1.3.0
Downloads
pdf
html
epub
On Read the Docs
Project Home
Builds

Free document hosting provided by Read the Docs.