mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-21 05:34:18 +08:00
91f10a1de1
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/32380 We'll clone the module first and then fold conv bn and return a new module Test Plan: . Imported from OSS Differential Revision: D19508033 fbshipit-source-id: 328e91a2c9420761c904a7f2b62dab4cfaaa31ac
108 lines
4.0 KiB
Python
108 lines
4.0 KiB
Python
from __future__ import absolute_import, division, print_function, unicode_literals
|
|
|
|
import torch
|
|
from .qconfig import QConfig
|
|
from torch.jit._recursive import wrap_cpp_module
|
|
|
|
class ConvPackedParams(torch.nn.Module):
|
|
def __init__(self):
|
|
super(ConvPackedParams, self).__init__()
|
|
wq = torch._empty_affine_quantized([1, 1, 1, 1], scale=1.0, zero_point=0, dtype=torch.qint8)
|
|
self.stride = [1, 1]
|
|
self.padding = [0, 0]
|
|
self.dilation = [1, 1]
|
|
self.groups = 1
|
|
self.set_weight_bias(wq, None)
|
|
|
|
@torch.jit.export
|
|
def set_conv_params(self, stride, padding, dilation, groups):
|
|
# type: (List[int], List[int], List[int], int) -> None
|
|
self.stride = stride
|
|
self.padding = padding
|
|
self.dilation = dilation
|
|
self.groups = groups
|
|
|
|
@torch.jit.export
|
|
def set_weight_bias(self, weight, bias):
|
|
# type: (torch.Tensor, Optional[torch.Tensor]) -> None
|
|
self._packed_params = torch.ops.quantized.conv2d_prepack(weight, bias, self.stride,
|
|
self.padding, self.dilation, self.groups)
|
|
|
|
@torch.jit.export
|
|
def _weight_bias(self):
|
|
return torch.ops.quantized.conv2d_unpack(self._packed_params)
|
|
|
|
def forward(self, x):
|
|
return x
|
|
|
|
@torch.jit.export
|
|
def __getstate__(self):
|
|
qweight, bias = self._weight_bias()
|
|
return (qweight,
|
|
bias,
|
|
self.stride,
|
|
self.padding,
|
|
self.dilation,
|
|
self.groups,
|
|
self.training)
|
|
|
|
@torch.jit.export
|
|
def __setstate__(self, state):
|
|
self.stride = state[2]
|
|
self.padding = state[3]
|
|
self.dilation = state[4]
|
|
self.groups = state[5]
|
|
self.set_weight_bias(state[0],
|
|
state[1])
|
|
self.training = state[6]
|
|
|
|
linear_packed_params = None
|
|
conv_packed_params = None
|
|
if 'fbgemm' in torch.backends.quantized.supported_engines:
|
|
linear_packed_params = torch.jit.script(torch.nn.quantized.modules.linear.LinearPackedParams())._c
|
|
conv_packed_params = torch.jit.script(ConvPackedParams())._c
|
|
|
|
def _check_is_script_module(model):
|
|
if not isinstance(model, torch.jit.ScriptModule):
|
|
raise ValueError('input must be a script module, got: ' + str(type(model)))
|
|
|
|
def prepare_script(model, qconfig_dict, inplace=False):
|
|
_check_is_script_module(model)
|
|
if not inplace:
|
|
model = model.copy()
|
|
model = wrap_cpp_module(torch._C._jit_pass_insert_observers(model._c,
|
|
'forward',
|
|
qconfig_dict,
|
|
False))
|
|
return model
|
|
|
|
def convert_script(model, inplace=False):
|
|
_check_is_script_module(model)
|
|
if not inplace:
|
|
model = model.copy()
|
|
model = wrap_cpp_module(torch._C._jit_pass_insert_quant_dequant(model._c, 'forward', False))
|
|
if 'fbgemm' in torch.backends.quantized.supported_engines:
|
|
torch._C._jit_pass_insert_prepack_unpack(model._c)
|
|
return model
|
|
|
|
# TODO: non-scriptable QConfig will be supported later
|
|
def script_qconfig(qconfig):
|
|
return QConfig(
|
|
activation=torch.jit.script(qconfig.activation())._c,
|
|
weight=torch.jit.script(qconfig.weight())._c)
|
|
|
|
def quantize_script(model, qconfig_dict, run_fn, run_args, inplace=False):
|
|
_check_is_script_module(model)
|
|
if not model._c._has_method('forward'):
|
|
raise ValueError('input script module does not have forward method')
|
|
assert not inplace, "We don't support inplace right now"
|
|
if not inplace:
|
|
model = model.copy()
|
|
scripted_qconfig_dict = {k: script_qconfig(v) for k, v in qconfig_dict.items()}
|
|
torch._C._jit_pass_dedup_module_uses(model._c)
|
|
model = wrap_cpp_module(torch._C._jit_pass_fold_convbn(model._c))
|
|
model = prepare_script(model, scripted_qconfig_dict, True)
|
|
run_fn(model._c._get_method('forward'), *run_args)
|
|
model = convert_script(model, True)
|
|
return model
|