7.3.5. FX Based Quantization

FX is a mechanism for torch to process computational graphs. Due to the existence of computational graphs, FX based quantization has the following advantages over eagle mode :

• Automatic matching of fuse pattern can be realized, and users no longer need to manually execute the fuse process.

• It provides possibility for other optimization based on calculation chart.

On the whole, the process of FX quantification is highly consistent with the eagle mode . Generally, you only need to replace the interface:

• prepare_calibration -> prepare_calibration_fx

• prepare_qat -> prepare_qat_fx

• convert -> convert_fx

Among them, prepare_calibration_fx and prepare_qat_fx interfaces are integrated with automatic fuse processes. In addition, we also provide a separate fuse_fx interface for debugging or research.

7.3.5.1. Restriction

FX uses symbolic execution to record operations in the model. This method has many limitations. See the relevant sections of the official document for details limitations-of-symbolic-tracing.

If operations not supported by FX, such as control flow, are used in the user model, they can be wrapped as a function or method as a whole by wrapping. FX will no longer pay attention to their internal logic, but will retain their calls as they are.

We have extended torch.fx.wrap to support more packaging forms. For details, see the interface document of fx.fx_helper.wrap .

The following example illustrates the use of wrap :

from torch import nn
import torch
from torch.nn import functional as F
from horizon_plugin_pytorch.quantization import QuantStub
from horizon_plugin_pytorch.quantization.quantize_fx import QuantizationTracer
from torch.quantization import DeQuantStub

from horizon_plugin_pytorch.quantization.fx.graph_module import GraphModuleWithAttr



class FxWrapExampleNet(nn.Module):
    def __init__(self):
        super(FxWrapExampleNet, self).__init__()
        self.quant = QuantStub()
        self.conv = nn.Conv2d(3, 3, 1)
        self.bn = nn.BatchNorm2d(3)
        self.relu = nn.ReLU()
        self.dequant = DeQuantStub()

    def forward(self, input):
        # The main body of the model needs to be fused, quantified, etc
        x = self.quant(input)
        x = self.conv(x)
        x = self.bn(x)
        x = self.relu(x)
        x = self.dequant(x)

        # Post processing, without quantification, and including conditional branches
        if self.training:
            print("Run softmax")
            return F.softmax(x, dim=1)
        else:
            print("Run argmax")
            return torch.argmax(x, dim=1)

model = FxWrapExampleNet()
tracer = QuantizationTracer([], [])
graph = tracer.trace(model)
graph.print_tabular()

graph_model = GraphModuleWithAttr(model, graph)
print(graph_model.code)

data = torch.rand(1, 3, 64, 64)
ret = graph_model(data)

Run softmax
opcode         name     target                                args        kwargs
-------------  -------  ------------------------------------  ----------  -------------------------------------------
placeholder    input_1  input                                 ()          {}
call_module    conv     conv                                  (input_1,)  {}
call_module    bn       bn                                    (conv,)     {}
call_module    relu     relu                                  (bn,)       {}
call_function  softmax  <function softmax at 0x7f0ac35be040>  (relu,)     {'dim': 1, '_stacklevel': 3, 'dtype': None}
output         output   output                                (softmax,)  {}



def forward(self, input):
    input_1 = input
    conv = self.conv(input_1);  input_1 = None
    bn = self.bn(conv);  conv = None
    relu = self.relu(bn);  bn = None
    softmax = torch.nn.functional.softmax(relu, dim = 1, _stacklevel = 3, dtype = None);  relu = None
    return softmax

As you can see, the post-processing torch.argmax and print statements in the model after trace have been discarded. If you need to keep the post-processing as it is, you can wrap it.

from horizon_plugin_pytorch.fx.fx_helper import wrap as fx_wrap

class FxWrapExampleNet(nn.Module):
    def __init__(self):
        super(FxWrapExampleNet, self).__init__()
        self.quant = QuantStub()
        self.conv = nn.Conv2d(3, 3, 1)
        self.bn = nn.BatchNorm2d(3)
        self.relu = nn.ReLU()
        self.dequant = DeQuantStub()

    # Wrap post-processing as a method
    @fx_wrap
    def _post_process(self, model_output):
        if self.training:
            print("Run softmax")
            return F.softmax(model_output, dim=1)
        else:
            print("Run argmax")
            return torch.argmax(model_output, dim=1)

    def forward(self, input):
        x = self.quant(input)
        x = self.conv(x)
        x = self.bn(x)
        x = self.relu(x)
        x = self.dequant(x)

        return self._post_process(x)

model = FxWrapExampleNet()
tracer = QuantizationTracer([], [])
graph = tracer.trace(model)
graph.print_tabular()

graph_model = GraphModuleWithAttr(model, graph)
print(graph_model.code)

ret = graph_model(data)
ret = graph_model.eval()(data)

opcode       name           target         args              kwargs
-----------  -------------  -------------  ----------------  --------
placeholder  input_1        input          ()                {}
call_module  conv           conv           (input_1,)        {}
call_module  bn             bn             (conv,)           {}
call_module  relu           relu           (bn,)             {}
get_attr     _self          _self          ()                {}
call_method  _post_process  _post_process  (_self, relu)     {}
output       output         output         (_post_process,)  {}



def forward(self, input):
    input_1 = input
    conv = self.conv(input_1);  input_1 = None
    bn = self.bn(conv);  conv = None
    relu = self.relu(bn);  bn = None
    _self = self._self
    _post_process = _self._post_process(relu);  _self = relu = None
    return _post_process

Run softmax
Run argmax

It can be seen that the packaged post-processing is called as a whole in the model after trace , the internal logic will be retained as it is, and the print statement can also print the content normally.