Add e2e test for conv+bn (#27348)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/27348
att
Test Plan:
python test/test_quantization.py
Imported from OSS
Differential Revision: D18182920
fbshipit-source-id: 40edc4d85903f979cd4755d6785d2842faa4d566
diff --git a/test/common_quantization.py b/test/common_quantization.py
index 9793628..442e24d 100644
--- a/test/common_quantization.py
+++ b/test/common_quantization.py
@@ -161,6 +161,7 @@
self.assertEqual(scripted_output, ref_output)
# Below are a series of neural net models to use in testing quantization
+# Single layer models
class SingleLayerLinearModel(torch.nn.Module):
def __init__(self):
super(SingleLayerLinearModel, self).__init__()
@@ -200,7 +201,6 @@
x = self.lstm(x)
return x
-
class ConvModel(torch.nn.Module):
def __init__(self):
super(ConvModel, self).__init__()
@@ -224,6 +224,33 @@
x = self.dequant(x)
return x
+class ConvBnModel(torch.nn.Module):
+ def __init__(self):
+ super(ConvBnModel, self).__init__()
+ self.conv = torch.nn.Conv2d(3, 5, 3, bias=False).to(dtype=torch.float)
+ self.bn = torch.nn.BatchNorm2d(5).to(dtype=torch.float)
+
+ def forward(self, x):
+ x = self.conv(x)
+ x = self.bn(x)
+ return x
+
+class AnnotatedConvBnModel(torch.nn.Module):
+ def __init__(self):
+ super(AnnotatedConvBnModel, self).__init__()
+ self.qconfig = default_qconfig
+ self.conv = torch.nn.Conv2d(3, 5, 3, bias=False).to(dtype=torch.float)
+ self.bn = torch.nn.BatchNorm2d(5).to(dtype=torch.float)
+ self.quant = QuantStub()
+ self.dequant = DeQuantStub()
+
+ def forward(self, x):
+ x = self.quant(x)
+ x = self.conv(x)
+ x = self.bn(x)
+ x = self.dequant(x)
+ return x
+
class TwoLayerLinearModel(torch.nn.Module):
def __init__(self):
super(TwoLayerLinearModel, self).__init__()
diff --git a/test/test_quantization.py b/test/test_quantization.py
index ce8225d..0490060 100644
--- a/test/test_quantization.py
+++ b/test/test_quantization.py
@@ -7,7 +7,7 @@
import torch.nn.intrinsic.quantized as nniq
import torch.nn.intrinsic.qat as nniqat
from torch.quantization import \
- QConfig, QConfigDynamic, default_observer, default_weight_observer, get_observer_dict,\
+ QConfigDynamic, get_observer_dict, default_weight_observer, \
quantize, prepare, convert, prepare_qat, quantize_qat, fuse_modules, \
quantize_dynamic, default_qconfig, default_debug_qconfig, default_qat_qconfig, \
default_dynamic_qconfig, HistogramObserver, MinMaxObserver, PerChannelMinMaxObserver,\
@@ -20,6 +20,7 @@
from common_quantization import QuantizationTestCase, \
AnnotatedSingleLayerLinearModel, SingleLayerLinearModel, \
AnnotatedConvModel, ConvModel, \
+ AnnotatedConvBnModel, ConvBnModel, \
SkipQuantModel, QuantStubModel, \
ModelForFusion, ModelWithSequentialFusion, ManualLinearQATModel, ManualConvLinearQATModel, \
ModelWithFunctionals, \
@@ -671,7 +672,7 @@
)
class GraphModePostTrainingQuantTest(QuantizationTestCase):
@_tmp_donotuse_dont_inline_everything
- def test_single_layer(self):
+ def test_single_liner(self):
r"""Compare the result of quantizing single linear layer in
eager mode and graph mode
"""
@@ -686,9 +687,7 @@
self.calib_data)
qconfig_dict = {
- '': QConfig(
- activation=default_observer,
- weight=default_weight_observer)
+ '': default_qconfig
}
model_script = quantize_script(
torch.jit.script(linear_model),
@@ -726,6 +725,33 @@
result_script = model_script(self.img_data[0][0])
self.assertEqual(result_eager, result_script)
+ @unittest.skip("This doesn't work right now, re-enable after fold_convbn is fixed")
+ def test_conv_bn(self):
+ r"""Compare the result of quantizing conv + bn layer in
+ eager mode and graph mode
+ """
+ # eager mode
+ conv_model = AnnotatedConvBnModel().eval()
+ conv_model_to_script = ConvBnModel().eval()
+ # copy the weight from eager mode so that we can
+ # compare the result of the two quantized models later
+ conv_model_to_script.conv.weight = torch.nn.Parameter(conv_model.conv.weight.detach())
+ fuse_modules(conv_model, ['conv', 'bn'], inplace=True)
+ model_eager = quantize(conv_model, default_eval_fn,
+ self.img_data)
+ qconfig_dict = {
+ '': default_qconfig
+ }
+ model_script = quantize_script(
+ torch.jit.script(conv_model_to_script),
+ qconfig_dict,
+ default_eval_fn,
+ [self.img_data],
+ inplace=False)
+ result_eager = model_eager(self.img_data[0][0])
+ result_script = model_script(self.img_data[0][0])
+ self.assertEqual(result_eager, result_script)
+
@unittest.skip("quantization for inlined linear is not working right now")
def test_nested(self):
# Eager mode
