File size: 577 Bytes
18dd6ad |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 |
# Copyright (c) OpenMMLab. All rights reserved.
import torch
import torch.nn as nn
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=1.0):
super(Scale, self).__init__()
self.scale = nn.Parameter(torch.tensor(scale, dtype=torch.float))
def forward(self, x):
return x * self.scale
|