mirror of
https://github.com/saymrwulf/pytorch.git
synced 2026-05-15 21:00:47 +00:00
7ea5fa3dd4
Summary: After https://github.com/pytorch/pytorch/pull/77608 `example_inputs` is required input for `prepare_fx` and `prepare_qat_fx`. This makes quantizing submodules harder, so we added this utility function to get a dictionary from fqn to submodule example_inputs Example Call: ``` example_inputs = (tensor0,) get_fqn_to_example_inputs(m, example_inputs) ``` Example output: ``` { "linear1": (tensor1,), "linear2": (tensor2,), "sub": (tensor3,), "sub.linear1": (tensor4,), ... } ``` Test Plan: python test/test_quantization.py TestUtils Reviewers: Subscribers: Tasks: Tags: Pull Request resolved: https://github.com/pytorch/pytorch/pull/78286 Approved by: https://github.com/dzdang
128 lines
4.6 KiB
Python
128 lines
4.6 KiB
Python
# Owner(s): ["oncall: quantization"]
|
|
|
|
import torch
|
|
from torch.testing._internal.common_utils import TestCase
|
|
from torch.ao.quantization.utils import get_fqn_to_example_inputs
|
|
|
|
|
|
class TestUtils(TestCase):
|
|
def _test_get_fqn_to_example_inputs(self, M, example_inputs, expected_fqn_to_dim):
|
|
m = M().eval()
|
|
fqn_to_example_inputs = get_fqn_to_example_inputs(m, example_inputs)
|
|
for fqn, expected_dims in expected_fqn_to_dim.items():
|
|
assert fqn in expected_fqn_to_dim
|
|
example_inputs = fqn_to_example_inputs[fqn]
|
|
for example_input, expected_dim in zip(example_inputs, expected_dims):
|
|
assert example_input.dim() == expected_dim
|
|
|
|
def test_get_fqn_to_example_inputs_simple(self):
|
|
class Sub(torch.nn.Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.linear1 = torch.nn.Linear(5, 5)
|
|
self.linear2 = torch.nn.Linear(5, 5)
|
|
|
|
def forward(self, x):
|
|
x = self.linear1(x)
|
|
x = self.linear2(x)
|
|
return x
|
|
|
|
class M(torch.nn.Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.linear1 = torch.nn.Linear(5, 5)
|
|
self.linear2 = torch.nn.Linear(5, 5)
|
|
self.sub = Sub()
|
|
|
|
def forward(self, x):
|
|
x = self.linear1(x)
|
|
x = self.linear2(x)
|
|
x = self.sub(x)
|
|
return x
|
|
|
|
expected_fqn_to_dim = {
|
|
"": (2,),
|
|
"linear1": (2,),
|
|
"linear2": (2,),
|
|
"sub": (2,),
|
|
"sub.linear1": (2,),
|
|
"sub.linear2": (2,)
|
|
}
|
|
example_inputs = (torch.rand(1, 5),)
|
|
self._test_get_fqn_to_example_inputs(M, example_inputs, expected_fqn_to_dim)
|
|
|
|
def test_get_fqn_to_example_inputs_default_kwargs(self):
|
|
""" Test that we can get example inputs for functions with default keyword arguments
|
|
"""
|
|
class Sub(torch.nn.Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.linear1 = torch.nn.Linear(5, 5)
|
|
self.linear2 = torch.nn.Linear(5, 5)
|
|
|
|
def forward(self, x, key1=torch.rand(1), key2=torch.rand(1)):
|
|
x = self.linear1(x)
|
|
x = self.linear2(x)
|
|
return x
|
|
|
|
class M(torch.nn.Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.linear1 = torch.nn.Linear(5, 5)
|
|
self.linear2 = torch.nn.Linear(5, 5)
|
|
self.sub = Sub()
|
|
|
|
def forward(self, x):
|
|
x = self.linear1(x)
|
|
x = self.linear2(x)
|
|
# only override `key2`, `key1` will use default
|
|
x = self.sub(x, key2=torch.rand(1, 2))
|
|
return x
|
|
|
|
expected_fqn_to_dim = {
|
|
"": (2,),
|
|
"linear1": (2,),
|
|
"linear2": (2,),
|
|
# second arg is `key1`, which is using default argument
|
|
# third arg is `key2`, override by callsite
|
|
"sub": (2, 1, 2),
|
|
"sub.linear1": (2,),
|
|
"sub.linear2": (2,)
|
|
}
|
|
example_inputs = (torch.rand(1, 5),)
|
|
self._test_get_fqn_to_example_inputs(M, example_inputs, expected_fqn_to_dim)
|
|
|
|
def test_get_fqn_to_example_inputs_complex_args(self):
|
|
""" Test that we can record complex example inputs such as lists and dicts
|
|
"""
|
|
class Sub(torch.nn.Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.linear1 = torch.nn.Linear(5, 5)
|
|
self.linear2 = torch.nn.Linear(5, 5)
|
|
|
|
def forward(self, x, list_arg, dict_arg):
|
|
x = self.linear1(x)
|
|
x = self.linear2(x)
|
|
return x
|
|
|
|
class M(torch.nn.Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.linear1 = torch.nn.Linear(5, 5)
|
|
self.linear2 = torch.nn.Linear(5, 5)
|
|
self.sub = Sub()
|
|
|
|
def forward(self, x):
|
|
x = self.linear1(x)
|
|
x = self.linear2(x)
|
|
x = self.sub(x, [x], {"3": x})
|
|
return x
|
|
|
|
example_inputs = (torch.rand(1, 5),)
|
|
m = M().eval()
|
|
fqn_to_example_inputs = get_fqn_to_example_inputs(m, example_inputs)
|
|
assert "sub" in fqn_to_example_inputs
|
|
assert isinstance(fqn_to_example_inputs["sub"][1], list)
|
|
assert isinstance(fqn_to_example_inputs["sub"][2], dict) and \
|
|
"3" in fqn_to_example_inputs["sub"][2]
|