mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-06 12:20:52 +01:00
f742ceaa46
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/27093 Add `insert_prepack_unpack` and `fold_prepack` to `convert_script` Test Plan: . Imported from OSS Differential Revision: D17678262 fbshipit-source-id: 4bfd6681af6fce226cc77aed8dd84066cbd8ed17
76 lines
2.7 KiB
Python
76 lines
2.7 KiB
Python
from __future__ import absolute_import, division, print_function, unicode_literals
|
|
|
|
import torch
|
|
from .QConfig import QConfig
|
|
|
|
class PackedParams(torch.nn.Module):
|
|
def __init__(self):
|
|
super(PackedParams, self).__init__()
|
|
w = torch.rand((5, 5), dtype=torch.float)
|
|
wq = torch.quantize_per_tensor(w, 2.0, 0, torch.qint8)
|
|
self.set_weight_bias(wq, torch.rand(5))
|
|
|
|
@torch.jit.export
|
|
def set_weight_bias(self, weight, bias):
|
|
# type: (torch.Tensor, Optional[torch.Tensor]) -> None
|
|
self._packed_params = torch.ops.quantized.linear_prepack(weight, bias)
|
|
|
|
@torch.jit.export
|
|
def _weight_bias(self):
|
|
return torch.ops.quantized.linear_unpack(self._packed_params)
|
|
|
|
def forward(self, x):
|
|
return x
|
|
|
|
@torch.jit.export
|
|
def __getstate__(self):
|
|
return self._weight_bias()
|
|
|
|
@torch.jit.export
|
|
def __setstate__(self, state):
|
|
self.set_weight_bias(state[0], state[1])
|
|
|
|
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()
|
|
torch._C._jit_pass_insert_observers(model._c,
|
|
'forward',
|
|
qconfig_dict,
|
|
True)
|
|
return model
|
|
|
|
def convert_script(model, inplace=False):
|
|
_check_is_script_module(model)
|
|
if not inplace:
|
|
model = model.copy()
|
|
torch._C._jit_pass_insert_quant_dequant(model._c, 'forward', True)
|
|
if 'fbgemm' in torch.backends.quantized.supported_engines:
|
|
_packed_params_scripted = torch.jit.script(PackedParams())._c
|
|
torch._C._jit_pass_insert_prepack_unpack(model._c)
|
|
torch._C._jit_pass_fold_prepack(model._c, _packed_params_scripted)
|
|
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')
|
|
if not inplace:
|
|
model = model.copy()
|
|
scripted_qconfig_dict = {k: script_qconfig(v) for k, v in qconfig_dict.items()}
|
|
torch._C._jit_pass_fold_convbn(model._c)
|
|
prepare_script(model, scripted_qconfig_dict, True)
|
|
run_fn(model._c._get_method('forward'), *run_args)
|
|
convert_script(model, True)
|
|
return model
|