You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
68 lines
1.7 KiB
68 lines
1.7 KiB
# Copyright (c) OpenMMLab. All rights reserved.
|
|
import numpy as np
|
|
import torch
|
|
import torch.nn as nn
|
|
from mmcv.cnn import build_conv_layer, build_upsample_layer
|
|
|
|
|
|
def test_build_upsample_layer():
|
|
layer1 = nn.ConvTranspose2d(
|
|
in_channels=3,
|
|
out_channels=10,
|
|
kernel_size=3,
|
|
stride=2,
|
|
padding=1,
|
|
output_padding=1,
|
|
bias=False)
|
|
|
|
layer2 = build_upsample_layer(
|
|
dict(type='deconv'),
|
|
in_channels=3,
|
|
out_channels=10,
|
|
kernel_size=3,
|
|
stride=2,
|
|
padding=1,
|
|
output_padding=1,
|
|
bias=False)
|
|
layer2.load_state_dict(layer1.state_dict())
|
|
|
|
input_shape = (1, 3, 32, 32)
|
|
inputs = _demo_inputs(input_shape)
|
|
out1 = layer1(inputs)
|
|
out2 = layer2(inputs)
|
|
assert torch.equal(out1, out2)
|
|
|
|
|
|
def test_build_conv_layer():
|
|
layer1 = nn.Conv2d(
|
|
in_channels=3, out_channels=10, kernel_size=3, stride=1, padding=1)
|
|
|
|
layer2 = build_conv_layer(
|
|
cfg=dict(type='Conv2d'),
|
|
in_channels=3,
|
|
out_channels=10,
|
|
kernel_size=3,
|
|
stride=1,
|
|
padding=1)
|
|
|
|
layer2.load_state_dict(layer1.state_dict())
|
|
|
|
input_shape = (1, 3, 32, 32)
|
|
inputs = _demo_inputs(input_shape)
|
|
out1 = layer1(inputs)
|
|
out2 = layer2(inputs)
|
|
assert torch.equal(out1, out2)
|
|
|
|
|
|
def _demo_inputs(input_shape=(1, 3, 64, 64)):
|
|
"""Create a superset of inputs needed to run backbone.
|
|
|
|
Args:
|
|
input_shape (tuple): input batch dimensions.
|
|
Default: (1, 3, 64, 64).
|
|
Returns:
|
|
Random input tensor with the size of input_shape.
|
|
"""
|
|
inps = np.random.random(input_shape)
|
|
inps = torch.FloatTensor(inps)
|
|
return inps
|
|
|