mirror of
https://github.com/saymrwulf/pytorch.git
synced 2026-05-14 20:57:59 +00:00
d47a9004c8
Summary: Action following https://github.com/pytorch/pytorch/issues/66232 Pull Request resolved: https://github.com/pytorch/pytorch/pull/66829 Reviewed By: albanD Differential Revision: D31928812 Pulled By: janeyx99 fbshipit-source-id: 8116b7f3728df8632278b013007c06ecce583862
95 lines
2.9 KiB
Python
95 lines
2.9 KiB
Python
# Owner(s): ["oncall: mobile"]
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.quantized as nnq
|
|
import torch.utils.bundled_inputs
|
|
from torch.ao.quantization import (
|
|
default_qconfig,
|
|
float_qparams_weight_only_qconfig,
|
|
)
|
|
|
|
# graph mode quantization based on fx
|
|
from torch.ao.quantization.quantize_fx import (
|
|
prepare_fx,
|
|
convert_fx,
|
|
)
|
|
from torch.testing._internal.common_quantization import NodeSpec as ns
|
|
from torch.testing._internal.common_quantization import (
|
|
QuantizationLiteTestCase,
|
|
LinearModelWithSubmodule,
|
|
)
|
|
|
|
|
|
class TestLiteFuseFx(QuantizationLiteTestCase):
|
|
|
|
# Tests from:
|
|
# ./caffe2/test/quantization/fx/test_quantize_fx.py
|
|
|
|
def test_embedding(self):
|
|
class M(torch.nn.Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.emb = torch.nn.Embedding(num_embeddings=10, embedding_dim=12)
|
|
|
|
def forward(self, indices):
|
|
return self.emb(indices)
|
|
|
|
model = M().eval()
|
|
indices = torch.randint(low=0, high=10, size=(20,))
|
|
|
|
quantized_node = ns.call_module(nnq.Embedding)
|
|
configs = [
|
|
(float_qparams_weight_only_qconfig, ns.call_module(nnq.Embedding)),
|
|
(None, ns.call_module(nn.Embedding)),
|
|
(default_qconfig, ns.call_module(nn.Embedding)),
|
|
]
|
|
|
|
for qconfig, node in configs:
|
|
qconfig_dict = {"": qconfig}
|
|
m = prepare_fx(model, qconfig_dict)
|
|
m = convert_fx(m)
|
|
self._compare_script_and_mobile(m, input=indices)
|
|
|
|
def test_conv2d(self):
|
|
class M(torch.nn.Module):
|
|
def __init__(self):
|
|
super(M, self).__init__()
|
|
self.conv1 = nn.Conv2d(1, 1, 1)
|
|
self.conv2 = nn.Conv2d(1, 1, 1)
|
|
|
|
def forward(self, x):
|
|
x = self.conv1(x)
|
|
x = self.conv2(x)
|
|
return x
|
|
|
|
m = M().eval()
|
|
qconfig_dict = {"": default_qconfig, "module_name": [("conv1", None)]}
|
|
m = prepare_fx(m, qconfig_dict)
|
|
data = torch.randn(1, 1, 1, 1)
|
|
m = convert_fx(m)
|
|
# first conv is quantized, second conv is not quantized
|
|
self._compare_script_and_mobile(m, input=data)
|
|
|
|
def test_submodule(self):
|
|
# test quantizing complete module, submodule and linear layer
|
|
configs = [
|
|
{},
|
|
{"module_name": [("subm", None)]},
|
|
{"module_name": [("fc", None)]},
|
|
]
|
|
for config in configs:
|
|
model = LinearModelWithSubmodule().eval()
|
|
qconfig_dict = {
|
|
"": torch.ao.quantization.get_default_qconfig("qnnpack"),
|
|
**config,
|
|
}
|
|
model = prepare_fx(model, qconfig_dict)
|
|
quant = convert_fx(model)
|
|
|
|
x = torch.randn(5, 5)
|
|
self._compare_script_and_mobile(quant, input=x)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
run_tests()
|