322 lines
14 KiB
Python
322 lines
14 KiB
Python
#
|
|
# SPDX-FileCopyrightText: Copyright (c) 1993-2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
#
|
|
|
|
|
|
"""tests of QuantPooling module.
|
|
Most tests check the functionality of all the combinations in Quant Pooling against the corresponding functionalities
|
|
in tensor_quant.
|
|
"""
|
|
import pytest
|
|
import numpy as np
|
|
|
|
import torch
|
|
import torch.nn.functional as F
|
|
|
|
from pytorch_quantization import tensor_quant
|
|
from pytorch_quantization.nn.modules import quant_pooling
|
|
|
|
# make everything run on the GPU
|
|
torch.set_default_tensor_type('torch.cuda.FloatTensor')
|
|
|
|
np.random.seed(1234)
|
|
torch.manual_seed(1234)
|
|
|
|
# pylint:disable=missing-docstring, no-self-use
|
|
class TestQuantMaxPool1d():
|
|
|
|
def test_raise(self):
|
|
with pytest.raises(ValueError) as excinfo:
|
|
quant_pooling_object = quant_pooling.QuantMaxPool1d(kernel_size=3, stride=1,
|
|
quant_desc_input=
|
|
tensor_quant.QuantDescriptor(fake_quant=False))
|
|
assert "Only fake quantization is supported" in str(excinfo.value)
|
|
|
|
# Quantizing activations
|
|
def test_input_fake_quant(self):
|
|
quant_pooling_object = quant_pooling.QuantMaxPool1d(kernel_size=3, stride=1)
|
|
|
|
test_input = torch.randn(1, 5, 5, dtype=torch.double)
|
|
|
|
quant_input = tensor_quant.fake_tensor_quant(test_input, torch.max(torch.abs(test_input)))
|
|
|
|
out1 = F.max_pool1d(quant_input, 3, 1, 0, 1, False, False)
|
|
out2 = quant_pooling_object(test_input)
|
|
np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
|
|
|
|
class TestQuantMaxPool2d():
|
|
|
|
def test_raise(self):
|
|
with pytest.raises(ValueError) as excinfo:
|
|
quant_pooling_object = quant_pooling.QuantMaxPool2d(kernel_size=3, stride=1,
|
|
quant_desc_input=
|
|
tensor_quant.QuantDescriptor(fake_quant=False))
|
|
assert "Only fake quantization is supported" in str(excinfo.value)
|
|
|
|
# Quantizing activations
|
|
def test_input_fake_quant(self):
|
|
quant_pooling_object = quant_pooling.QuantMaxPool2d(kernel_size=3, stride=1)
|
|
|
|
test_input = torch.randn(1, 5, 5, 5, dtype=torch.double)
|
|
|
|
quant_input = tensor_quant.fake_tensor_quant(test_input, torch.max(torch.abs(test_input)))
|
|
|
|
out1 = F.max_pool2d(quant_input, 3, 1, 0, 1, False, False)
|
|
out2 = quant_pooling_object(test_input)
|
|
np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
|
|
|
|
def test_input_variable_bits(self):
|
|
# Repeat checking the output for variable number of bits to QuantDescriptor
|
|
for bits in [2, 4, 6]:
|
|
quant_desc_input = tensor_quant.QuantDescriptor(num_bits=bits)
|
|
|
|
quant_pooling.QuantMaxPool2d.set_default_quant_desc_input(quant_desc_input)
|
|
quant_pooling_object = quant_pooling.QuantMaxPool2d(kernel_size=3, stride=1)
|
|
|
|
test_input = torch.randn(1, 5, 5, 5, dtype=torch.double)
|
|
|
|
quant_input = tensor_quant.fake_tensor_quant(test_input, torch.max(torch.abs(test_input)), bits)
|
|
|
|
out1 = F.max_pool2d(quant_input, 3, 1, 0, 1, False, False)
|
|
out2 = quant_pooling_object(test_input)
|
|
np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
|
|
|
|
def test_input_fake_quant_disable(self):
|
|
quant_pooling_object = quant_pooling.QuantMaxPool2d(kernel_size=3, stride=1)
|
|
|
|
test_input = torch.randn(1, 5, 5, 5, dtype=torch.double)
|
|
|
|
quant_pooling_object.input_quantizer.disable()
|
|
|
|
out1 = F.max_pool2d(test_input, 3, 1, 0, 1, False, False)
|
|
out2 = quant_pooling_object(test_input)
|
|
np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
|
|
|
|
def test_input_multi_axis(self):
|
|
quant_desc_input = tensor_quant.QuantDescriptor(num_bits=8, axis=(0, 1))
|
|
|
|
quant_pooling.QuantMaxPool2d.set_default_quant_desc_input(quant_desc_input)
|
|
quant_pooling_object = quant_pooling.QuantMaxPool2d(kernel_size=3, stride=1)
|
|
|
|
test_input = torch.randn(16, 7, 5, 5, dtype=torch.double)
|
|
input_amax = torch.amax(torch.abs(test_input), dim=(2, 3), keepdim=True)
|
|
quant_input = tensor_quant.fake_tensor_quant(test_input, input_amax)
|
|
|
|
out1 = F.max_pool2d(quant_input, 3, 1, 0, 1, False, False)
|
|
out2 = quant_pooling_object(test_input)
|
|
np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
|
|
|
|
class TestQuantMaxPool3d():
|
|
|
|
def test_raise(self):
|
|
with pytest.raises(ValueError) as excinfo:
|
|
quant_pooling_object = quant_pooling.QuantMaxPool3d(kernel_size=3, stride=1,
|
|
quant_desc_input=
|
|
tensor_quant.QuantDescriptor(fake_quant=False))
|
|
assert "Only fake quantization is supported" in str(excinfo.value)
|
|
|
|
# Quantizing activations
|
|
def test_input_fake_quant(self):
|
|
quant_pooling_object = quant_pooling.QuantMaxPool3d(kernel_size=3, stride=1)
|
|
|
|
test_input = torch.randn(5, 5, 5, 5, dtype=torch.double)
|
|
|
|
quant_input = tensor_quant.fake_tensor_quant(test_input, torch.max(torch.abs(test_input)))
|
|
|
|
out1 = F.max_pool3d(quant_input, 3, 1, 0, 1, False, False)
|
|
out2 = quant_pooling_object(test_input)
|
|
np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
|
|
|
|
class TestQuantAvgPool1d():
|
|
|
|
def test_raise(self):
|
|
with pytest.raises(ValueError) as excinfo:
|
|
quant_pooling_object = quant_pooling.QuantAvgPool1d(kernel_size=3, stride=1,
|
|
quant_desc_input=
|
|
tensor_quant.QuantDescriptor(fake_quant=False))
|
|
assert "Only fake quantization is supported" in str(excinfo.value)
|
|
|
|
# Quantizing activations
|
|
def test_input_fake_quant(self):
|
|
quant_pooling_object = quant_pooling.QuantAvgPool1d(kernel_size=3, stride=1)
|
|
|
|
test_input = torch.randn(1, 5, 5, dtype=torch.double)
|
|
|
|
quant_input = tensor_quant.fake_tensor_quant(test_input, torch.max(torch.abs(test_input)))
|
|
|
|
out1 = F.avg_pool1d(quant_input, 3, 1, 0, False, True)
|
|
out2 = quant_pooling_object(test_input)
|
|
np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
|
|
|
|
class TestQuantAvgPool2d():
|
|
|
|
def test_raise(self):
|
|
with pytest.raises(ValueError) as excinfo:
|
|
quant_pooling_object = quant_pooling.QuantAvgPool2d(kernel_size=3, stride=1,
|
|
quant_desc_input=
|
|
tensor_quant.QuantDescriptor(fake_quant=False))
|
|
assert "Only fake quantization is supported" in str(excinfo.value)
|
|
|
|
# Quantizing activations
|
|
def test_input_fake_quant(self):
|
|
quant_pooling_object = quant_pooling.QuantAvgPool2d(kernel_size=3, stride=1)
|
|
|
|
test_input = torch.randn(1, 5, 5, 5, dtype=torch.double)
|
|
|
|
quant_input = tensor_quant.fake_tensor_quant(test_input, torch.max(torch.abs(test_input)))
|
|
|
|
out1 = F.avg_pool2d(quant_input, 3, 1, 0, False, True, None)
|
|
out2 = quant_pooling_object(test_input)
|
|
np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
|
|
|
|
def test_input_variable_bits(self):
|
|
# Repeat checking the output for variable number of bits to QuantDescriptor
|
|
for bits in [2, 4, 6]:
|
|
quant_desc_input = tensor_quant.QuantDescriptor(num_bits=bits)
|
|
|
|
quant_pooling.QuantAvgPool2d.set_default_quant_desc_input(quant_desc_input)
|
|
quant_pooling_object = quant_pooling.QuantAvgPool2d(kernel_size=3, stride=1)
|
|
|
|
test_input = torch.randn(1, 5, 5, 5, dtype=torch.double)
|
|
|
|
quant_input = tensor_quant.fake_tensor_quant(test_input, torch.max(torch.abs(test_input)), bits)
|
|
|
|
out1 = F.avg_pool2d(quant_input, 3, 1, 0, False, True, None)
|
|
out2 = quant_pooling_object(test_input)
|
|
np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
|
|
|
|
def test_input_fake_quant_disable(self):
|
|
quant_pooling_object = quant_pooling.QuantAvgPool2d(kernel_size=3, stride=1)
|
|
|
|
test_input = torch.randn(1, 5, 5, 5, dtype=torch.double)
|
|
|
|
quant_pooling_object.input_quantizer.disable()
|
|
|
|
out1 = F.avg_pool2d(test_input, 3, 1, 0, False, True, None)
|
|
out2 = quant_pooling_object(test_input)
|
|
np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
|
|
|
|
class TestQuantAvgPool3d():
|
|
|
|
def test_raise(self):
|
|
with pytest.raises(ValueError) as excinfo:
|
|
quant_pooling_object = quant_pooling.QuantAvgPool3d(kernel_size=3, stride=1,
|
|
quant_desc_input=
|
|
tensor_quant.QuantDescriptor(fake_quant=False))
|
|
assert "Only fake quantization is supported" in str(excinfo.value)
|
|
|
|
# Quantizing activations
|
|
def test_input_fake_quant(self):
|
|
quant_pooling_object = quant_pooling.QuantAvgPool3d(kernel_size=3, stride=1)
|
|
|
|
test_input = torch.randn(5, 5, 5, 5, dtype=torch.double)
|
|
|
|
quant_input = tensor_quant.fake_tensor_quant(test_input, torch.max(torch.abs(test_input)))
|
|
|
|
out1 = F.avg_pool3d(quant_input, 3, 1, 0, False, True, None)
|
|
out2 = quant_pooling_object(test_input)
|
|
np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
|
|
|
|
class TestQuantAdaptiveAvgPool1d():
|
|
|
|
def test_raise(self):
|
|
with pytest.raises(ValueError) as excinfo:
|
|
quant_pooling_object = quant_pooling.QuantAdaptiveAvgPool1d(output_size=3,
|
|
quant_desc_input=
|
|
tensor_quant.QuantDescriptor(fake_quant=False))
|
|
assert "Only fake quantization is supported" in str(excinfo.value)
|
|
|
|
# Quantizing activations
|
|
def test_input_fake_quant(self):
|
|
quant_pooling_object = quant_pooling.QuantAdaptiveAvgPool1d(output_size=3)
|
|
|
|
test_input = torch.randn(1, 5, 5, dtype=torch.double)
|
|
|
|
quant_input = tensor_quant.fake_tensor_quant(test_input, torch.max(torch.abs(test_input)))
|
|
|
|
out1 = F.adaptive_avg_pool1d(quant_input, 3)
|
|
out2 = quant_pooling_object(test_input)
|
|
np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
|
|
|
|
class TestQuantAdaptiveAvgPool2d():
|
|
|
|
def test_raise(self):
|
|
with pytest.raises(ValueError) as excinfo:
|
|
quant_pooling_object = quant_pooling.QuantAdaptiveAvgPool2d(output_size=3,
|
|
quant_desc_input=
|
|
tensor_quant.QuantDescriptor(fake_quant=False))
|
|
assert "Only fake quantization is supported" in str(excinfo.value)
|
|
|
|
# Quantizing activations
|
|
def test_input_fake_quant(self):
|
|
quant_pooling_object = quant_pooling.QuantAdaptiveAvgPool2d(output_size=3)
|
|
|
|
test_input = torch.randn(1, 5, 5, 5, dtype=torch.double)
|
|
|
|
quant_input = tensor_quant.fake_tensor_quant(test_input, torch.max(torch.abs(test_input)))
|
|
|
|
out1 = F.adaptive_avg_pool2d(quant_input, 3)
|
|
out2 = quant_pooling_object(test_input)
|
|
np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
|
|
|
|
def test_input_variable_bits(self):
|
|
# Repeat checking the output for variable number of bits to QuantDescriptor
|
|
for bits in [2, 4, 6]:
|
|
quant_desc_input = tensor_quant.QuantDescriptor(num_bits=bits)
|
|
|
|
quant_pooling.QuantAdaptiveAvgPool2d.set_default_quant_desc_input(quant_desc_input)
|
|
quant_pooling_object = quant_pooling.QuantAdaptiveAvgPool2d(output_size=3)
|
|
|
|
test_input = torch.randn(1, 5, 5, 5, dtype=torch.double)
|
|
|
|
quant_input = tensor_quant.fake_tensor_quant(test_input, torch.max(torch.abs(test_input)), bits)
|
|
|
|
out1 = F.adaptive_avg_pool2d(quant_input, 3)
|
|
out2 = quant_pooling_object(test_input)
|
|
np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
|
|
|
|
def test_input_fake_quant_disable(self):
|
|
quant_pooling_object = quant_pooling.QuantAdaptiveAvgPool2d(output_size=3)
|
|
|
|
test_input = torch.randn(1, 5, 5, 5, dtype=torch.double)
|
|
|
|
quant_pooling_object.input_quantizer.disable()
|
|
|
|
out1 = F.adaptive_avg_pool2d(test_input, 3)
|
|
out2 = quant_pooling_object(test_input)
|
|
np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
|
|
|
|
class TestQuantAdaptiveAvgPool3d():
|
|
|
|
def test_raise(self):
|
|
with pytest.raises(ValueError) as excinfo:
|
|
quant_pooling_object = quant_pooling.QuantAdaptiveAvgPool3d(output_size=3,
|
|
quant_desc_input=
|
|
tensor_quant.QuantDescriptor(fake_quant=False))
|
|
assert "Only fake quantization is supported" in str(excinfo.value)
|
|
|
|
# Quantizing activations
|
|
def test_input_fake_quant(self):
|
|
quant_pooling_object = quant_pooling.QuantAdaptiveAvgPool3d(output_size=3)
|
|
|
|
test_input = torch.randn(5, 5, 5, 5, dtype=torch.double)
|
|
|
|
quant_input = tensor_quant.fake_tensor_quant(test_input, torch.max(torch.abs(test_input)))
|
|
|
|
out1 = F.adaptive_avg_pool3d(quant_input, 3)
|
|
out2 = quant_pooling_object(test_input)
|
|
np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
|