mirror of
https://github.com/open-mmlab/mmcv.git
synced 2025-06-03 21:54:52 +08:00
45111e193d
* add building bricks of cnn * add unit tests * use registry for building bricks * minor updates * add scale layer * add test for scale * add doc string Co-authored-by: Jiarui XU <xvjiarui0826@gmail.com>
22 lines
511 B
Python
22 lines
511 B
Python
import torch
|
|
|
|
from mmcv.cnn.bricks import Scale
|
|
|
|
|
|
def test_scale():
|
|
# test default scale
|
|
scale = Scale()
|
|
assert scale.scale.data == 1.
|
|
assert scale.scale.dtype == torch.float
|
|
x = torch.rand(1, 3, 64, 64)
|
|
output = scale(x)
|
|
assert output.shape == (1, 3, 64, 64)
|
|
|
|
# test given scale
|
|
scale = Scale(10.)
|
|
assert scale.scale.data == 10.
|
|
assert scale.scale.dtype == torch.float
|
|
x = torch.rand(1, 3, 64, 64)
|
|
output = scale(x)
|
|
assert output.shape == (1, 3, 64, 64)
|