A lightweight pruning method for optical flow models for edge deployment

Abstract

The present invention discloses a light-weight trimming method for an optical flow model for edge deployment, comprising the following steps: (1) obtaining a data set Flying Chairs as a training set and a validation set and performing preprocessing; (2) inputting the preprocessed data set into a global matching optical flow neural network, i.e., a GMFlow neural network, for training; (3) trimming the GMFlow neural network, i.e., performing light-weight processing; (4) obtaining a test set data MPI Sintel Dataset, testing the result obtained in step (3), and selecting the optimal solution; the present invention provides ideas for further developing the deployment of optical flow estimation on edge devices.

Description

Technical Field

[0001] The present invention relates to the technical field of optical flow estimation, and in particular to a lightweight clipping method of an optical flow model for edge deployment. Background Art

[0002] Optical flow estimation is a key research direction in computer vision. It has important applications in video interpolation, target tracking, and moving object segmentation. By inferring pixel-level motion information in continuous image frames, the position and speed of moving objects can be determined, and real-time tracking of targets can be achieved. Although deep learning has improved the accuracy of optical flow estimation, the resulting complex network architecture and increased computing time have put forward higher requirements on device computing power. With the development of fields such as autonomous driving and target tracking, mobile devices and embedded devices have become increasingly important in these applications. Edge devices achieve real-time response through local computing, avoid data transmission delays and bandwidth limitations, protect data privacy, and can be customized according to application requirements. However, mobile devices and embedded devices with low power consumption and limited computing resources are difficult to perfectly implement optical flow estimation algorithms, limiting their application value; there are the following problems: a large number of iterative operations greatly increase the time of optical flow estimation, which makes it difficult to optimize the model and deploy practical applications; the deployment of neural networks on edge devices needs to consider the amount of computing and model size. Summary of the invention

[0003] Purpose of the invention: The purpose of the present invention is to provide a lightweight pruning method for optical flow models for edge deployment, to solve the problem of how to optimize and lightweight the network to achieve efficient deployment on edge devices.

[0004] Technical solution: The light-weight trimming method of an optical flow model for edge deployment described in the present invention comprises the following steps:

[0005] (1) Obtain the Flying Chairs dataset as a training set and a validation set and perform preprocessing;

[0006] (2) Input the preprocessed data set into the global matching optical flow neural network, i.e., GMFlow neural network, for training;

[0007] (3) Pruning the GMFlow neural network to make it lightweight;

[0008] (4) Obtain the test set data MPI Sintel Dataset, test the results obtained in step (3), and select the optimal solution.

[0009] Furthermore, in step (1), the preprocessing is as follows: the Flying Chairs dataset is cropped to a uniform size.

[0010] Furthermore, in step (2), the GMFlow neural network includes a feature extraction module, a feature enhancement module, a feature matching module and a flow propagation module.

[0011] Furthermore, the feature extraction module is as follows: the input of the GMFlow neural network is two videos of the same size. In the feature extraction stage, the network first uses a weight-shared convolutional neural network to extract features from the two images; the convolutional neural network is composed of multiple residual blocks, each of which is composed of a convolutional layer, an activation function layer, and a normalization layer. The neural network completes downsampling while extracting features from the input image, and then transmits the feature map to the feature enhancement module.

[0012] Furthermore, the feature enhancement module is a Transformer module for establishing similarity relationships between features. The self-attention mechanism and the cross-attention mechanism are stacked in the Transformer module to obtain the relationship between features between images.

[0013] Furthermore, the feature matching module is specifically as follows: after completing the enhancement of the features, the similarity of the two feature maps is calculated pixel by pixel through a correlation layer for global feature matching, and a correlation matrix is ​​generated to compare the feature similarities.

[0014] Furthermore, the flow propagation module is specifically as follows: a softmax layer is used to determine the correspondence between pixels to obtain the result of optical flow estimation; and the matching optical flow estimation result is propagated to the unmatched pixels through the similarity of features.

[0015] Furthermore, step (3) is specifically as follows: pruning the neural network of the Transformer module in the feature enhancement, that is, the number of Transformer module variants is 2 layers or the feedforward neural network in the Transformer module is removed.

[0016] An electronic device described in the present invention includes a memory, a processor, and a computer program stored in the memory and executable on the processor. When the computer program is loaded into the processor, it implements any one of the methods for lightweight trimming of optical flow models for edge deployment.

[0017] A storage medium described in the present invention stores a computer program, and when the computer program is executed by a processor, any one of the methods for lightweight trimming of optical flow models for edge deployment is implemented.

[0018] Beneficial effects: Compared with the prior art, the present invention has the following significant advantages: two different lightweight schemes are implemented for GMFlow, and then the two lightweight networks are trained and tested with the final results. Finally, GMFlow_without_ffn achieves a network size reduction of 49.6% of the original size with a 24.5% reduction in accuracy. BRIEF DESCRIPTION OF THE DRAWINGS

[0019] Figure 1 The overall architecture of GMFlow of the present invention;

[0020] Figure 2 The change of loss in 100k trainings of the lightweight network of the present invention;

[0021] Figure 3 The change of EPE during 100k trainings of the lightweight network of the present invention. DETAILED DESCRIPTION

[0022] The technical solution of the present invention is further described below in conjunction with the accompanying drawings.

[0023] like Figure 1 As shown, an embodiment of the present invention provides a lightweight tailoring method for an optical flow model for edge deployment, which selects a neural network model suitable for the edge device according to the type and hardware specifications of the edge device, completes the lightweighting of the neural network, and performs a performance evaluation on the lightweight neural network, and finally converts the neural network into a format supported by the edge device; the method includes the following steps:

[0024] (1) Obtain the Flying Chairs dataset as a training set and a validation set and perform preprocessing. The preprocessing is as follows: Crop the Flying Chairs dataset to a uniform size.

[0025] (2) Input the preprocessed data set into the GMFlow neural network for training; the GMFlow neural network includes a feature extraction module, a feature enhancement module, a feature matching module, and a flow propagation module. The feature extraction module is as follows:

[0026] The input of the GMFlow neural network is two videos of the same size. In the feature extraction stage, the network first uses a weight-sharing convolutional neural network to extract features from the two images. The convolutional neural network consists of multiple residual blocks, each of which consists of a convolutional layer, an activation function layer, and a normalization layer. The neural network completes downsampling while extracting features from the input image, and then transmits the feature map to the feature enhancement module. The formula is as follows:

[0027] ;

[0028] Among them, C is the correlation matrix, which represents the correlation between the feature maps of the previous and next frames; F1 and F2 are feature maps extracted from the previous and next frames of the video, and each row represents the feature vector of a position in the feature map; by transposing F2, the features of the two feature maps are multiplied by matrix multiplication to calculate the feature similarity of different positions in the feature map; D represents the dimension of the feature vector of the feature map.

[0029] The feature enhancement module is a Transformer module used to establish similar relationships between features. In the Transformer module, the self-attention mechanism and the cross-attention mechanism are stacked to obtain the relationship between features of the images. The details are as follows: First, the window size is split to match the feature size in the image, that is, the input feature map of size H*W is split into K* K windows, each of which is H / K* H / K in size, and self-attention and cross-attention are calculated in each local window.

[0030] The feature matching module is as follows: after the enhancement of the features is completed, the similarity of the two feature maps is calculated pixel by pixel through a correlation layer for global feature matching, and a correlation matrix is ​​generated to compare the feature similarities.

[0031] The flow propagation module is specifically as follows: for the occluded areas in the image and the areas outside the boundary, due to the high similarity between the optical flow field and the image, the softmax layer is used to determine the correspondence between pixels to obtain the optical flow estimation result; the matching optical flow estimation result is propagated to the unmatched pixels through the similarity of features.

[0032] (3) Prune or lightweight the GMFlow neural network; the details are as follows:

[0033] In the feature extraction stage, GMFlow uses a 7*7 convolution kernel in the first layer of the network for feature extraction, and a 3*3 convolution kernel in the residual block for multi-layer feature extraction. At the same time, in the residual block, GMFlow mainly uses convolution layers, activation layers, and normalization layers, and does not rely heavily on fully connected layers to integrate feature information. Therefore, there is not much redundancy in the number of parameters and calculations of the convolution layer in GMFlow, so the main focus is on pruning the neural network for the Transformer module in feature enhancement.

[0034] In addition, the GMFlow neural network establishes global feature relationships by stacking 6 layers of Swin Transformer module variants. Therefore, the neural network can be pruned by reducing the number of Swin Transformer module variants to minimize the network size while ensuring the accuracy of optical flow estimation. In addition, in each Swin Transformer module variant, the feedforward neural network uses a large number of parameters and greatly increases the time for neural network optical flow estimation. Therefore, the network size can be reduced by removing the feedforward neural network.

[0035] Through the above analysis, two methods are adopted to achieve lightweight GMFlow, namely, the number of Transformer module variants is 2 layers or the feedforward neural network in the Transformer module is removed.

[0036] (4) Obtain the test set data MPI Sintel Dataset, test the results obtained in step (3), and select the optimal solution. Specifically: During the training process, the parameter settings are as follows: the learning rate is set to 0.0004, the batch_size is set to 16, the training round is set to 100K, the validation set is Flying Chairs, the validation frequency is set to 10K, and the frequency of saving weights is set to 10K. The optimizer used in the experiment is AdamW, and the measurement indicator used is the commonly used metric in optical flow estimation. The endpoint error (EPE) is the average Euclidean distance between the optical flow prediction result and the true result. In order to better understand the performance improvement, the present invention also reports the EPE of different motion amplitudes, that is, using , and to represent the EPE on 0-10, 10-40 and more than 40 pixels. The loss function used in training is L, is a discount factor, set to 0.9, Represents the error between the model prediction value and the true value. The formula is as follows:

[0037] ;

[0038] Here, N represents the total number of time steps.

[0039] The change of loss during training is as follows Figure 2 The changes of EPE during training are shown in Figure 3 The EPE of different range of motion during training is shown in Table 1.

[0040] Table 1 EPE of two lightweight models at different ranges of motion during training

[0041] ;

[0042] like Figure 2 and Figure 3 As shown in the figure, the blue line represents the GMFlow without the feedforward neural network in the Transformer variant module, but the number of Transformer variant modules is still 6 layers, which can be recorded as the global matching optical flow GMFlow_without_ffn without the feedforward neural network. The orange line represents the GMFlow with 2 layers of Transformer variant modules used in the feature enhancement module and the feedforward neural network retained, which can be recorded as the global matching optical flow GMFlow_2layers with two layers of Transformer retained. By analyzing the data, it can be found that the loss and EPE of GMFlow_2layers are better than GMFlow_without_ffn during training, and it also has better performance in EPE of different motion amplitudes. At the same time, GMFlow_2layers still performs better on the verification data set. In addition, GMFlow_2layers has a faster convergence speed during training, but overall, the time spent by the two lightweight networks to reach the global optimal point in training is similar. The test dataset is the MPI Sintel Dataset, an authoritative benchmark test set in the field of optical flow estimation. The dataset is divided into two parts: clean and final. The final dataset contains images with motion blur, while the clean dataset does not. GMFlow_without_ffn and GMFlow_2layers are tested on the clean and final parts of the MPI Sintel Dataset, respectively. The test results are shown in Table 2.

[0043] Table 2 EPE of the two lightweight models and the original model on the test set

[0044] ;

[0045] The EPE results at different motion amplitudes are shown in Table 3.

[0046] Table 3 EPE of the two lightweight models and the original model at different motion ranges in the test

[0047] ;

[0048] According to the overall EPE in the test results, GMFlow_2layers performs better than GMFlow_without_ffn on the MPI Sintel Dataset, but the difference between the two is very small. When focusing on the EPE of different amplitudes, the situation changes. GMFlow_without_ffn has a better performance in the range of 0-10 pixels, but GMFlow_2layers has a greater advantage in the range of more than 40 pixels, and in the range of 10-40 pixels, the two lightweight networks have a win or loss. As for the difference in performance over a large range of pixels, although GMFlow_without_ffn has 6 layers of Transformer modules superimposed, due to the lack of a feedforward neural network, it cannot enhance features with a large spatial span, and thus cannot grasp their similarity relationship, so it performs poorly on large displacements. GMFlow_2layers uses 2 complete Transformer modules, so it can perform more accurate similarity calculations on features. In the estimation of optical flow within a small range of pixels, GMFlow_without_ffn uses multiple layers of Transformer stacking. Although it lacks the integration of feature information by the feedforward neural network, it still achieves more accurate estimation after multiple iterations. However, GMFlow_2layers performs worse due to the disadvantage of the number of Transformer modules. After pruning the GMFlow network, the parameters of GMFlow_without_ffn and GMFlow_2layers are shown in Table 4.

[0049] Table 4 Comparison between lightweight model and original model

[0050] ;

[0051] By comparison, GMFlow was successfully lightweighted, with the accuracy losses of GMFlow_without_ffn and GMFlow_2layers being 24.5% and 24.2% respectively, successfully reducing the model size to 49.6% and 55.1% of the original. After considering the accuracy of optical flow estimation and the network volume, GMFlow_without_ffn will be used as the deployment model in the implementation of this invention, and the lightweight network will be converted into an ONNX intermediate model. The two tools ONNX Simplifier and Optimizer provided by the ONNX framework will be used to optimize the ONNX model, reduce the model complexity, and check whether the model complies with the ONNX specification. Convert the ONNX model to a model supported by the NCNN framework, and obtain the .para and .bin files under the NCNN framework. Depending on the edge devices, different applications can be developed to complete the implementation of the optical flow estimation network.

[0052] An embodiment of the present invention also provides an electronic device, including a memory, a processor, and a computer program stored in the memory and executable on the processor, wherein when the computer program is loaded into the processor, the method for lightweight trimming of optical flow models for edge deployment described in any one of the above is implemented.

[0053] An embodiment of the present invention further provides a storage medium storing a computer program, wherein the computer program, when executed by a processor, implements any one of the methods for lightweight trimming of optical flow models for edge deployment.

Claims

1. A lightweight pruning method for optical flow model for edge deployment, characterized in that: The following steps are involved: (1) Obtain the Flying Chairs dataset as a training set and a validation set and perform preprocessing. The preprocessing is as follows: Crop the Flying Chairs dataset to a uniform size; (2) The preprocessed data set is input into the global matching optical flow neural network, i.e., GMFlow neural network, for training; the GMFlow neural network includes a feature extraction module, a feature enhancement module, a feature matching module, and a flow propagation module; the feature extraction module is as follows: the input of the GMFlow neural network is two videos of the same size before and after. In the feature extraction stage, the network first uses a weight-sharing convolutional neural network to extract features from the two images; the convolutional neural network consists of multiple residual blocks, each of which consists of a convolutional layer, an activation function layer, and a normalization layer. The neural network completes downsampling while extracting features from the input image, and then transmits the feature map to the feature map. The feature enhancement module is a Transformer module used to establish similarity between features. In the Transformer module, the self-attention mechanism and the cross-attention mechanism are stacked to obtain the relationship between features between images. The feature matching module is as follows: after completing the feature enhancement, the similarity of the two feature maps is calculated pixel by pixel through a correlation layer for global feature matching, and a correlation matrix is ​​generated to compare feature similarity. The flow propagation module is as follows: the softmax layer is used to determine the correspondence between pixels, so as to obtain the result of optical flow estimation; the matched optical flow estimation result is propagated to the unmatched pixels through the similarity of the features. (3) The GMFlow neural network is pruned, i.e., lightweight. Specifically, in the feature extraction stage, GMFlow uses a 7*7 convolution kernel in the first layer of the network for feature extraction, while the convolution kernel in the residual block uses a 3*3 convolution kernel for multi-layer feature extraction. At the same time, in the residual block, GMFlow mainly uses convolution layers, activation layers, and normalization layers, and prunes the neural network by reducing the number of Swin Transformer module variants. The neural network of the Transformer module in feature enhancement is pruned, i.e., the number of Transformer module variants is 2 layers or the feedforward neural network in the Transformer module is removed. (4) Obtain the test set data MPI Sintel Dataset, test the results obtained in step (3), and select the optimal solution; specifically: use GMFlow_without_ffn as the deployment model, convert the lightweight network into an ONNX intermediate model, and use the two tools ONNX Simplifier and Optimizer provided by the ONNX framework to optimize the ONNX model; check whether the model complies with the ONNX specification; convert the ONNX model into a model supported by the NCNN framework, and obtain the .para and .bin files under the NCNN framework. According to different edge devices, different applications can be developed to complete the implementation of the optical flow estimation network.

2. An electronic device comprising a memory, a processor, and a computer program stored in the memory and executable on the processor, characterized in that: When the computer program is loaded into a processor, the method for lightweight pruning of optical flow models for edge deployment according to claim 1 is implemented.

3. A storage medium storing a computer program, characterized in that: When the computer program is executed by a processor, the method for lightweight pruning of optical flow models for edge deployment according to claim 1 is implemented.

Images