A method for constructing a multi-task two-stage structure based on a convolutional neural network
By constructing a multi-task two-stage structure based on convolutional neural networks, the problem of difficulty in reusing multi-task image data in existing technologies is solved, achieving efficient training on the same image data and flexible switching between sub-tasks, thus reducing the training burden of the model.
Patent Information
- Authority / Receiving Office
- CN · China
- Patent Type
- Patents(China)
- Current Assignee / Owner
- HANGZHOU ZHUOXI INST OF BRAIN & INTELLIGENCE
- Filing Date
- 2022-08-24
- Publication Date
- 2026-07-07
AI Technical Summary
Existing technologies lack suitable machine learning models for performing multiple tasks on the same image data, resulting in the inability to effectively reuse sub-tasks and increasing the training burden.
A multi-task two-stage structure based on convolutional neural networks is constructed, including a left branch model, a right branch model, and long connections. By combining segmentation heads, detection heads, and classification heads, a U-shaped convolutional neural network is used for multi-task training and transfer training to achieve reuse between sub-tasks.
It improves the efficiency of multi-task projects, reduces the burden of model training, supports pre-training and transfer training, and enables flexible switching between data-similar subtasks.
Smart Images
Figure CN116152145B_ABST
Abstract
Description
Technical Field
[0001] This invention relates to the field of machine learning technology, and specifically to a method for constructing a two-stage multi-task structure based on a convolutional neural network. Background Technology
[0002] Deep learning technology is developing rapidly, and the U-shaped convolutional neural network model has particularly prominent characteristics. First, it is shaped like a capital U, consisting of two branches, one for encoding and the other for decoding the image. Each branch is a multi-layered structure, capable of extracting and combining image features. Furthermore, long-distance connections exist between the branches, ensuring the full protection and flow of image feature information at different stages. It is generally believed that deeper convolutional networks are better at processing large images, but less effective at handling finer details, and vice versa. Image recognition technology based on neural network models has been widely and extensively used in various fields.
[0003] However, in practical applications, it has been found that large projects often consist of multiple sub-tasks. These sub-tasks often share high similarity in image data but differ in image segmentation and recognition, making it difficult to reuse machine learning models effectively. For example, in processing brain CT medical data, sub-tasks include infarction detection, fracture detection, hemorrhage segmentation, and tumor segmentation. These sub-tasks are distinct but all performed on image data of the same patient's brain region. Therefore, designing multiple different models to handle these sub-tasks separately is cumbersome. Furthermore, since each sub-task is a separate model, there is no information sharing between them, increasing the burden of training the machine learning model. Therefore, it is necessary to research a multi-task machine learning model that can be easily transferred and reused to achieve a technical solution for reusing sub-tasks with similar data. Summary of the Invention
[0004] The technical problem this invention aims to solve is the current lack of suitable machine learning models for performing multiple tasks on the same image data. This invention proposes a method for constructing a two-stage multi-task structure based on convolutional neural networks, which enables reuse among sub-tasks with similar data, thereby improving the overall project efficiency.
[0005] To solve the above technical problems, the present invention adopts the following technical solution: a method for constructing a multi-task two-stage structure based on a convolutional neural network, wherein the convolutional neural network is a U-shaped convolutional neural network, including a left branch model and a right branch model, with long connections between the left branch model and the right branch model, which are used for image encoding and decoding respectively, including the following steps:
[0006] A segmentation head is established, the input of which is the output segmentation map of a U-shaped convolutional neural network;
[0007] A detection head is established, and the input of the detection head is the multi-layer feature map on the right side of the U-shaped convolutional neural network;
[0008] The segmentation head and detection head are used as a one-stage model;
[0009] A classification head is established, whose inputs are the feature map of the bottom layer of the U-shaped convolutional neural network, the output of the segmentation head, and the output of the detection head. The output of the classification head is used as the final segmentation map, and the classification head is used as a two-stage model.
[0010] Acquire several sample images, input the sample images into the U-shaped convolutional neural network, calculate the loss function using the output of the classification head, train the parameters of the U-shaped convolutional neural network, the segmentation head, the detection head, and the classification head until the output of the classification head reaches a preset accuracy, and complete the construction of the multi-task two-stage structure.
[0011] Preferably, the segmentation head is left unused, and the classification head is directly connected to the output segmentation map of the U-shaped convolutional neural network, or the segmentation head is established as a parameterless connection.
[0012] Preferably, an ASPP pyramid structure front-end is added to the detection head. The ASPP pyramid structure front-end is connected between the detection head and the multi-layer feature map on the right side of the U-shaped convolutional neural network. The input of the ASPP pyramid structure front-end is the multi-layer feature map on the right side of the U-shaped convolutional neural network, and the output of the ASPP pyramid structure front-end is used as the input of the detection head.
[0013] Preferably, a post-processing module is also included. The post-processing module is located before the classification head. The feature map of the bottom layer of the U-shaped convolutional neural network, the output of the segmentation head, and the output of the detection head are connected to the post-processing module. The output of the post-processing module is the image cropped out of the detection box, and the output of the post-processing module is used as the input of the classification head.
[0014] Preferably, the classification head is a ResNet18 model.
[0015] Preferably, when training the parameters of the U-shaped convolutional neural network, the segmentation head, the detection head, and the classification head, the long connections of the U-shaped convolutional neural network are disabled during the training of the first few sample images.
[0016] Preferably, when training the parameters of the U-shaped convolutional neural network, the segmentation head, the detection head, and the classification head, the U-shaped convolutional neural network is first trained separately using several sample images. Then, the segmentation head, the detection head, and the classification head are connected to the U-shaped convolutional neural network, and the parameters of the U-shaped convolutional neural network, the segmentation head, the detection head, and the classification head are trained using sample images.
[0017] As a preferred method, the method for obtaining sample images is as follows: acquire several images, randomly cover part of the image, output the restored image by the classification head, and use the difference between the restored image and the original image as the loss function.
[0018] Preferably, the difference between the restored image and the original image is the sum of the differences between all pixels in the restored image and the original image within the covered area.
[0019] Preferably, a transfer training method is also included for applying the already trained U-shaped convolutional neural network and two-stage structure to new image recognition, the transfer training method comprising:
[0020] Several sample images are reacquired, and a portion of the sample images is randomly covered. The reconstructed image is output by the classification head, and the difference between the reconstructed image and the original image is used as the loss function.
[0021] The parameters of the left encoder of the U-shaped convolutional neural network are fixed, and only the parameters of the left decoder, segmentation head, detection head and classification head of the U-shaped convolutional neural network are trained until the output of the classification head reaches the preset accuracy.
[0022] End training or open the training of the parameters of the left encoder of the U-shaped convolutional neural network, and continue training until the output of the classification head reaches the preset accuracy.
[0023] The beneficial technical effects of the present invention include: the two-stage structure provided by the present invention can flexibly stack models, support pre-training and transfer training, and can easily switch between multiple sub-tasks with the same or similar data, effectively reducing the burden of model training and improving project implementation efficiency.
[0024] Other features and advantages of the present invention will be disclosed in detail in the following detailed description and accompanying drawings. Attached Figure Description
[0025] The invention will be further described below with reference to the accompanying drawings:
[0026] Figure 1 This is a schematic diagram of a multi-task two-stage structure according to an embodiment of the present invention.
[0027] Figure 2 This is a schematic diagram of the multi-task two-stage structure for an embodiment of the present invention.
[0028] Figure 3 This is a schematic diagram of the multi-task two-stage structure transfer training method according to an embodiment of the present invention.
[0029] The components are: 1. Left branch model, 2. Long connection, 3. Right branch model, 4. Segmentation head, 5. ASPP pyramid structure front head, 6. Detection head, 7. Post-processing module, and 8. Classification head. Detailed Implementation
[0030] The technical solutions of the embodiments of the present invention will be explained and described below with reference to the accompanying drawings. However, the following embodiments are only preferred embodiments of the present invention and not all of them. Other embodiments obtained by those skilled in the art based on the embodiments in the implementation methods without creative effort are all within the protection scope of the present invention.
[0031] In the following description, terms such as “inner,” “outer,” “upper,” “lower,” “left,” and “right” are used only to indicate orientation or positional relationship for the convenience of describing the embodiments and simplifying the description, and are not intended to indicate or imply that the device or element referred to must have a specific orientation, or be constructed and operated in a specific orientation, and therefore should not be construed as a limitation of the present invention.
[0032] A method for constructing a two-stage multi-task structure based on convolutional neural networks is provided in the appendix. Figure 1 The convolutional neural network used in this embodiment is a U-shaped convolutional neural network, which includes a left branch model 1 and a right branch model 3. There is a long connection 2 between the left branch model 1 and the right branch model 3, which are used for image encoding and decoding, respectively.
[0033] Please see the appendix Figure 2 The method for constructing a two-stage multi-task structure based on convolutional neural networks includes the following steps:
[0034] Step A01) Create segmentation head 4. The input of segmentation head 4 is the output segmentation map of the U-shaped convolutional neural network.
[0035] Step A02) Establish detection head 6. The input of detection head 6 is the multi-layer feature map on the right side of the U-shaped convolutional neural network;
[0036] Step A03) Segmentation head 4 and detection head 6 are used as a first-stage model;
[0037] Step A04) Establish classification head 8. The input of classification head 8 is the feature map of the last layer of the U-shaped convolutional neural network, the output of segmentation head 4 and the output of detection head 6. The output of classification head 8 is used as the final segmentation map. Classification head 8 is used as a two-stage model.
[0038] Step A05) Obtain several sample images, input the sample images into the U-shaped convolutional neural network, calculate the loss function based on the output of the classification head 8, train the parameters of the U-shaped convolutional neural network, the segmentation head 4, the detection head 6 and the classification head 8 until the output of the classification head 8 reaches the preset accuracy, and complete the construction of the multi-task two-stage structure.
[0039] Step A06) Empty segmentation head 4, classify head 8 directly connects to the output segmentation map of the U-shaped convolutional neural network, or establish segmentation head 4 as a parameterless connection.
[0040] Please refer to the appendix again. Figure 1 In this embodiment, an ASPP pyramid structure pre-head 5 is added to the detection head 6. The ASPP pyramid structure pre-head 5 is connected between the detection head 6 and the multi-layer feature maps on the right side of the U-shaped convolutional neural network. The input of the ASPP pyramid structure pre-head 5 is the multi-layer feature maps on the right side of the U-shaped convolutional neural network, and the output of the ASPP pyramid structure pre-head 5 is used as the input of the detection head 6. The detection head 6 utilizes the multi-layer feature maps on the right side of the U-shaped structure. Since the objects in detection tasks typically vary in size, this approach is reasonable. For small object detection tasks, it is necessary to remove, for example, the bottom layer feature map as input, and vice versa, providing great flexibility. Furthermore, since the output sizes of the multiple layers on the right side of the U-shaped structure differ, an ASPP pyramid structure or its various variations are used as a tool to unify the output size. Of course, this structure itself also has the ability to process feature images.
[0041] Please refer to the appendix again. Figure 1 This embodiment also includes a post-processing module 7, which is positioned before the classification head 8. The feature map of the bottom layer of the U-shaped convolutional neural network, the output of the segmentation head 4, and the output of the detection head 6 are connected to the post-processing module 7. The output of the post-processing module 7 is the image within the cropped detection box, and it serves as the input to the classification head 8. The model selection for the classification head 8 is very flexible, such as a lightweight ResNet18. In this embodiment, the classification head 8 is a ResNet18 model. The segmentation head 4, detection head 6, classification head 8, and ASPP pyramid structure pre-head 5 are all models containing several layers of neurons with connections between them.
[0042] When training the parameters of the U-shaped convolutional neural network, segmenter 4, detector 6, and classifier 8, long connection 2 of the U-shaped convolutional neural network is disabled during the training of the first few sample images. This prevents the possibility of parameter cheating in the subsequent right-side branch model 3. When training the parameters of the U-shaped convolutional neural network, segmenter 4, detector 6, and classifier 8, the U-shaped convolutional neural network is first trained separately using several sample images. Then, segmenter 4, detector 6, and classifier 8 are connected to the U-shaped convolutional neural network, and the parameters of the U-shaped convolutional neural network, segmenter 4, detector 6, and classifier 8 are mechanically trained using sample images.
[0043] The method for obtaining sample images is as follows: acquire several images, randomly occlude a portion of each image, and output the reconstructed image from the classification head. The difference between the reconstructed image and the original image is used as the loss function. This method can be combined with various pre-training techniques. This embodiment uses occlusion reconstruction as an example. Because the input and output dimensions of a U-shaped convolutional neural network model are the same, self-supervised learning becomes possible. A common approach is to randomly occlude a random region of an image and then let the model learn how to reconstruct the occluded image region. In this process, the model acquires the ability to extract features. The difference between the reconstructed image and the original image is the sum of the differences between all pixels in the reconstructed image and the original image within the occluded region.
[0044] It also includes transfer learning methods for using already trained U-shaped convolutional neural networks and two-stage structures for new image recognition; please refer to the appendix. Figure 3 Transfer training methods include:
[0045] Step B01) Reacquire several sample images, randomly cover part of the sample images, and output the restored image from the classification head 8. The difference between the restored image and the original image is used as the loss function.
[0046] Step B02) Fix the parameters of the left encoder of the U-shaped convolutional neural network, and train only the parameters of the left decoder, segmentation head 4, detection head 6 and classification head 8 of the U-shaped convolutional neural network until the output of classification head 8 reaches the preset accuracy.
[0047] Step B03) End training or open the training of the parameters of the left encoder of the U-shaped convolutional neural network and continue training until the output of the classification head 8 reaches the preset accuracy.
[0048] The beneficial technical effects of this embodiment include: the two-stage structure provided by the present invention can flexibly stack models, support pre-training and transfer training, and can easily switch between multiple sub-tasks with the same or similar data, effectively reducing the burden of model training and improving project implementation efficiency.
[0049] The above description is merely a specific embodiment of the present invention, but the scope of protection of the present invention is not limited thereto. Those skilled in the art should understand that the present invention includes, but is not limited to, the contents described in the accompanying drawings and the specific embodiments above. Any modifications that do not depart from the functional and structural principles of the present invention will be included within the scope of the claims.
Claims
1. A method for constructing a multi-task two-stage structure based on a convolutional neural network, wherein the convolutional neural network is a U-shaped convolutional neural network, including a left branch model and a right branch model, with long connections between the left branch model and the right branch model, respectively used for image encoding and decoding, characterized in that, Includes the following steps: A segmentation head is established, the input of which is the output segmentation map of a U-shaped convolutional neural network; A detection head is established, and the input of the detection head is the multi-layer feature map on the right side of the U-shaped convolutional neural network; The segmentation head and detection head are used as a one-stage model; A classification head is established, whose inputs are the feature map of the bottom layer of the U-shaped convolutional neural network, the output of the segmentation head, and the output of the detection head. The output of the classification head is used as the final segmentation map, and the classification head is used as a two-stage model. Acquire several sample images, input the sample images into the U-shaped convolutional neural network, calculate the loss function using the output of the classification head, train the parameters of the U-shaped convolutional neural network, the segmentation head, the detection head and the classification head until the output of the classification head reaches the preset accuracy, and complete the construction of the multi-task two-stage structure; The method for obtaining sample images is as follows: acquire several images, randomly cover part of the image, output the restored image by the classification head, and use the difference between the restored image and the original image as the loss function.
2. The method for constructing a two-stage multi-task structure based on a convolutional neural network according to claim 1, characterized in that, The segmentation head is left unused, and the classification head is directly connected to the output segmentation map of the U-shaped convolutional neural network, or the segmentation head is established as a parameterless connection.
3. A method for constructing a two-stage multi-task structure based on a convolutional neural network according to claim 1 or 2, characterized in that, An ASPP pyramid structure front-end is added to the detection head. The ASPP pyramid structure front-end is connected between the detection head and the multi-layer feature map on the right side of the U-shaped convolutional neural network. The input of the ASPP pyramid structure front-end is the multi-layer feature map on the right side of the U-shaped convolutional neural network, and the output of the ASPP pyramid structure front-end is used as the input of the detection head.
4. A method for constructing a two-stage multi-task structure based on a convolutional neural network according to claim 1 or 2, characterized in that, It also includes a post-processing module, which is set before the classification head. The feature map of the bottom layer of the U-shaped convolutional neural network, the output of the segmentation head, and the output of the detection head are connected to the post-processing module. The output of the post-processing module is the image cropped out of the detection box, and the output of the post-processing module is used as the input of the classification head.
5. A method for constructing a two-stage multi-task structure based on a convolutional neural network according to claim 1 or 2, characterized in that, The classification head is a ResNet18 model.
6. A method for constructing a two-stage multi-task structure based on a convolutional neural network according to claim 1 or 2, characterized in that, When training the parameters of the U-shaped convolutional neural network, segmentation head, detection head, and classification head, the long connections of the U-shaped convolutional neural network are disabled during the training of the first few sample images.
7. A method for constructing a two-stage multi-task structure based on a convolutional neural network according to claim 1 or 2, characterized in that, When training the parameters of the U-shaped convolutional neural network, the segmentation head, the detection head, and the classification head, the U-shaped convolutional neural network is first trained separately using several sample images. Then, the segmentation head, the detection head, and the classification head are connected to the U-shaped convolutional neural network, and the parameters of the U-shaped convolutional neural network, the segmentation head, the detection head, and the classification head are trained using sample images.
8. The method for constructing a two-stage multi-task structure based on a convolutional neural network according to claim 1, characterized in that, The difference between the restored image and the original image is the sum of the differences between all pixels in the restored image and the original image within the covered area.
9. A method for constructing a two-stage multi-task structure based on a convolutional neural network according to claim 1, characterized in that, It also includes a transfer training method for using a trained U-shaped convolutional neural network and a two-stage structure for new image recognition, the transfer training method including: Several sample images are reacquired, and a portion of the sample images is randomly covered. The reconstructed image is output by the classification head, and the difference between the reconstructed image and the original image is used as the loss function. The parameters of the left encoder of the U-shaped convolutional neural network are fixed, and only the parameters of the left decoder, segmentation head, detection head and classification head of the U-shaped convolutional neural network are trained until the output of the classification head reaches the preset accuracy. End training or open the training of the parameters of the left encoder of the U-shaped convolutional neural network, and continue training until the output of the classification head reaches the preset accuracy.