A video super-resolution inference method and device based on TensorRT

Abstract

The application discloses a video super-resolution inference method and device based on TensorRT. First, a video super-resolution network is trained. Then, network operators are analyzed, common operators are sorted out, unsupported operators are written, and the operator fusion is performed on the optimizable operators. Finally, for the complete video super-resolution network structure, a serialized engine is written and generated by using TensorRT to perform inference. The application proposes a new method for the landing of video super-resolution inference, that is, the inference framework of TensorRT is used for deployment and acceleration, a common operator library is established to align the TensorRT operators and Pytorch operators, the deployment process is optimized, and the inference under the training framework has better speed and lower memory occupation.

Description

TECHNICAL FIELD

[0001] The application belongs to the field of video super-resolution in computer vision, and particularly relates to an optimization scheme for using TensorRT to perform inference on a video super-resolution method, which can reduce the inference time of a model, reduce the memory occupation, and provide an effective solution for the landing of video super-resolution. BACKGROUND

[0002] Video super-resolution is derived from image super-resolution, and the purpose is to recover a high-resolution (HR) image from one or more low-resolution (LR) images. Their differences are also obvious, since a video is composed of multiple frames, that is, the former usually utilizes the information between frames for repair, which also brings additional challenges, that is, the amount of information is larger, a more complex network structure is needed, more model parameters are needed, and more memory is occupied. These challenges bring great difficulties to the landing of the model and the use of actual scenes.

[0003] TensorRT is a high-performance inference framework produced by NVIDIA, which is specially applied to GPU devices and embedded devices for inference. TensorRT includes the deep learning landing technologies we usually see: model quantization, dynamic memory optimization and other optimization technologies. More importantly, the inference code is directly run on the graphics card using CUDA language, and all code libraries only include C++ and CUDA. After we run the trained code using this optimization library, the running time and memory and memory occupation can be greatly reduced.

[0004] There are three ways to build a model supported by TensorRT itself. The first way is to use the onnx format, which can convert the model structure and model parameters of onnx to TensorRT for inference acceleration through the onnx interpreter of TensorRT itself. This way has the advantages of being simple and easy to operate, and the disadvantage is that it is difficult to handle unsupported operators. The second way is to import a caffe or TensorFlow model. TensorRT comes with interpreters for these two models, which can be more efficient than the onnx method, but the disadvantage is that the model needs to be converted into caffe or TensorFlow, and most of the current networks are under the Pytorch framework. The last way is to use C++ native API to build the network, according to the deep learning operator API provided by TensorRT, such as convolution, pooling, activation function, etc. The advantage of this way is that the inference speed is faster, and more optimization can be done inside. The disadvantage is that it is time-consuming and complex to build a model. The application optimizes the deployment based on the third way, so as to optimize the inference deployment of the video super-resolution network in a simpler way with better effect. SUMMARY

[0005] To address the shortcomings of existing technologies, this invention provides a video super-resolution inference method and apparatus based on TensorRT.

[0006] Therefore, the purpose of this invention is to improve the speed of video super-resolution, reduce the GPU memory required by the video super-resolution network, and optimize the usage of TensorRT. This invention proposes a simplified TensorRT deployment scheme for video super-resolution and optimizes a commonly used operator in video super-resolution. The input of this invention is a long video, and the output is the super-resolution long video. The video super-resolution network is trained, and then the trained model and computation graph are exported and saved. Subsequently, the operators in the video super-resolution network are analyzed, commonly used operators are organized, optimizable operators are fused, and finally, the video super-resolution network is transformed and applied layer by layer.

[0007] A video super-resolution inference method based on TensorRT includes the following steps:

[0008] Step (1) Train the video super-resolution network;

[0009] Step (2) Analyze network operators, organize commonly used operators, write unsupported operators, and perform operator fusion for optimizable operators.

[0010] Step (3) For the complete video super-resolution network structure, use TensorRT to write and generate a serialization engine for inference;

[0011] Furthermore, step (1) is implemented as follows;

[0012] The selected video super-resolution network model is trained to obtain the deployed model using the PyTorch deep learning training framework. Additionally, mmediting is used to configure the model parameters and training dataset required for training Basicvsr.

[0013] Furthermore, step (2) is implemented as follows;

[0014] Commonly used operators need to be organized because TensorRT's operator API differs from PyTorch's. Aligning the operators of both facilitates subsequent calls. Operators not supported by TensorRT need to be written manually. For unsupported operators, a custom TensorRT plugin is required, which needs to inherit from the IPluginV2DynamicExt and PluginCreator classes and implement the member functions of these classes, including the CUDA functions required for inference and the serialization function. Optimizable operators should be merged into a single operator to reduce memory accesses during inference and accelerate overall computation.

[0015] Further, step (3) is implemented as follows:

[0016] According to the used video super-resolution network, the corresponding operator is added layer by layer. Then the generation engine is generated, and in the generation process, TensorRT is used for compilation and optimization, which can automatically select a more reasonable and faster operator implementation. After the engine is constructed, the optimal scheme is obtained, and when performing inference, the engine needs to be deserialized first and then the video sequence is input. In the code, the input data is transmitted from the CPU to the GPU, and the inference calculation is performed on the GPU. Finally, the data is transmitted from the GPU to the CPU.

[0017] Further, the video super-resolution network uses Basicvsr.

[0018] Further, the operator flow_warp not supported by TensorRT is in the execution process of Pytorch: first create a grid, add the optical flow to the grid, and then normalize the value to (-1, 1). Then call the grid_sample operator with input parameters as input, processed optical flow data and interpolation mode. The execution process of the grid_sample operator in Pytorch is to first denormalize, then find the value of the output corresponding position index on the optical flow, and use the value to perform linear interpolation on the input to get the result at the output.

[0019] When using the custom operator of TensorRT, cuda code is needed to complete the forward execution process on Pytorch as described above. By analyzing the execution process of the flow_warp operator on Pytorch, it can be known that the normalization and denormalization of the optical flow data are opposite operations, which are omitted. In addition, the grid of the optical flow data can be calculated by each thread on the cuda corresponding to the horizontal and vertical axis indexes and added to the optical flow. The above operator fusion process can reduce matrix transformation and operation. Finally, the linear interpolation operation can obtain the output result.

[0020] A video super-resolution inference device based on TensorRT includes a video super-resolution network training module, a network operator analysis processing module, a serialized engine construction module, and an inference module.

[0021] The video super-resolution network training module is used to train the video super-resolution network to obtain a deployed model, and uses the Pytorch deep learning training framework. In addition, the mmediting configuration is used to configure the model parameters and training data set required for training Basicvsr.

[0022] The network operator analysis processing module is used to analyze the network operator, organize the commonly used operators, write the unsupported operators, and optimize the operators.

[0023] Further, the network operator analysis processing module is specifically implemented as follows:

[0024] The commonly used operators need to be sorted out because the operator API of TensorRT is different from that of Pytorch. After aligning the operators of the two, subsequent calls are facilitated. Operators not supported by TensorRT need to be written. For unsupported operators, the TensorRT plugin needs to be customized, the IPluginV2DynamicExt and PluginCreator classes need to be inherited, and the member functions in the class need to be implemented, including the cuda functions required for inference and serialization functions. For the optimizable operator, fusion into an operator is performed, which reduces the number of memory accesses in the inference process and speeds up the overall calculation.

[0025] The serialization engine construction module adds the corresponding operators layer by layer according to the video super-resolution network used. Then the engine is generated, and in the generation process, TensorRT is used for compilation and optimization to automatically select a more reasonable and faster operator implementation method.

[0026] After the engine is constructed, the optimal solution is obtained, and when the inference module is used for inference, the engine needs to be deserialized first and then the video sequence is input. In the code, the input data is transmitted from the CPU to the GPU end, the inference calculation is performed on the GPU, and finally the data is transmitted from the GPU end to the CPU end.

[0027] Further, the network operator analysis processing module for writing unsupported operators is specifically implemented as follows:

[0028] The flow_warp operator not supported by TensorRT in the execution process of Pytorch is as follows: first, a grid is created, the optical flow is added to the grid, and then the value is normalized to (-1, 1). Then the grid_sample operator is called with the input, the processed optical flow data, and the interpolation mode. The execution process of the grid_sample operator in Pytorch is to first denormalize, then find the value of the output corresponding position index on the optical flow, and use the value to perform linear interpolation on the input to obtain the result at the output.

[0029] In the use of TensorRT custom operator, the cuda code needs to be written to complete the foregoing forward execution process on Pytorch. By analyzing the execution process of the flow_warp operator on Pytorch, it can be known that the normalization and denormalization of the optical flow data are opposite operations, which are omitted, and the grid of the optical flow data can be calculated by each thread on the cuda corresponding to the horizontal and vertical axis indexes and added to the optical flow. The above operator fusion process can reduce matrix transformation and operation. Finally, the linear interpolation operation can obtain the output result.

[0030] The application has the following advantages:

[0031] (1) The application proposes a new method of video super-resolution inference landing, that is, using TensorRT inference framework for deployment and acceleration, establishing a commonly used operator library to align TensorRT operators and Pytorch operators, optimizing the deployment process, and being able to have better speed and lower memory occupation than inference under the training framework.

[0032] (2) The application proposes a more concise implementation of the flow_warp operator on TensorRT. In this layer, the flow_warp can be implemented faster, and operator fusion and cuda acceleration are performed. BRIEF DESCRIPTION OF DRAWINGS

[0033] Figure 1 The specific use process of the TensorRT embodiment of the application is as follows:

[0034] Figure 2 The flow_warp operator calculation graph of the embodiment of the application is as follows. DETAILED DESCRIPTION

[0035] The specific implementation of the application patent will be described in detail below in combination with the drawings of the patent.

[0036] A video super-resolution inference method based on TensorRT includes the following three implementation steps:

[0037] Step (1) training the video super-resolution network;

[0038] In this stage, we need to train the model to get the deployed model, using the Basicvsr method as an example of our video super-resolution method, which can meet our needs in terms of parameter size and model performance. Using the Pytorch deep learning training framework, we can efficiently build the network and perform forward and backward propagation of gradients and update the model parameters. In addition, using mmediting, an open-source image and video editing tool code library based on Pytorch, we can efficiently configure the model parameters and training data sets required for training Basicvsr, and support multi-GPU parallel training to efficiently complete the model training task we need.

[0039] Step (2) analyzes the network operator, organizes common operators, writes unsupported operators, and optimizes operators.

[0040] In the network of Baiscvsr, most of the operators are common operators such as convolution, pooling, and activation functions. For these operators, they need to be organized because the operator API of TensorRT is different from that of Pytorch. After aligning the two operators, it will be easier to call them later, and it will be more convenient and efficient to deploy the video super-resolution network. However, there are also operators that are not common and are not supported by TensorRT, such as the flow_warp operator in the Spynet subnetwork included in Basicvsr. This subnetwork is used to extract the optical flow of adjacent frames to provide the optical flow map for subsequent use of inter-frame information. The flow_warp operator in this subnetwork is not supported by TensorRT and needs to be written by ourselves. For unsupported operators, we need to customize the TensorRT plugin, which needs to inherit the IPluginV2DynamicExt and PluginCreator classes and implement the member functions in this class, including the cuda functions required for inference and serialization functions, etc. For the optimizable operators such as convolution and activation functions, these two operators can be fused into one operator, which reduces the number of memory accesses and speeds up the overall calculation during inference.

[0041] Step (3) for the complete video super-resolution network structure, use TensorRT to write and generate a serialized engine for inference;

[0042] After the second step of analyzing the network operator and writing the unsupported operator, the network can be built. According to the video super-resolution network Basicvsr used, the corresponding operator is added layer by layer. Then the generation engine can be generated. In the generation process, TensorRT is used for compilation and optimization, which will automatically select a more reasonable and faster operator implementation method. After the engine is built, the optimal solution is obtained. When reasoning, the engine needs to be deserialized first, and then the video sequence is input. In the code, the input data is transmitted from the CPU to the GPU end, and the reasoning calculation is performed on the GPU. Finally, the data is transmitted from the GPU end to the CPU end.

[0043] A video super-resolution reasoning device based on TensorRT, comprising a video super-resolution network training module, a network operator analysis processing module, a serialized engine construction module, and a reasoning module.

[0044] The video super-resolution network training module is used to train the video super-resolution network to obtain the deployed model, using the Pytorch deep learning training framework. In addition, the mmediting configuration is used to configure the model parameters and training data set required for training Basicvsr.

[0045] The network operator analysis processing module is used to analyze the network operator, sort the commonly used operators, write the unsupported operators, and optimize the operators.

[0046] The network operator analysis processing module is implemented as follows:

[0047] The commonly used operators need to be sorted because the operator API of TensorRT is different from the operator API of Pytorch. After aligning the two operators, it is convenient for subsequent calling. The unsupported operators need to be written. For the unsupported operators, the TensorRT plug-in needs to be customized, the IPluginV2DynamicExt and PluginCreator classes need to be inherited, and the member functions in the class need to be implemented, including the cuda function required for reasoning and the serialization function. The optimized operators are fused into one operator, which reduces the number of memory accesses and speeds up the overall calculation during the reasoning process.

[0048] The unsupported operator is written and implemented as follows:

[0049] The flow_warp operator not supported by TensorRT is in the execution flow of Pytorch: first, a grid is created, the optical flow is added to the grid, and then the value is normalized to (-1, 1). Then, the grid_sample operator is called, and the input parameters are input, processed optical flow data, and interpolation mode. The execution flow of the grid_sample operator in Pytorch is to first denormalize, then find the value of the corresponding position index of the output on the optical flow, and perform linear interpolation on the input to obtain the result at the output.

[0050] When using the custom operator of TensorRT, the cuda code needs to be written to complete the foregoing forward execution process on Pytorch. By analyzing the execution flow of the flow_warp operator on Pytorch, it can be known that the normalization and denormalization of the optical flow data are opposite operations, which are omitted, and the grid of the optical flow data can be calculated through the horizontal and vertical axis indexes corresponding to each thread on the cuda and added to the optical flow. The above operator fusion process can reduce a lot of matrix transformation and operation. Finally, the linear interpolation operation can obtain the output result.

[0051] The serialization engine construction module adds the corresponding operator layer by layer according to the used video super-resolution network. Then, the engine is generated, and in the generation process, TensorRT is used for compilation and optimization, which can automatically select a more reasonable and faster operator implementation method.

[0052] After the engine is constructed, the optimal scheme is obtained, and when the inference module is used for inference, the engine needs to be deserialized first, and then the video sequence is input. In the code, the input data is transmitted from the CPU to the GPU end, the inference calculation is performed on the GPU, and finally the data is transmitted from the GPU end to the CPU end.

[0053] Embodiment

[0054] Figure 1 The TensorRT specific use flow of the embodiment of the application is as follows:

[0055] The application mainly includes the following implementation stages:

[0056] Step 1. Training the Basicvsr network:

[0057] Download the REDS4 dataset using the toolkit code in mmediting, and divide the training set and test set. Set the configuration file of training Basicvsr in mmediting, the main parameters include the training set and test set directory in Data, and set the workers on each GPU and the samples on each GPU according to the memory of the training device. In addition, the learning strategy needs to be modified, the total number of iterations is 300k, and the training strategy is CosineRestart. Train the model using two or four graphics cards and the parameters configured in the previous step. During the training process, the program will continuously output the changes of loss and the results of the test set at each stage, and the results of the test set are used to simply judge whether the training is problematic. The program can complete the training of the model for two to three days. Save and export the model. Compare the results of the test set to select the best trained model, save the model, and save the network structure and model parameters when exporting to become a disk file, which is convenient for reading and using in subsequent steps.

[0058] Step 2. Align TensorRt operators and Pytorch operators:

[0059] The API inconsistency between TensorRT operators and Pytorch operators brings many difficulties in the deployment process, and it is necessary to be familiar with the API of both to normally call. Therefore, for commonly used operators, alignment can be performed to facilitate subsequent deployment and optimization processes. Here, the commonly used convolution operator is used as an example. In the convolution API of Pytorch, the input is the input channel number, output channel number, kernel size, stride size, padding size, etc. While the input of TensorRT's convolution API is input, convolution weight, convolution bias, output channel number, kernel size, stride size, padding size, layer name, etc. We need to implement a function that inputs the convolution API of Pytorch, and calls the API of TensorRT inside the function, because the network of the model is trained by Pytorch, and Pytorch is more commonly used. After completing the function, the operator can be used very conveniently and quickly in subsequent calls.

[0060] Step 3. Write the TensorRT implementation of the flow_warp operator:

[0061] The execution flow of the flow_warp operator in Pytorch (i.e., the forward execution process) is as follows: a grid is first created, the optical flow is added to the grid, and then the value is normalized to (-1, 1). Then, the grid_sample operator is called with the input, the processed optical flow data, and the interpolation mode as input parameters. The execution flow of the grid_sample operator in Pytorch is as follows: first, the de-normalization is performed, then the value of the output corresponding position index on the optical flow is found, and the value is used for linear interpolation on the input to obtain the result at the output.

[0062] When using the custom operator of TensorRT, we need to write the cuda code to complete the forward execution process described above on Pytorch. By analyzing the execution flow of the flow_warp operator on Pytorch, it can be known that the normalization and de-normalization of the optical flow data are opposite operations and can be omitted. In addition, the grid of the optical flow data can be calculated through the horizontal and vertical axis indexes corresponding to each thread on the cuda and added to the optical flow. The above operator fusion process can reduce a lot of matrix transformation and operation. Finally, the linear interpolation operation can be performed to obtain the output result.

[0063] Figure 2 The Flow_warp operator calculation graph for the embodiment of the application.

[0064] Step 4. Build the Basicvsr network and perform inference:

[0065] According to the above steps, we have obtained the alignment of commonly used operators and the implementation of the flow_warp operator which is not supported by TensorRT. Next, we use C++ to call the operator to build the network structure of Basicvsr. And implement the forward inference. First, analyze the network structure of Basicvsr, which mainly includes convolution layers and activation function layers, and some residual connection structures. Therefore, for commonly used operators, we mainly use TensorRT's addConvolutionNd(), addPooling(), addSoftMax(), etc. to implement the building of the Basicvsr network. For the flow_warp operator, directly call the plug-in we implemented. TensorRT can optimize each layer of the network, including horizontal and vertical layer fusion, and operator acceleration, etc. After performing the optimization operation, the model can be serialized to the local, and when actually using the inference, only the serialized file needs to be deserialized to achieve the purpose of efficient inference.

[0066] The above description is further detailed in connection with specific / preferred embodiments of the present application, and should not be construed as limiting the specific implementation of the present application to these descriptions. For those skilled in the art to which the present application belongs, without departing from the concept of the present application, a number of substitutions or modifications can be made to the described embodiments, and these substitutions or modifications should be considered to belong to the protection scope of the present application.

[0067] The part of the present application not described in detail belongs to the technology known to those skilled in the art.

Claims

1. A TensorRT-based video super-resolution inference method, characterized in that, Comprise the following steps: Step (1) training the video super-resolution network; Step (2) analyzing the network operator, organizing common operators, writing unsupported operators, and optimizing operators; the specific implementation is as follows: Common operators need to be organized, the operator API of TensorRT is different from the operator API of Pytorch, after aligning the two operators, it is convenient for subsequent calling; The operator that TensorRT does not support needs to be written; for the unsupported operator, the TensorRT plugin needs to be customized, the IPluginV2DynamicExt and PluginCreator classes need to be inherited, and the member functions in the class need to be implemented, including the cuda function required for inference and the serialization function; the optimized operator fusion becomes an operator, which reduces the number of memory accesses in the inference process and speeds up the overall calculation; Step (3) for the complete video super-resolution network structure, use TensorRT to write and generate a serialized engine for inference; the specific implementation is as follows: According to the used video super-resolution network, add the corresponding operator layer by layer; then generate the engine, and in the generation process, compile and optimize it through TensorRT, which will automatically select a more reasonable and faster operator implementation method; After building the engine, the optimal solution is obtained, and when inferring, the engine needs to be deserialized first, and then the video sequence is input, the input data is transmitted from the CPU to the GPU in the code, the inference calculation is performed on the GPU, and finally the data is transmitted from the GPU to the CPU.

2. The TensorRT-based video super-resolution inference method according to claim 1, wherein, Step (1) the specific implementation is as follows: Train the selected video super-resolution network model to obtain the deployed model, use the Pytorch deep learning training framework; in addition, use mmediting to configure the model parameters and training data set required for training Basicvsr.

3. The video super-resolution inference method based on TensorRT according to any one of claims 1-2, characterized in that, The video super-resolution network uses Basicvsr.

4. The TensorRT-based video super-resolution inference method according to claim 1, wherein, The flow_warp operator that TensorRT does not support in the execution process of Pytorch is: first create a grid, add the optical flow to the grid, and then normalize the value to (-1, 1); then call the grid_sample operator with input, processed optical flow data and interpolation mode as input parameters; the execution process of the grid_sample operator in Pytorch is to first denormalize, then find the value of the output corresponding position index on the optical flow, and use the value to perform linear interpolation on the input to obtain the result at the output. When using the custom operator of TensorRT, the cuda code needs to be written to complete the forward execution process on Pytorch as described above; by analyzing the execution process of the flow_warp operator on Pytorch, it can be known that the normalization and denormalization of the optical flow data are opposite operations, which are omitted, in addition, the grid of the optical flow data can be calculated by each thread on the cuda corresponding to the horizontal and vertical axis indexes and added to the optical flow, the above operator fusion process can reduce matrix transformation and operation; finally, the linear interpolation operation can be performed to obtain the output result.

5. A TensorRT-based video super-resolution inference apparatus, characterized in that, The video super-resolution network training module, the network operator analysis processing module, the serialization engine construction module, and the inference module are included. The video super-resolution network training module is used to train the video super-resolution network to obtain a deployed model, using the Pytorch deep learning training framework; in addition, the model parameters required for training Basicvsr and the training data set are configured using mmediting; The network operator analysis processing module is used to analyze the network operator, organize common operators, write unsupported operators, and optimize operators. The serialization engine construction module adds corresponding operators layer by layer according to the used video super-resolution network; then generates an engine, and in the generation process, TensorRT is used for compilation and optimization to automatically select a more reasonable and faster operator implementation method; After the engine is constructed, the optimal solution is obtained, and when the inference module is used for inference, the engine needs to be deserialized first and then the video sequence is input; in the code, the input data is transmitted from the CPU to the GPU end, the inference calculation is performed on the GPU, and finally the data is transmitted from the GPU end to the CPU end; The network operator analysis processing module is implemented as follows: Common operators need to be organized, the operator API of TensorRT is different from the operator API of Pytorch, after aligning the two operators, it is convenient for subsequent calling; Unsupported operators of TensorRT need to be written; for unsupported operators, the TensorRT plugin needs to be customized, the IPluginV2DynamicExt and PluginCreator classes need to be inherited, and the member functions in the class need to be implemented, including the cuda function required for inference and the serialization function; for the optimized operator, it is fused into an operator, which reduces the number of memory accesses and speeds up the overall calculation during inference.

6. The video super-resolution inference apparatus based on TensorRT according to claim 5, wherein, The network operator analysis processing module writes the unsupported operators as follows: The operator flow_warp, which is not supported by TensorRT, has the following execution flow in Pytorch: first, a grid is created, and the optical flow is added to the grid, and then the value is normalized to (-1, 1); then, the grid_sample operator is called, and the input parameters are the input, the processed optical flow data, and the interpolation mode; the execution flow of the grid_sample operator in Pytorch is to first denormalize, then find the value of the output corresponding position index on the optical flow, and then use the value to perform linear interpolation on the input to obtain the output result; When using the custom operator of TensorRT, the cuda code needs to be written to complete the foregoing forward execution process on Pytorch; through the analysis of the execution flow of the flow_warp operator on Pytorch, it can be known that the normalization and denormalization of the optical flow data are opposite operations, which are omitted, and the grid of the optical flow data can be calculated through the horizontal and vertical axis indexes corresponding to each thread on the cuda and added to the optical flow; the above operator fusion process can reduce the matrix transformation and operation; finally, the linear interpolation operation can obtain the output result.

Images