Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

[Torch][Quantized] Fix converting serialized quantized models #5839

Merged
merged 3 commits into from
Jun 18, 2020
Merged

[Torch][Quantized] Fix converting serialized quantized models #5839

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
ae7fdc1
[Torch] Fix converting serialized quantized models
masahi Jun 18, 2020
081d318
clean up dtype check
masahi Jun 18, 2020
8bd86ed
comment clean up
masahi Jun 18, 2020
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
46 changes: 29 additions & 17 deletions python/tvm/relay/frontend/pytorch.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -115,6 +115,11 @@ def inplace_add_to_add(op_name):
return False


def _is_quantized_tensor(data, prelude):
ty = _infer_type_with_prelude(data, prelude)
return ty.dtype == "uint8"


# operator implementation
def _elemwise(name):
def _impl(inputs, input_types):
Expand Down Expand Up @@ -530,10 +535,11 @@ def _impl(inputs, input_types):
return _impl


def _relu():
def _relu(prelude):
def _impl(inputs, input_types):
data = inputs[0]
if input_types[0] == "quint8":
# See the comment in adaptive avg pool2d
if _is_quantized_tensor(data, prelude):
assert len(inputs) == 3, "Input quant param not found in op inputs"
input_zero_point = _expr.const(inputs[2], dtype="int32")
return qnn_torch.quantized_relu(data, input_zero_point)
Expand Down Expand Up @@ -595,15 +601,18 @@ def _impl(inputs, input_types):
return _op.log(_op.tensor.sigmoid(data))
return _impl

def _adaptive_avg_pool_2d():
def _adaptive_avg_pool_2d(prelude):
def _impl(inputs, input_types):
data = inputs[0]
output_size = _infer_shape(inputs[1])

def func(x):
return _op.nn.adaptive_avg_pool2d(x, output_size=output_size)

if input_types[0] == "quint8":
# If a quantized Torch module is saved and loaded back, dtype will be dropped
# input_types[0] can be float even though the input is a quantized tensor
# To reliably determine input types, we use Relay's type inference result
if _is_quantized_tensor(data, prelude):
masahi marked this conversation as resolved.
Show resolved Hide resolved
return qnn_torch.apply_with_upcast(data, func)

return func(data)
Expand Down Expand Up @@ -1108,7 +1117,7 @@ def _impl(inputs, input_types):
return _op.log(_op.exp(inputs[0] * beta) + _expr.const(1.)) / beta
return _impl

def _avg_pool2d():
def _avg_pool2d(prelude):
def _impl(inputs, input_types):
data = inputs[0]

Expand All @@ -1130,7 +1139,8 @@ def func(x):
ceil_mode=ceil_mode,
count_include_pad=count_include_pad)

if input_types[0] == "quint8":
# See the comment in adaptive avg pool2d
if _is_quantized_tensor(data, prelude):
return qnn_torch.apply_with_upcast(data, func)

return func(data)
Expand Down Expand Up @@ -1254,7 +1264,7 @@ def _impl(inputs, input_types):

return _impl

def _mean():
def _mean(prelude):
def _impl(inputs, input_types):
data = inputs[0]

Expand All @@ -1274,7 +1284,8 @@ def _impl(inputs, input_types):
def func(x):
return _op.mean(x, axis, keepdims, exclude)

if input_types[0] == "quint8":
# See the comment in adaptive avg pool2d
if _is_quantized_tensor(data, prelude):
assert len(inputs) == 6, "Input quant param not found in op inputs"
input_scale = _expr.const(inputs[4])
input_zero_point = _expr.const(inputs[5])
Expand Down Expand Up @@ -1492,7 +1503,7 @@ def _impl(inputs, input_types):

return _impl

def _upsample(method):
def _upsample(method, prelude):
def _impl(inputs, input_types):
if isinstance(inputs[1], _expr.Var):
out_size = _infer_shape(inputs[1])
Expand All @@ -1516,7 +1527,8 @@ def _impl(inputs, input_types):
def func(x):
return _op.image.resize(x, out_size, "NCHW", method, coord_trans)

if input_types[0] == "quint8":
# See the comment in adaptive avg pool2d
if _is_quantized_tensor(data, prelude):
import torch
from packaging import version

Expand Down Expand Up @@ -1835,8 +1847,8 @@ def _get_convert_map(prelude):
"aten::take" : _take(),
"aten::where" : _where(),
"aten::topk" : _topk(),
"aten::relu" : _relu(),
"aten::relu_" : _relu(),
"aten::relu" : _relu(prelude),
"aten::relu_" : _relu(prelude),
"aten::prelu" : _prelu(),
"aten::leaky_relu" : _leaky_relu(),
"aten::elu" : _elu(),
Expand All @@ -1845,7 +1857,7 @@ def _get_convert_map(prelude):
"aten::gelu" : _gelu(),
"aten::selu" : _selu(),
"aten::log_sigmoid" : _log_sigmoid(),
"aten::adaptive_avg_pool2d" : _adaptive_avg_pool_2d(),
"aten::adaptive_avg_pool2d" : _adaptive_avg_pool_2d(prelude),
"aten::adaptive_max_pool2d" : _adaptive_max_pool_2d(),
"aten::max_pool2d" : _maxpool_2d(),
"aten::max_pool2d_with_indices" : _maxpool_2d_with_indices(),
Expand Down Expand Up @@ -1874,13 +1886,13 @@ def _get_convert_map(prelude):
"aten::log_softmax" : _log_softmax(),
"aten::sigmoid" : _sigmoid(),
"aten::softplus" : _softplus(),
"aten::avg_pool2d" : _avg_pool2d(),
"aten::avg_pool2d" : _avg_pool2d(prelude),
"aten::avg_pool3d" : _avg_pool3d(),
"aten::dropout" : _dropout(),
"aten::dropout_" : _dropout(),
"aten::feature_dropout" : _dropout(),
"aten::alpha_dropout" : _dropout(),
"aten::mean" : _mean(),
"aten::mean" : _mean(prelude),
"aten::chunk" : _chunk(prelude),
"aten::matmul" : _matmul(prelude),
"aten::expand" : _expand(),
Expand Down Expand Up @@ -1932,8 +1944,8 @@ def _get_convert_map(prelude):
"aten::isnan" : _unary("isnan"),
"aten::clamp" : _clamp(),
"aten::detach" : _identity(),
"aten::upsample_bilinear2d" : _upsample("bilinear"),
"aten::upsample_nearest2d" : _upsample("nearest_neighbor"),
"aten::upsample_bilinear2d" : _upsample("bilinear", prelude),
"aten::upsample_nearest2d" : _upsample("nearest_neighbor", prelude),
"aten::upsample_trilinear3d" : _upsample3d("trilinear"),
"aten::upsample_nearest3d" : _upsample3d("nearest_neighbor"),
"aten::expand_as" : _expand_as(),
Expand Down
45 changes: 42 additions & 3 deletions tests/python/frontend/pytorch/qnn_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -63,7 +63,7 @@ def get_qconfig(per_channel):
weight=default_weight_observer)


def quantize_model(model, inp, per_channel=False, dummy=True):
def quantize_model(model, inp, per_channel=False):
model.fuse_model()
model.qconfig = get_qconfig(per_channel)
torch.quantization.prepare(model, inplace=True)
Expand Down Expand Up @@ -243,6 +243,18 @@ def fuse_model(self):
pass


class AdaptiveAvgPool2d(nn.Module):
def __init__(self):
super().__init__()
self.pool = QuantWrapper(nn.AdaptiveAvgPool2d((1, 1)))

def forward(self, x):
return self.pool(x)

def fuse_model(self):
pass


def test_quantized_modules():
imagenet_ishape = (1, 3, 224, 224)

Expand Down Expand Up @@ -280,7 +292,7 @@ def test_quantized_modules():
raw_module.eval()
inp = torch.rand(ishape)

quantize_model(raw_module, inp, per_channel=per_channel, dummy=True)
quantize_model(raw_module, inp, per_channel=per_channel)
script_module = torch.jit.trace(raw_module, inp).eval()

with torch.no_grad():
Expand Down Expand Up @@ -376,7 +388,7 @@ def get_imagenet_input():
inp = get_imagenet_input()
pt_inp = torch.from_numpy(inp)

quantize_model(raw_model, pt_inp, per_channel=per_channel, dummy=False)
quantize_model(raw_model, pt_inp, per_channel=per_channel)
script_module = torch.jit.trace(raw_model, pt_inp).eval()

with torch.no_grad():
Expand Down Expand Up @@ -465,3 +477,30 @@ def get_imagenet_input():
mean abs_diff: 0.054197952
558 in 1000 raw outputs identical.
"""


def test_serialized_modules():
ishape = (1, 16, 64, 64)
raw_module = AdaptiveAvgPool2d().eval()
inp = torch.rand(ishape)

quantize_model(raw_module, inp)
script_module = torch.jit.trace(raw_module, inp).eval()

fname = "tmp.pt"
torch.jit.save(script_module, fname)
loaded = torch.jit.load(fname)
os.remove(fname)

with torch.no_grad():
pt_result = loaded(inp.clone()).numpy()

input_name = "input"
runtime = get_tvm_runtime(loaded, input_name, ishape)
runtime.set_input(input_name, inp.numpy().copy())
runtime.run()
tvm_result = runtime.get_output(0).asnumpy()

num_identical = np.sum(tvm_result == pt_result)
match_ratio = num_identical / float(np.prod(tvm_result.shape))
assert match_ratio > 0.2