def test_standardize_zeros() -> None:
"""
Make sure there are no divide-by-zero errors.
"""
size = (5, 3, 3, 3)
weights = torch.zeros(size)
result = WeightStandardizedConv2d.standardize(weights)
assert result.shape == weights.shape
assert torch.allclose(result, torch.zeros(size=size))
def test_standardize_ones() -> None:
"""
Smoke test for normalization.
"""
size = (5, 3, 3, 3)
weights = torch.ones(size)
result = WeightStandardizedConv2d.standardize(weights)
assert result.shape == weights.shape
assert torch.allclose(result, torch.zeros(size=size))
def __init__(self,
in_channels: int,
out_channels: int,
bottleneck_channels: int,
num_groups: int,
downsample_stride: int = 1):
super().__init__()
self.gn1 = nn.GroupNorm(num_groups=num_groups,
num_channels=in_channels)
self.conv1 = WeightStandardizedConv2d(in_channels=in_channels,
out_channels=bottleneck_channels,
kernel_size=(1, 1),
stride=(1, 1),
bias=False)
self.gn2 = nn.GroupNorm(num_groups=num_groups,
num_channels=bottleneck_channels)
self.conv2 = WeightStandardizedConv2d(in_channels=bottleneck_channels,
out_channels=bottleneck_channels,
kernel_size=(3, 3),
stride=(downsample_stride,
downsample_stride),
padding=(1, 1),
bias=False)
self.gn3 = nn.GroupNorm(num_groups=num_groups,
num_channels=bottleneck_channels)
self.conv3 = WeightStandardizedConv2d(in_channels=bottleneck_channels,
out_channels=out_channels,
kernel_size=(1, 1),
stride=(1, 1),
bias=False)
self.relu = nn.ReLU(inplace=True)
if in_channels != out_channels:
self.downsample: Optional[WeightStandardizedConv2d] = \
WeightStandardizedConv2d(in_channels=in_channels,
out_channels=out_channels,
kernel_size=(1, 1),
stride=(downsample_stride, downsample_stride),
bias=False)
else:
self.downsample = None
def test_standardize_rows() -> None:
"""
We find mean and variance for each filter, so a filter filled with a constant value should normalize to 0.
"""
size = (5, 3, 3, 3)
weights = torch.ones(size)
for i in range(weights.shape[0]):
weights[i] = i
result = WeightStandardizedConv2d.standardize(weights)
assert result.shape == weights.shape
assert torch.allclose(result, torch.zeros(size=size))
def __init__(self,
num_groups: int = 32,
num_classes: int = 21843,
num_blocks_in_layer: Tuple[int, int, int,
int] = (3, 4, 23, 3),
width_factor: int = 1):
super().__init__()
self.initial = nn.Sequential(
WeightStandardizedConv2d(in_channels=3,
out_channels=64 * width_factor,
kernel_size=(7, 7),
stride=(2, 2),
padding=(3, 3),
bias=False),
nn.ConstantPad2d(padding=1, value=0),
nn.MaxPool2d(kernel_size=3, stride=2, padding=0, dilation=1))
self.conv_stack = nn.Sequential(
ResNetV2Layer(in_channels=64 * width_factor,
out_channels=256 * width_factor,
bottleneck_channels=64 * width_factor,
num_groups=num_groups,
downsample_stride=1,
num_blocks=num_blocks_in_layer[0]),
ResNetV2Layer(in_channels=256 * width_factor,
out_channels=512 * width_factor,
bottleneck_channels=128 * width_factor,
num_groups=num_groups,
downsample_stride=2,
num_blocks=num_blocks_in_layer[1]),
ResNetV2Layer(in_channels=512 * width_factor,
out_channels=1024 * width_factor,
bottleneck_channels=256 * width_factor,
num_groups=num_groups,
downsample_stride=2,
num_blocks=num_blocks_in_layer[2]),
ResNetV2Layer(in_channels=1024 * width_factor,
out_channels=2048 * width_factor,
bottleneck_channels=512 * width_factor,
num_groups=num_groups,
downsample_stride=2,
num_blocks=num_blocks_in_layer[3]))
self.linear = nn.Sequential(
nn.GroupNorm(num_groups=num_groups,
num_channels=2048 * width_factor),
nn.ReLU(inplace=True), nn.AdaptiveAvgPool2d(output_size=1),
nn.Conv2d(in_channels=2048 * width_factor,
out_channels=num_classes,
kernel_size=(1, 1),
bias=True))
def test_standardize_random() -> None:
"""
Test normalization on arbitrary weights.
"""
size = (5, 3, 3, 3)
random = np.random.randint(low=0, high=100, size=size).astype(np.float)
weights = torch.from_numpy(random)
expected = np.copy(random)
mean = expected.mean(axis=(1, 2, 3), keepdims=True)
# this test also makes sure we are not using an unbiased estimate of variance in the convolution layer:
# Torch uses unbiased estimates by default
var = expected.var(axis=(1, 2, 3), ddof=0, keepdims=True)
expected = (expected - mean) / np.sqrt(var + eps)
expected = torch.from_numpy(expected)
result = WeightStandardizedConv2d.standardize(weights)
assert result.shape == weights.shape
assert torch.allclose(result, expected)
def test_conv_constant_filter() -> None:
"""
Test with filters filled with a constant value: they'll normalize to zero
"""
in_channels = 3
out_channels = 5
batch_size = 1
image_size = (batch_size, in_channels, 20, 20)
kernel_size = (3, 3)
expected_output_size = (batch_size, out_channels,
image_size[2] - kernel_size[0] + 1,
image_size[3] - kernel_size[1] + 1)
conv_layer = WeightStandardizedConv2d(in_channels=in_channels,
out_channels=out_channels,
kernel_size=kernel_size,
bias=True)
conv_layer.weight.data.fill_(1)
assert conv_layer.bias is not None # for mypy
conv_layer.bias.data.fill_(0) # we aren't normalizing bias
image = torch.ones(size=image_size)
result = conv_layer(image)
assert result.shape == expected_output_size
assert torch.allclose(result, torch.zeros(size=expected_output_size))