Neural network structure tracking method based on torch FX
By embedding tensor into proxy as an attribute in the Torch FX module, the problem of its inability to process judgment statements is solved, and complete neural network structure tracking and topology map generation are achieved.
Patent Information
- Application Number
- CN202410352716.1
- Authority / Receiving Office
- CN · China
- Patent Type
- Applications(China)
- Current Assignee / Owner
- Filing Date
- 2024-03-26
- Publication Date
- 2025-09-26
AI Technical Summary
The existing Torch FX module does not support judgment statements when tracing neural network structures, resulting in tracing failure and error reporting.
Combine tensor with proxy, embed tensor into proxy as an attribute, call operator function of proxy with that of tensor, and put the returned tensor into the proxy to be returned, so that proxy itself has numerical concept and solves the specific response when judging statements.
It is possible to process judgment statements during Torch FX module tracing, fully trace the neural network structure, and generate the correct topology structure.
Smart Images
Figure CN120706477A_ABST
Abstract
Description
Technical Field
[0001] The present invention belongs to the technical field of neural network model structure optimization, and in particular relates to a neural network structure tracking method based on torch FX. Background Art
[0002] The quantization of neural networks generally uses a specific structure as the quantization unit, such as convolution-batch normalization-activation three operators as a quantization unit. In order to match this specific structure, it is necessary to analyze the entire network structure. Torch's FX is a module that tracks the network model on the Python side. It can track the structure of the entire neural network, thereby replacing or rewriting operators. However, it does not support model definitions with judgment statements. The FX is a toolkit for developers to use to convert instances. FX consists of three main components: a symbol tracker, an intermediate representation, and Python code generation; among them:
[0003] The symbol tracker implements Python's "symbolic execution" code. It provides fake values, called proxies, through the code. Operations are recorded on these proxies.
[0004] The intermediate representation is a container of operations recorded during symbolic tracing. It consists of a list representing the function inputs, the call site (function, method, or instance), and the return value.
[0005] Python code generation makes FX a Python-to-Python (or module-to-module) conversion toolkit. For each Graph IR, valid Python code that matches the Graph semantics can be created. This functionality is encapsulated in GraphModule, which is a torch.nn.Module instance that contains a Graph and a forward method generated from the Graph.
[0006] In summary, this pipeline of components (symbolic tracing -> intermediate representation -> conversion -> Python code generation) constitutes the FX Python-to-Python conversion pipeline. Furthermore, these components can be used independently. For example, symbolic tracing can be used alone to capture the form of code for analysis (rather than conversion). Code generation can be used to programmatically generate models, such as from a configuration file.
[0007] However, judgment statements are not supported when tracing network structures.
[0008] Specifically, first of all, the Torch FX principle is as follows:
[0009] PyTorch is a module in the Python environment and a neural network training framework. Most neural networks use Torch as their training environment. Users use Python statements on their computers to build the network's computational graph according to the Torch definition format and train it. If you want to optimize the network structure or quantize the network operators, you need to analyze the network structure in detail. The FX module automatically tracks Python code and converts the neural network into a topological graph structure that is easy for users to modify. Figure 1 As shown, the basic usage of fx is shown:
[0010] The first step is to import the torch module;
[0011] The second step is to define a network of type MyModule;
[0012] The third step is to instantiate the network;
[0013] Step 4: Import the torch.fx module.
[0014] Step 5: Call the symbolic_trace interface to analyze the trace network;
[0015] Step 6: Print out the tracked model structure.
[0016] The principle of Torch's FX module is to use a proxy to replace the tensor for network inference, rewriting all the computation interfaces that the original tensor passes through.
[0017] Generally, there are the following types of interfaces: Currently, the pytorch framework only has these interfaces. This method only supports the interfaces in pytorch. These interfaces are the operators that make up the neural network.
[0018] Module: torch natively defined module class, such as Conv2d, Linear;
[0019] Function: general function methods, addition, subtraction, multiplication, division, etc.;
[0020] Torch_function: Torch exposes tensor-specific computations, such as torch.cat and torch.add. The input and output of all these interfaces are designed to be used on tensors and then perform specific functions. FX hijacks these interfaces, preventing them from performing specific functions and instead generating operator information and saving it as a graph.
[0021] For example:
[0022] There is a module of type conv2d, which can be defined as conv = torch.nn.Conv2d(). Its specific function is to perform a two-dimensional convolution operation on the input tensor.
[0023] Call conv using Python statements:
[0024] out = conv(x), where x is the input and out is the output. When this statement is executed, the __call__() member function of conv is automatically called to perform a two-dimensional convolution on the input x and return the result out.
[0025] The above is a very common way to use Torch. The call statements here are written in a Python script. If you want to quantize the conv function, you can't modify the file definition. Therefore, FX records the operator operations and generates a graph, and then modifies the operator based on the graph.
[0026] like Figure 2 As shown, when FX tracks the behavior of the tensor, it replaces the tensor with a proxy and sends it to the network for execution. At the same time, it hijacks the __call__ function of conv; but this is not limited to the conv operator. Operators of other module types are all like this (defined in pytorch), and the __call__ function is uniformly hijacked. The call function is equivalent to a transit function. When the left and right modules are executed, the call function is called first, and then the forward function itself is called inside the call function. In this way, when executing the same statement out = conv(x), the input x type is proxy, and the function executed is the hijacked function module_call_wrapper. The function of the hijacking function is to record the input and output as well as the module type, generate node nodes, store the nodes in the graph, and establish topological connections. The final output is also the out of the proxy, which is used as input for the next layer.
[0027] Similarly, by hijacking all interfaces that the tensor passes through and running the network completely again, you can get the topology of the entire neural network.
[0028] Not only the module interface, but also the general operation interface must be hijacked:
[0029] Assume that there is a piece of code as follows:
[0030] x=x*3+y;
[0031] It can be decomposed into two operators:
[0032] x=x*3;
[0033] x=x+y;
[0034] When both statements are executed, the operator member function of x is called.
[0035] Originally, x and y were both tensors, and tensors contained the operator functions __add__ and __mul__. These functions add or multiply two input values. However, tensors do not have the ability to record information.
[0036] Now replace tensor with proxy and rewrite __add__ and __mul__ in proxy. The rewritten functions no longer perform operations but only record operator information. The __add__ function now generates a node whose inputs are x and 3 and whose operation is addition. Similarly, __mul__ generates a node whose inputs are x and y and whose operation is multiplication. The generated node is stored in the graph and the new proxy is returned.
[0037] However, the above native FX has a disadvantage: using a proxy instead of a tensor can cause errors in some operations that require specific values.
[0038] for example:
[0039] x=x+3
[0040] x=x if x>0else 0
[0041] The second statement in this case is a check statement, which determines whether the value of x is greater than 0. If fx is used for tracing again at this point, the __bool__ member function in the proxy will be called. However, the proxy has no concept of specific values and cannot return a Boolean value based on the input. Therefore, it is difficult to determine whether the input is greater than 0.
[0042] Additionally, commonly used technical terms include:
[0043] Torch is an open source training framework;
[0044] fx is a functional module of torch;
[0045] Tensor is the data in the neural network, usually a tensor;
[0046] Proxy is an agent defined in fx, which is used to flow in the neural network and replace tensor during tracking;
[0047] A node is a node in the fx module used to record operators, recording inputs, outputs, and specific operations of operators; a graph is a computational graph that contains multiple nodes, and the nodes are connected into a directed acyclic graph based on inputs and outputs. Summary of the Invention
[0048] To address the aforementioned issues, this application aims to address the shortcomings of existing proxies by combining tensors with proxies, embedding the tensor as an attribute within the proxy. The proxy's operator function calls the tensor's operator function and places the returned tensor into the proxy. This allows the proxy to have a sense of numerical values, allowing operations to be performed using the tensor in its attributes when needed. This article introduces a method that adds functionality to the existing module to achieve a certain level of dynamic support.
[0049] Specifically, the present invention provides a neural network structure tracking method based on torch FX, the method comprising:
[0050] Combine tensor with proxy, embed tensor into proxy as an attribute. Here, tensor and proxy are both classes. Since the member variables of a class can also be classes, tensor is used as a member variable of proxy.
[0051] The operator function of Proxy calls the operator function of Tensor, turning the original operation of Tensor into the operation of Proxy, and the operation of Proxy includes the operation of Tensor;
[0052] And put the returned tensor into the proxy that is about to be returned; so that the proxy itself has a numerical concept, and the tensor in the attribute is used for calculation when needed, that is, when a calculation call function is needed, the tensor in the proxy attribute is used and placed in the member variable of the proxy;
[0053] Furthermore, during the tracking phase, the data circulating in the network changes from tensors to proxies, and the input and output of each operator are proxies. The original computing interface is hijacked, and the tensor is taken from the input proxy within the interface to calculate the original output, and then a new proxy is generated based on the original output. During this process, the computing nodes are recorded in the graph in the form of nodes. The running function of the original computing node is wrapped and redirected to the proxy version to achieve the purpose of hijacking. When the network is executed, the corresponding proxy version function will be called.
[0054] The method further comprises the following steps:
[0055] S1. Hijack the original function during the tracing phase. Wrap the original compute node's running function and redirect it to the proxy version to achieve the hijacking purpose. When executing the network, the corresponding proxy version of the function will be called.
[0056] S2. Start tracing, combine the tensor with the proxy, and embed the tensor into the proxy as an attribute. During the tracing phase, the data flowing in the network changes from the tensor to the proxy, and the input and output of each operator are the proxy.
[0057] During the S3.trace process, the generated operation nodes are recorded and recorded in the graph. The new interface logic is to take the tensor from the input proxy to calculate the original output, and then generate a new proxy based on the original output. During this process, the calculation nodes are recorded in the graph in the form of nodes.
[0058] In step S2, assuming that the operation of the original tensor x, y is y=x*x, the operation of the proxy py, px becomes py.y=px.x*px.x.
[0059] In step S3, assuming that the operation statement of the original tensor is y=x+2, when this statement is executed, the operator function __add__ of the tensor will be called;
[0060] The calling rules here are the specifications of the programming language, and all operations are function calls;
[0061] Which function to call is determined by the parameters of the operation. Here, the parameters of this statement are x and 2, and the return value is y. Since x is a tensor, the __add__ function defined in the tensor class will be called; in the graph tracking stage, it needs to become a proxy of the tensor to execute, so the __add__ function in the proxy class will be called;
[0062] When the __add__ function of Proxy is calculated, it will take the tensor in the proxy x attribute, which is the member variable, add 2 to it, and put it into the member variable of proxy y.
[0063] The method is applicable when encountering a judgment statement during FX module tracing, i.e., the graph tracing stage. This method introduces the tensor into the proxy and assigns a specific value to the proxy. When encountering the judgment statement, a specific response can be made based on the current value.
[0064] Therefore, the advantage of this application is that it solves the problem in the prior art that when encountering a judgment statement during FX module tracking, it is often impossible to continue tracking and can only report an error. This method introduces tensor into proxy and assigns a specific value to the proxy, so that specific response can be made according to the current value when encountering a judgment statement. BRIEF DESCRIPTION OF THE DRAWINGS
[0065] The drawings described herein are used to provide a further understanding of the present invention, constitute a part of this application, and do not constitute a limitation of the present invention.
[0066] Figure 1 This is a schematic diagram of the basic usage of FX in the prior art.
[0067] Figure 2 It is a schematic diagram that hijacks the __call__ function of Module and uses a function whose input and output are proxy to replace a function whose input and output are tensor, so as to achieve the purpose of using proxy to replace tensor.
[0068] Figure 3 This is a code diagram that saves the value of tensor as attribute t_value in the proxy constructor.
[0069] Figure 4 This is a code diagram of the network definition in this application example.
[0070] Figure 5 This is a code diagram showing that hijacking the __call__ function of the Module class in this application example can achieve the purpose of modifying the Module call.
[0071] Figure 6 This is the function that generates the proxy in the example of this application, and generates a schematic diagram of the corresponding node node added to the graph.
[0072] Figure 7 This is a schematic diagram of the process of converting the original simple tensor reasoning into a computational flow graph in this application.
[0073] Figure 8 This is the __abs__ function of the proxy in this application, which hijacks the __abs__ function of the tensor. The content is to infer the abs value of the original tensor and then generate the code diagram of the abs node.
[0074] Figure 9 It is a flow chart of the present application method.
[0075] Figure 10 This is a code diagram for symbolic tracing of PyTorch models using the FX module.
[0076] Figure 11 is used Figure 10 Schematic diagram of the code for outputting the PyTorch model.
[0077] Figure 12 This is a schematic diagram of the steps for network quantization using torch's fx module.
[0078] Figure 13 This is a code diagram for network quantization using torch.fx and torch.quantization.
[0079] Figure 14 This is a code diagram of a simple PyTorch model with a conditional branch that depends on the input data. DETAILED DESCRIPTION
[0080] In order to more clearly understand the technical content and advantages of the present invention, the present invention is now further described in detail with reference to the accompanying drawings.
[0081] The present invention relates to structural tracking of a neural network model, which facilitates subsequent optimization of the structure or quantification of the model.
[0082] The method includes: introducing tensor into proxy to solve the dynamic graph tracking problem. The method further includes:
[0083] Introduce tensor into proxy, that is, combine tensor with proxy and embed tensor into proxy as an attribute. Here, tensor and proxy are both classes (the name of an abstract concept in object-oriented programming). Since the member variables of a class can also be classes, tensor is used as a member variable of proxy here.
[0084] The operator function of the proxy calls the operator function of the tensor, so that the proxy itself has a numerical concept, and the original tensor operation becomes the proxy operation, and the proxy operation includes the tensor operation;
[0085] And put the returned tensor into the proxy to be returned. When the function needs to be called for operation, the tensor in the proxy attribute will be used and put into the member variable of the proxy.
[0086] Furthermore, during the tracking phase, the data circulating in the network is changed from tensor to proxy, and the input and output of each operator are both proxies. The original computing interface is hijacked, and the tensor is taken from the input proxy within the interface to calculate the original output, and then a new proxy is generated based on the original output. During this process, the computing nodes are recorded in the graph in the form of nodes. The running function of the original computing node is wrapped and redirected to the proxy version to achieve the purpose of hijacking. When the network is executed, the corresponding proxy version function will be called.
[0087] The method further comprises the following steps, Figure 9 As shown:
[0088] S1. Hijack the original function during the tracing phase. Wrap the original compute node's running function and redirect it to the proxy version to achieve the hijacking purpose. When executing the network, the corresponding proxy version of the function will be called.
[0089] S2. Start tracing, combine the tensor with the proxy, and embed the tensor into the proxy as an attribute. During the tracing phase, the data flowing in the network changes from the tensor to the proxy, and the input and output of each operator are the proxy.
[0090] During the S3.trace process, the generated operation nodes are recorded and recorded in the graph. The new interface logic is to take the tensor from the input proxy to calculate the original output, and then generate a new proxy based on the original output. During this process, the calculation nodes are recorded in the graph in the form of nodes.
[0091] In the step S2,
[0092] Assuming that the original tensor x, y operation is y = x * x, then the operation of proxy py, px becomes py.y = px.x * px.x;
[0093] Hijack the original computing interface, take the tensor from the input proxy in the interface to calculate the original output, and then generate a new proxy based on the original output. In this process, the computing nodes are recorded in the graph in the form of nodes;
[0094] In the step S3,
[0095] Assuming that the original tensor's operation statement is y=x+2, when this statement is run, the tensor's operator function __add__ will be called;
[0096] The calling rules here are the specifications of the programming language, and all operations are function calls;
[0097] Which function to call is determined by the parameters of the operation. Here, the parameters of this statement are x and 2, and the return value is y. Since x is a tensor, the __add__ function defined in the tensor class will be called; data in neural networks are generally tensors, and pytorch tensor data is defined as tensor;
[0098] During the graph tracing phase, a proxy of the tensor is required to execute, so the __add__ function in the proxy class is called. The purpose of this method is to analyze the network structure through graph tracing. During the calculation, the __add__ function of the proxy takes the tensor in the proxy x attribute, which is the member variable, adds 2, and puts it into the member variable of the proxy y.
[0099] Wrap the running function of the original computing node and redirect it to the proxy version to achieve the purpose of hijacking. When executing the network, the corresponding proxy version function will be called.
[0100] The specific principle is as follows: the value of tensor is saved as attribute t_value in the constructor of node node, such as Figure 3 The figure below shows the proxy class's initialization function, which passes in t_value as a proxy attribute. The node and graph here refer to the node and network described earlier. The function GraphAppendingTracer() is used to obtain a tracer. Introducing node and graph makes it easy to access them at any time.
[0101] Before the proxy generates a new node in the hijacked interface, it first takes out the tensor of the input node for specific calculation and then uses the result to generate a new node.
[0102] For example: There is a network definition such as Figure 4 As shown:
[0103] The network is very simple, with only one fully connected layer and one activation layer. Now we need to trace this network, that is, to parse the network code and convert it into a directed acyclic network topology.
[0104] The original input and output of MyModule are tensors. When executing the network's forward function, the fully connected instance's forward function is called first, followed by the tensor's member function relu(). To track the tensor's computational flow, each compute node of the tensor must be hijacked. Hijacking can be understood as a hook: calling the hook function before calling the original function.
[0105] Since MyModule and Linear are both torch modules, hijacking the __call__ function of the Module class can achieve the purpose of modifying the Module call, such as Figure 5 shown.
[0106] First, the module_call_wrapper function uses functools to modify the function and wrap the original Module's __call__ function. At this time, when the Module is executed again, the module_call_wrapper function will be called instead of the __call__ function.
[0107] Because network tracing doesn't necessarily require a specific tensor, a proxy is used for proxy execution. If a tensor is present, it is used as the proxy's attribute t_value for network inference. In the call_module function, the tensor in the input args is first extracted and inference is performed using the original computation function to obtain the correct tensor. The computation then generates a node, records it in the graph, and generates a new proxy output.
[0108] A Tensor is a concrete tensor with specific data and shape. A Proxy is a proxy for a Tensor, without data or shape, and is an abstract representation of a Tensor. A Node is a computational node that a Tensor passes through, which records the operator type and parameters of the computational node. A Graph, composed of nodes, is the computational flow diagram of a neural network.
[0109] In the function that generates the proxy output, you can see the node generation process as follows Figure 6 As shown in the figure, node records the computation node type (kind), code location (target), parameters (args, kwargs), name (name), and description (type_expr). Each generated node is automatically added to the graph, recording the node's input and output relationships.
[0110] Specifically, it is the process of converting the original simple tensor reasoning into a computational flow graph, such as Figure 7 As shown,
[0111] Get tensor from Proxy, expressed as get tensor from proxy;
[0112] In order to obtain the tensor value contained in the output proxy in the hijacking function, the original inference function is first called to obtain the correct tensor value, which is expressed as Linear forward;
[0113] Create a node, expressed as create Node;
[0114] Add a node to the graph, expressed as Add node to graph;
[0115] Create an output proxy, expressed as create output proxy;
[0116] The above waits for module_call_warpper to be called.
[0117] Taking the abs operator as an example, when executing y=abs(x), the __abs__ member function in the proxy will be called, such as Figure 8 As shown in the figure, the input here is x, which is self itself. The function first performs the abs function on the tensor attribute to obtain the actual tensor output. The result t_value is then given to the constructor of the new node.
[0118] Since this method is based on the native PyTorch framework, operators not supported by PyTorch are not applicable to this method. PyTorch supports a very comprehensive set of operators, covering almost all neural network operators.
[0119] Specifically, the Torch FX module is a new toolkit introduced in PyTorch 1.8 for capturing the computational graph of a PyTorch model. The FX module provides a way to manipulate and analyze models without actually performing computations. It can be used for a variety of tasks, such as model optimization, visualization, and quantization.
[0120] The core concept of the FX module is symbolic tracing. Through symbolic tracing, FX can capture the computational graph of a PyTorch model and represent it as an intermediate representation (IR). This IR is a directed acyclic graph (DAG) where nodes represent operations and edges represent data flow. The following are some of the main components and functions of the FX module:
[0121] torch.fx.symbolic_trace: This function is used to perform symbolic tracing on a PyTorch model. It accepts a model instance as input and returns a GraphModule object that contains the computational graph of the model.
[0122] torch.fx.GraphModule: This class represents the model obtained through symbolic tracing. It contains the structure and parameters of the original model, as well as the graph attribute representing the computational graph.
[0123] torch.fx.Graph: This class represents the computational graph of the model. It contains information about nodes and edges. It can be accessed through the graph attribute of GraphModule.
[0124] torch.fx.Node: This class represents a node in a computational graph. Each node corresponds to an operation, such as a tensor operation or a module call. The node has attributes such as op, target, args, and kwargs, which represent the type, target, and parameters of the operation, respectively.
[0125] torch.fx.Interpreter: This class is used to interpret and execute GraphModule. Through Interpreter, you can run GraphModule without compiling and get intermediate results.
[0126] torch.fx.Transformer: This class is used to transform and optimize GraphModule. By inheriting the Transformer class and overriding its methods, you can implement custom graph transformation logic. torch.fx.Proxy: This class is used to represent tensors during symbolic tracking. It allows operations on tensors without actually creating them.
[0127] Here is a simple example showing how to use the FX module to perform symbolic tracing on a PyTorch model. Figure 10 As shown. Output, such as Figure 11 As shown in Figure 2, this example shows how to use the torch.fx.symbolic_trace function to perform symbolic tracing on a custom module and print out the generated computational graph.
[0128] The FX module also provides some advanced features and use cases, including:
[0129] Quantization: FX can be used to quantize PyTorch models. By converting the model to an FX graph, quantization points can be easily determined and quantization operations can be inserted. PyTorch provides the torch.quantization.fx module, which is built on top of FX to perform quantization-related graph transformations.
[0130] Visualization: FX graphs can be easily visualized to help understand the structure and computational flow of the model. You can use the torch.fx.Graph.print_tabular() method to print the graph in tabular form, or use a graph visualization tool such as Graphviz to generate a visual representation of the graph.
[0131] Model optimization: By analyzing and transforming FX graphs, various model optimization techniques can be implemented, such as operator fusion, constant folding, dead code elimination, etc. This can help improve the execution efficiency of the model and reduce memory usage.
[0132] Custom operators: FX allows you to insert custom operators into the graph to support specific hardware or acceleration libraries. By inheriting torch.fx.Interpreter and overriding related methods, you can implement the interpreted execution of custom operators.
[0133] Dynamic tensor shapes: FX supports dynamic tensor shapes, meaning that there is no need to specify a specific tensor shape during symbolic tracking. This is very useful for models that process variable-length sequences or dynamic batch sizes. Cross-language deployment: FX graphs can be exported to other formats, such as ONNX and TorchScript, to support the deployment of PyTorch models in other languages or environments.
[0134] In summary, the torch.fx module provides a powerful and flexible way to capture and manipulate the computational graph of a PyTorch model. It opens up new possibilities for model analysis, optimization, and deployment, allowing PyTorch models to better adapt to different hardware and application scenarios.
[0135] This concludes a detailed introduction to the torch.fx module. The FX module is a crucial toolkit within the PyTorch ecosystem, significantly enhancing PyTorch's capabilities in model optimization and deployment. As PyTorch continues to evolve, the FX module continues to improve and expand, providing users with more functionality and possibilities.
[0136] Using the torch fx module for network quantization can be divided into the following steps, such as Figure 12 Shown:
[0137] Prepare the model: First, you need to prepare a trained PyTorch model. This model should be defined through the torch.nn module and have loaded the trained weights.
[0138] Symbolic tracing: Use the torch.fx.symbolic_trace function to perform symbolic tracing on the model and generate a GraphModule object. This process captures the computational graph of the model but does not actually perform the calculation.
[0139] Determine the quantization configuration: Determine the quantization configuration, such as the quantization bit width, quantization mode (symmetric quantization or asymmetric quantization), quantization granularity (tensor quantization or channel quantization), etc. You can use the torch.quantization.QConfig class to specify the quantization configuration.
[0140] Prepare for quantization: Use the torch.quantization.prepare_fx function to prepare the GraphModule for quantization. This function inserts some auxiliary nodes into the computational graph to collect statistical information about activation values for subsequent calculation of quantization parameters.
[0141] Calibration: Use a portion of the training data to perform a forward pass on the model to collect statistics about the activation values. This process is called calibration. You can use a subclass of torch.quantization.ObserverBase (such as torch.quantization.MinMaxObserver) to observe and record the minimum and maximum activation values.
[0142] Conversion: Use the torch.quantization.convert_fx function to convert the GraphModule into a quantized model. This function calculates the quantization parameters based on the collected statistics and replaces the original floating-point operations with quantized operations.
[0143] Quantized inference: Use the quantized model for inference. The quantized model will use low-precision integer operations, thereby reducing memory bandwidth and computational overhead.
[0144] Here is a sample code, such as Figure 13 As shown, it demonstrates how to use torch.fx and torch.quantization to perform network quantization:
[0145] In this example, a simple convolutional neural network model is first defined. Then, symbolic tracing of the model is performed using torch.fx.symbolic_trace to generate a GraphModule. Next, a quantization configuration qconfig is defined, specifying the quantization method and parameters.
[0146] Use the torch.quantization.prepare_fx function to prepare the GraphModule for quantization and insert the necessary auxiliary nodes. Then, calibrate the model using the calibration data and collect activation statistics. This assumes that the calibration dataset calibration_data has been prepared.
[0147] After calibration is complete, use the torch.quantization.convert_fx function to convert the GraphModule into a quantized model. Finally, you can use the quantized model for inference. The input data will be automatically quantized and the quantized weights and activation values will be used for calculations.
[0148] It's important to note that quantization introduces a certain degree of accuracy loss, so a trade-off between model accuracy and performance is necessary when performing quantization. The accuracy of the quantized model can be improved by adjusting the quantization configuration, adding calibration data, and fine-tuning the quantized model.
[0149] In addition, not all operations support quantization, and some complex operations (such as convolution of dynamic size) may not be directly quantized. In this case, the model can be divided into quantizable and non-quantizable parts, and the quantizable part is quantized, while the non-quantizable part maintains the original floating-point precision.
[0150] The torch.fx and torch.quantization modules provide a flexible set of tools for quantizing PyTorch models. By generating computational graphs through symbolic tracing and performing quantization-related transformations on the graphs, models can be easily quantized and achieve significant performance improvements and memory savings.
[0151] The above is a detailed introduction to network quantization using the torch.fx module. Quantization is an important means of optimizing model performance, especially on resource-constrained devices (such as mobile and embedded devices). Quantization can reduce a model's memory footprint and accelerate inference while maintaining high accuracy. The combination of the torch.fx and torch.quantization modules provides powerful support for quantizing PyTorch models, making the quantization process more automated and user-friendly.
[0152] Torch's fx does not support dynamic control flow, which is the main improvement of this invention.
[0153] Let's take an example to illustrate that the Torch fx module does not support dynamic control flow. Suppose there is a simple PyTorch model that contains a conditional branch that depends on the input data. The following is an example code, such as Figure 14 As shown in this example, a simple model MyModel is defined, which contains three fully connected layers. In the forward propagation function forward, we decide whether to use fc2 or fc3 for subsequent calculations based on whether the sum of the output x of the first fully connected layer is greater than 0.
[0154] When trying to symbolically trace this model using torch.fx.symbolic_trace , I encounter an error:
[0155] Error occurred during symbolic tracing:Proxy object cannot be compared with a value. This is likely because the Tracer attempted to tracethrough a control flow statement (if, for, while etc.) that depends on a symbolic value. This error message tells us that the fx module cannot process control flow statements (such as if, for, while, etc.) that depend on symbolic values. In this example, the conditional branch if x.sum()>0 depends on the input data x, which is a symbolic value (Proxy object) during symbolic tracing and cannot be compared with a specific value.
[0156] This is a typical example of the fx module not supporting dynamic control flow. In this case, we need to make some modifications to the model or use other technologies such as TorchScript to handle dynamic control flow.
[0157] This invention solves the problem of dynamic control flow by placing tensors in a proxy and running them together. When the judgment function is traced, the actual data tensor at that time is used to calculate the sum and select the correct branch path. This solves the problem that the fx module does not support dynamic control flow.
[0158] The foregoing description is merely a preferred embodiment of the present invention and is not intended to limit the present invention. Those skilled in the art will readily appreciate that various modifications and variations of the present invention are possible. Any modifications, equivalent substitutions, or improvements made within the spirit and principles of the present invention are intended to be within the scope of protection of the present invention.
Claims
1. A neural network structure tracking method based on torch FX, characterized in that: The method comprises: Combine tensor with proxy, embed tensor into proxy as an attribute. Here, tensor and proxy are both classes. Since the member variables of a class can also be classes, tensor is used as a member variable of proxy. The operator function of Proxy calls the operator function of Tensor, turning the original operation of Tensor into the operation of Proxy, and the operation of Proxy includes the operation of Tensor; And put the returned tensor into the proxy that is about to be returned; so that the proxy itself has a numerical concept, and the tensor in the attribute is used for calculation when needed, that is, when a calculation call function is needed, the tensor in the proxy attribute is used and placed in the member variable of the proxy; Furthermore, during the tracking phase, the data circulating in the network changes from tensors to proxies, and the input and output of each operator are proxies. The original computing interface is hijacked, and the tensor is taken from the input proxy within the interface to calculate the original output, and then a new proxy is generated based on the original output. During this process, the computing nodes are recorded in the graph in the form of nodes. The running function of the original computing node is wrapped and redirected to the proxy version to achieve the purpose of hijacking. When the network is executed, the corresponding proxy version function will be called.
2. The method for tracking a neural network structure based on torch FX according to claim 1, wherein: The method further comprises the following steps: S1. Hijack the original function during the tracing phase. Wrap the original compute node's running function and redirect it to the proxy version to achieve the hijacking purpose. When executing the network, the corresponding proxy version of the function will be called. S2. Start tracing, combine the tensor with the proxy, and embed the tensor into the proxy as an attribute. During the tracing phase, the data flowing in the network changes from the tensor to the proxy, and the input and output of each operator are the proxy. During the S3.trace process, the generated operation nodes are recorded and recorded in the graph. The new interface logic is to take the tensor from the input proxy to calculate the original output, and then generate a new proxy based on the original output. During this process, the calculation nodes are recorded in the graph in the form of nodes.
3. The neural network structure tracking method based on torch FX according to claim 2, characterized in that: In step S2, assuming that the operation of the original tensor x, y is y=x*x, the operation of the proxy py, px becomes py.y=px.x*px.x.
4. The method for tracking a neural network structure based on torch FX according to claim 2, wherein: In the step S3, Assuming that the original tensor's operation statement is y=x+2, when this statement is run, the tensor's operator function __add__ will be called; The calling rules here are the specifications of the programming language, and all operations are function calls; Which function to call is determined by the parameters of the operation. Here, the parameters of this statement are x and 2, and the return value is y. Since x is a tensor, the __add__ function defined in the tensor class will be called; in the graph tracking stage, it needs to become a proxy of the tensor to execute, so the __add__ function in the proxy class will be called; When the __add__ function of Proxy is calculated, it will take the tensor in the proxy x attribute, which is the member variable, add 2 to it, and put it into the member variable of proxy y.
5. The method for tracking a neural network structure based on torch FX according to claim 2, wherein: The method is applicable when encountering a judgment statement during FX module tracing, i.e., the graph tracing stage. This method introduces the tensor into the proxy and assigns a specific value to the proxy. When encountering the judgment statement, a specific response can be made based on the current value.