Lightweight double-branch graph capsule network model applied to image classification

Abstract

The invention relates to the technical field of image processing and deep learning, in particular to a lightweight double-branch graph capsule network model applied to image classification. Comprising an input layer used for receiving a to-be-classified image; the double-branch feature extraction layer is connected with the input layer, consists of two parallel branches and is used for performing multi-scale feature extraction on the image to generate a feature map; the primary capsule layer is connected with the double-branch feature extraction layer and is used for converting the feature map into a plurality of primary capsule vectors; the multi-head attention graph pooling module is connected with the primary capsule layer and is used for performing consistent voting on the primary capsule vector to generate a digital capsule vector; the output layer is connected with the multi-head attention image pooling module and used for determining the classification result of the image according to the model length of the digital capsule vector, and the classification accuracy of the model on a complex image data set is remarkably improved.

Description

Technical Field

[0001] This invention relates to the fields of image processing and deep learning technology, specifically a lightweight bi-branch graph capsule network model for image classification. Background Technology

[0002] Classical Convolutional Neural Networks (CNNs) possess translation invariance, ensuring that the output remains the same regardless of where the target object is translated in the image. However, they lack rotation invariance, leading to misclassification if the target object is rotated by a certain angle. This is essentially because CNNs use individual neurons to store image features, representing image feature information as scalars, which fails to recognize the relationship between the whole and its parts. Furthermore, the pooling operation in CNNs can cause the loss of much valuable information. Therefore, CNNs require a large dataset for training to achieve good generalization performance.

[0003] To address the lack of semantic representation in CNNs, Sabour et al. first proposed Capsule Network (CapsNet) in November 2017. It encodes image features into capsules by introducing a pose matrix; each capsule is a vector with size and orientation. CapsNet uses a dynamic routing algorithm to construct the whole-part relationships of objects, giving the neural network rotation invariance. Furthermore, CapsNet achieves good generalization ability with a small dataset, thus overcoming the shortcomings of CNNs. While Sabour's Capsule Network achieves very high classification accuracy on the MNIST dataset, its accuracy is not high on datasets with complex image content such as FashionMNIST and CIFAR10. Therefore, numerous improvements have been made to Capsule Network image classification research both domestically and internationally.

[0004] The aim of improving the feature extraction part of capsule networks is to extract more features from images with fewer parameters, thereby obtaining high-quality capsules. Vanderson et al. proposed a parallelizable multi-lane capsule network (MLCN), where each lane independently acts on one dimension of the digital capsule layer, resulting in richer feature information. Chang et al. proposed a multi-lane capsule network with strict-squash (MLSCN), which, based on MLCN, introduces weights for each lane and designs a strict compression function to replace the original compression function, thereby filtering out lanes that negatively affect the classification results. Compared with the original capsule network, the classification accuracy in CIFAR10 is improved by about 17%. Xiang et al. proposed a multi-scale capsule network (MS-CapsNet), which uses multiple convolutional kernels of different sizes and lengths to extract features from the image at multiple scales, and then encodes the features at different scales into a multi-dimensional primary capsule layer. Therefore, MS-CapsNet converges more easily than the original capsule network. YANG et al. proposed a capsule network (Res2Netand Squeeze-and-Excitation Capsule Network, RS-CapsNet) consisting of a Res2Net module and a compressed excitation network. This network uses four sets of Res2Net blocks with small convolutional kernels to extract multi-scale features from images, thereby improving classification performance.

[0005] Jia et al. proposed the Residual Capsule Network (Res-CapsNet), which uses residual networks to extract features from images, increasing network depth while preventing network degradation. Li Jianqiao et al. proposed MFF-CapsNet, which fuses extracted shallow, mid-, and deep features to extract more comprehensive features, achieving accuracies of 94.20% and 87.21% on the Fashion-MNIST and CIFAR10 datasets, respectively. Wu Lin et al. proposed a lightweight multi-branch RA capsule network, which uses a residual attention module for image feature extraction and designs two lightweight branches of different depths to improve network classification performance. Ryu et al. used a three-dimensional convolutional capsule network for image feature extraction, achieving higher classification accuracy for hyperspectral images than most deep learning networks. Shiri et al. proposed the Convolutional Fully-Connected Capsule Network (CFC-CapsNet), which creates a fully connected convolutional layer to transform extracted low-level features into higher-dimensional, fewer capsules, thus reducing the number of network parameters.

[0006] Huang et al. proposed a dual attention mechanism capsule network (DA-CapsNet), which adds a first attention mechanism called Conv-Attention after the convolutional layer, and a second attention mechanism called Caps-Attention after the primary capsule layer. Xu et al. proposed a multi-scale capsule network with octave convolution (MSOctCaps), which combines octave convolution with capsule networks to extract multi-scale features from images, thereby improving classification accuracy.

[0007] The aim of improving the dynamic routing algorithm of capsule networks is to enhance network robustness and accelerate network training. Gu et al. demonstrated through multiple comparative experiments that the dynamic routing algorithm of capsule networks is not necessarily more robust than CNNs, and proposed an interpretable graph capsule network (Graph CapsNet), in which the routing algorithm employs a graph pooling method based on multi-head attention, achieving better classification performance with fewer parameters. Vittorio Mazzia et al. proposed a more efficient capsule network, employing a non-iterative, highly parallelizable routing algorithm that effectively suppresses useless capsules using a self-attention mechanism. Amira et al. proposed Weighted Capsule Fuzzy Gaussian Model-Routing (WCFGM-R), which introduces a novel pose loss function that provides the optimal inter-class separation distance, thereby improving the model's classification performance. In December 2019, Zhang et al. proposed a capsule network with bottleneck residual blocks and automatic progressive pruning. The bottleneck residual blocks can effectively reduce the number of parameters in the primary capsule layer, and the automatic progressive pruning prunes capsules that contribute less to the dynamic routing process. Song Yan et al. proposed an attention-based shared parameter capsule network, which introduces an attention mechanism into the routing algorithm to calculate the correlation between low-level capsules and reduces the number of parameters in the dynamic routing algorithm by sharing the transformation matrix.

[0008] While the above research has made significant breakthroughs, its classification performance is only good on a small number of simple datasets. There is still room for optimization on most complex datasets with noise and background information, such as CIFAR10 and SVHN classification tasks, which affects the quality of image classification. Summary of the Invention

[0009] To overcome the shortcomings of existing technologies, this invention proposes a lightweight bi-branch graph capsule network model for image classification, aiming to address the problems of insufficient feature extraction capability, high computational cost, and slow training speed of existing capsule networks in complex image classification tasks. This model achieves multi-scale feature extraction through a bi-branch structure and combines multi-head attention graph pooling to improve training efficiency and classification accuracy.

[0010] The technical invention adopted by this invention to solve its technical problem is: a lightweight bi-branch graph capsule network model described in this invention, applied to image classification, comprising: The input layer is used to receive the images to be classified. A dual-branch feature extraction layer, connected to the input layer, consists of two parallel branches and is used to extract multi-scale features from the image to generate a feature map. The primary capsule layer, connected to the dual-branch feature extraction layer, is used to convert the feature map into multiple primary capsule vectors; A graph pooling module with multi-head attention is connected to the primary capsule layer and is used to perform consensus voting on the primary capsule vectors to generate digital capsule vectors. The output layer, connected to the graph pooling module of the multi-head attention, is used to determine the classification result of the image based on the magnitude of the digital capsule vector.

[0011] Preferably, the dual-branch feature extraction layer includes: branch one, which employs a residual network with skip connections to extract deep features of the image; Branch 2 employs a residual octave convolution module to extract multi-scale features of the image by fusing high and low frequencies.

[0012] Preferably, the residual octave convolution module includes an initial octave convolution, a middle octave convolution, and a terminal octave convolution; wherein the input of the middle octave convolution includes a high-frequency feature map. and low-frequency feature maps Output and It can be calculated using the following formula:

[0013] Where f represents the convolution operation, W represents the convolution kernel weights, upsample represents the upsampling operation, and pool represents the pooling operation.

[0014] Preferably, the graph pooling module for multi-head attention performs the following steps: The feature maps of the primary capsule layer are divided into L groups, with each group containing feature maps P. i Through the transformation matrix W 1i Construct as a header X i The following formula is used: in, D in D represents the dimension of the capsules in the primary capsule layer. out For each image, the dimension of the capsules in the digital capsule layer is [dimension]. It consists of K×K nodes, each node representing a capsule. , where i represents the i-th head, j represents the j-th capsule on the i-th head, and the values ​​of each dimension of the same capsule are the values ​​of feature maps from different channels but at the same position; Calculate the attention matrix A for each graph, where A is the element of the attention matrix. ab Calculated using the following formula:

[0015] in, a and b represent the index values ​​of the nodes in the graph, P m P n This represents the row and column coordinates corresponding to the two nodes with index values ​​a and b in the graph. P m P n From (1,1) to (K, K), For hyperparameters; The attention coefficient matrix for each head is calculated using the attention matrix A and the transformation matrix W2, using the following formula:

[0016] in, These are trainable parameters, where M is the number of capsules in the digit capsule layer, i.e., the number of categories in the dataset to be classified. , indicating L heads with K 2 The relative importance of each primary capsule to M numerical capsules; The prediction result for each head pair of digital capsules is calculated using the following formula:

[0017] in, Generate M dimensions of D for each head out Digital capsules; The average of the prediction results for L heads is calculated, and then compressed using the squash function to obtain the final digital capsule vector.

[0018] Among them, .

[0019] Preferably, the squash function is defined as:

[0020] The squash function is used to compress M high-level capsules in the digital capsule layer into a vector V with the same direction but a length of 0-1.

[0021] Preferably, the model uses a combination of an edge loss function and a reconstruction loss function as the total loss function; wherein, the edge loss function L k Defined as:

[0022] in, This indicates whether category k exists; the value is 1 if it exists, and 0 otherwise. + Let m be the upper bound of the probability, with a value of 0.9. - This is the lower bound of the probability, taking a value of 0.1. It is a proportional coefficient that adjusts the weight of the two, and its value is 0.5.

[0023] Preferably, the convolution kernel size used in the dual-branch feature extraction layer and the primary capsule layer is 3×3, to replace the 9×9 convolution kernel in the original capsule network, thereby reducing the number of network parameters and computational load.

[0024] Preferably, branch one and branch two independently generate their respective digital capsule prediction results through the graph pooling module of the multi-head attention, and the voting results of the two are merged and averaged at the digital capsule layer to output the final digital capsule vector.

[0025] An image classification method, employing the lightweight bi-branch graph capsule network model described above, includes the following steps: S1: Receive the image to be classified; S2: Multi-scale feature extraction of the image is performed through a dual-branch feature extraction layer; S3: Convert the feature map into a primary capsule vector; S4: The primary capsule vectors are voted on by the graph pooling module with multi-head attention to generate digital capsule vectors; S5: Determine the classification result of the image based on the magnitude of the digital capsule vector.

[0026] An electronic device includes a memory, a processor, and a computer program stored in the memory and executable on the processor, wherein the processor executes the computer program to implement the image classification method as described in claim 9.

[0027] The beneficial effects of this invention are as follows: 1. The lightweight bi-branch graph capsule network model for image classification described in this invention achieves multi-scale feature extraction of images through a bi-branch structure (residual network and residual octave convolution), which can construct higher quality capsules and significantly improve the classification accuracy of the model on complex image datasets (such as CIFAR10 and SVHN), thus significantly improving the quality of image classification.

[0028] 2. The lightweight bi-branch graph capsule network model for image classification described in this invention replaces the dynamic routing algorithm of the original capsule network with a non-iterative multi-head attention graph pooling method, which significantly reduces the amount of computation and parameters and speeds up the training process.

[0029] 3. The lightweight bi-branch graph capsule network model for image classification described in this invention exhibits superior performance on four benchmark datasets: MNIST, FashionMNIST, CIFAR10, and SVHN. In particular, its generalization ability and robustness are significantly better than existing models on complex datasets containing noise and background information.

[0030] 4. The lightweight bi-branch graph capsule network model for image classification described in this invention replaces 9×9 convolution kernels with multiple 3×3 small-sized convolution kernels and combines a bi-branch parallel structure, which improves performance while effectively controlling the number of model parameters and computational complexity. Attached Figure Description

[0031] The present invention will be further described below with reference to the accompanying drawings and embodiments. Figure 1 This is a diagram of the capsule network architecture in this invention; Figure 2 This is the structure of the residual block in this invention. Figure 1 ; Figure 3 This is the structure of the residual block in this invention. Figure 2 ; Figure 4 This is a schematic diagram illustrating the working principle of octave convolution in this invention; Figure 5 This is a diagram of the octave convolution kernel structure in this invention; Figure 6 This is a schematic diagram of residual octave convolution in this invention; Figure 7 This is a diagram of the MHA-GP module in this invention; Figure 8 This is an architecture diagram of the bi-branch graph capsule network in this invention; Figure 9 This is a sample image of the dataset in this invention; Figure 10 This is a graph showing the accuracy of different models tested on four datasets in this invention. Detailed Implementation

[0032] The technical invention of the present invention will now be clearly and completely described with reference to the accompanying drawings. Obviously, the described embodiments are only some embodiments of the present invention, and not all embodiments. The components of the present invention described and shown in the accompanying drawings can be arranged and designed in various different configurations. Therefore, the following detailed description of the embodiments of the present invention provided in the accompanying drawings is not intended to limit the scope of the claimed invention, but merely to illustrate selected embodiments of the invention.

[0033] Among the known existing technologies, capsule networks have made the following key improvements compared to convolutional neural networks: (1) Abandoning pooling layers: Capsule networks no longer use pooling layers to reduce spatial dimensionality. While pooling layers can reduce computation, they may lead to the loss of spatial information. Capsule networks preserve spatial information through pose matrices; (2) Abandoning fully connected layers: Capsule networks do not rely on fully connected layers to build classifiers, but instead use dynamic routing algorithms to process the hierarchical structure of images; (3) Introduction of pose matrix: Capsule networks encode image features as vectors, which enables the network to capture richer spatial information. This helps improve the network's performance in handling spatial relationships and pose transformations; (4) Use of dynamic routing algorithm: Capsule network uses dynamic routing algorithm and voting mechanism to model the whole-part relationship of objects; The working principle of capsule networks is as follows: Figure 1 As shown, it consists of an encoder and a decoder, and is used for MNIST handwritten digit classification. The encoder consists of three parts: a convolutional layer, a primary capsule layer, and a digit capsule layer. The input is a 28×28 image to be classified, and the output is ten 1×16 vectors, where 10 represents a 10-classification task. The class corresponding to the longest vector is the classification result of the capsule network for the input image. The decoder consists of three fully connected layers. Its input is the encoder's output, i.e., a 10×16 vector, and its output is the image reconstructed by the three fully connected layers, with the same size as the original image input to the encoder.

[0034] In this system, the capsule network employs an edge loss function. The encoder outputs the probability for each class, and the edge loss and reconstruction loss are combined as its loss function. The reconstruction loss function is the Euclidean distance between the reconstructed image from the decoder and the original input image. The edge loss function is as follows:

[0035] in, This indicates whether category k exists; the value is 1 if it exists, and 0 otherwise. + Let m be the upper bound of the probability, with a value of 0.9. - This is the lower bound of the probability, taking a value of 0.1. It is a proportional coefficient that adjusts the weight of the two, and its value is 0.5.

[0036] For residual networks, as image classification tasks become increasingly complex and network depth increases, the stochastic gradient descent strategy can no longer find the global optimum, leading to network degradation. To address this issue, existing technologies propose the basic concept and structure of residual networks. Residual networks are mainly composed of residual blocks, which have two structures, such as... Figure 2 As shown, the input of the residual block is X, and the output is... Where X is the identity mapping and F(X) is the residual function to be learned, when the input and output dimensions of the residual block are the same, its residual block structure is as follows: Figure 2 As shown; when the residual block input is m-dimensional and the output is n-dimensional, the dimensions of the input and output should be unified at the jump connection, that is, the dimensions should be adjusted to n-dimensional through a 1×1 convolution operation, as shown. Figure 3 As shown.

[0037] In particular, for octave convolution, the input and output feature maps of the convolutional layer, like natural images, also have high and low frequency components. Existing technologies divide the feature maps after convolution into frequencies, retaining the resolution of high-frequency components and reducing the resolution of low-frequency components containing redundant information to half. This not only reduces the amount of computation and improves the efficiency of convolution, but also expands the receptive field of low-frequency components, thereby enabling the extraction of multi-scale feature information.

[0038] The structure of octave convolution is as follows Figure 4 As shown, , These represent the high-frequency and low-frequency feature maps of the octave convolution input, respectively. , The high-frequency and low-frequency feature maps output by the octave convolution are represented, respectively. The length and width of the low-frequency feature map in the same layer are both half that of the high-frequency feature map. Because the high-frequency and low-frequency feature maps have different sizes, traditional convolution operations cannot directly fuse them. Therefore, an octave convolution that allows the fusion of high-frequency and low-frequency feature maps is proposed. A complete octave convolution operation requires four convolution operations: two convolutions of the same frequency and two convolutions across frequencies. Finally, the high-frequency feature map is output. for and Similarly, the low-frequency feature map is output as a result of the merging process. for and The result of the merger.

[0039] The structure of an octave convolution kernel is as follows: Figure 5 As shown, the convolution kernel An octave convolution kernel consists of four parts, namely... , , and In the convolution between features of the same frequency, the input and output feature maps have the same size. and For traditional convolution operations, First, a convolution operation is performed, followed by a 2x upsampling, thus achieving cross-frequency information exchange from low to high frequencies. Similarly, for... First, average pooling is performed to achieve a 2x downsampling, followed by convolution to achieve cross-frequency information exchange from high to low frequencies. , These represent the ratios of the low-frequency feature maps of the input and output, respectively.

[0040] Based on this, this invention proposes a lightweight bi-branch graph capsule network model for image classification. The capsule network is improved from two aspects: the feature extraction layer and the dynamic routing algorithm. The feature extraction layer adopts a bi-branch structure, and the graph pooling method with multi-head attention is used between the primary capsule layer and the digital capsule layer, as detailed below.

[0041] First, for the dual-branch structure, a lightweight parallel dual-branch feature extraction layer is used to extract features from the image at multiple scales, thereby constructing high-quality capsules and improving the robustness of the network. Branch one employs a residual network with a convolutional kernel size of 3×3, such as... Figure 2-3 As shown, stacking smaller convolutional kernels and introducing residual skip connections not only improves the model's nonlinear expressiveness but also prevents network degradation. Branch two uses a residual octave convolutional module with a kernel size of 3×3, as shown... Figure 6 As shown, feature extraction is performed. By using cross-frequency convolution, not only is the problem of feature information redundancy effectively overcome, but multi-scale feature information is also extracted. Each branch contains a feature extraction layer, a primary capsule layer, and a high-level capsule layer.

[0042] The residual octave convolution module comprises three structures: Initial Octave Convolution, Middle Octave Convolution, and Final Octave Convolution. The input and output of the Middle Octave Convolution layer both contain high-frequency feature maps. and low-frequency feature maps The scaling factor of the feature map It also includes a residual skip connection structure, that is, for the input of the middle octave convolutional layers... , Its output , As shown in the following formula:

[0043] Among them, the output of the high-frequency feature map It consists of three parts, namely from the high-frequency input. Feature maps obtained after convolution operations, low-frequency input Feature maps obtained after convolution and upsampling operations, and identity mappings Output of low-frequency feature map It also consists of three parts, namely from the low-frequency input. Feature maps obtained after convolution operations, high-frequency inputs Feature maps obtained after pooling and convolution operations, and identity mappings The length and width of the high-frequency feature map are both twice that of the low-frequency feature map in the same layer.

[0044] The input to the initial octave convolutional layer contains only high-frequency feature maps, and the output is the same as that of the middle octave convolutional layers. , The input to the terminal octave convolutional layer is the same as that of the previous layer, but the output only contains high-frequency feature maps. ,Right now , .

[0045] In existing technologies, capsule networks have been studied. Through multiple ablation experiments, it has been shown that dynamic routing algorithms are not necessarily more robust than convolutional neural networks. Therefore, existing technologies have been improved to propose an interpretable object recognition model, namely graph capsule networks.

[0046] Graph Capsule Networks introduce a multi-head attention graph pooling (MHAGP) module between the primary capsule layer and the digital capsule layer, such as... Figure 7 As shown, the extracted feature maps are divided into L groups, and then transformed into L graphs using a transformation matrix W1. Each graph represents a head, and each head contains K×K nodes. Each node in each head is a primary capsule. The spatial importance of the nodes in each head is labeled using a multi-head attention (MHA) module. Then, another transformation matrix W2 is used to generate the votes of each head for M higher-level capsules in the digital capsule layer. Finally, the average of the votes from the L digital capsule layers is calculated, which is the graph pooling method, to obtain the votes of the primary capsules for the digital capsules.

[0047] Graph pooling methods for multi-head attention, such as Figure 7 As shown, L sets of feature maps are formed after convolution in the primary capsule layer, and the size of each set of feature maps is K×K×1×D. in， D in Let P represent the number of channels in each feature map group, and K represent the width and height of the feature map. Assume the i-th feature map group in the primary capsule layer is P. i After transformation matrix W 1i Feature map P i Modeled as a head X i As shown in the formula below: in D inD represents the dimension of the capsules in the primary capsule layer. out For each image, the dimension of the capsules in the digital capsule layer is [dimension]. It consists of K×K nodes, each node representing a capsule. , where i represents the i-th head, j represents the j-th capsule on the i-th head, and the values ​​of each dimension of the same capsule are the values ​​of feature maps from different channels but at the same position.

[0048] The MHA module is an attention matrix A, which models the spatial relationships of the graph constructed by the primary capsules. Specifically, for a given K... 2 For a graph with n nodes, the attention matrix is ​​calculated as shown in the following formula: in a and b represent the index values ​​of the nodes in the graph, P m P n This represents the row and column coordinates corresponding to the two nodes with index values ​​a and b in the graph. That is, P m P n From (1,1) to (K, K), This is a hyperparameter.

[0049] For the constructed L heads, after initially recalibrating the importance of each capsule using the attention matrix A, the attention coefficient matrix of each head is then calculated using the transformation matrix W2 and the normalization function softmax, as shown in the following formula:

[0050] in , are trainable parameters, and M is the number of capsules in the digit capsule layer, i.e., the number of categories in the dataset to be classified. , indicating L heads with K 2 The importance of each primary capsule to the M digit capsules was determined by recalibrating the importance of each capsule. Some primary capsules that represent background information and other factors that affect classification performance had attention coefficients that were almost zero, indicating that these primary capsules did not contribute to the classification results. The prediction result S of the i-th head pair for the final digital capsule layer M advanced capsules i It is calculated using the following formula: in That is, each head generates M dimensions with dimension D. out The digital capsules, and the final output of the digital capsule layer are L head pairs, M dimensions, and D. outThe average of the voting results for the advanced capsules is used, and the M advanced capsules are compressed using the squash function, i.e., graph pooling, as shown in the following formula.

[0051] Among them, The squash function compresses M high-level capsules in the digital capsule layer into a vector V with the same direction but a length of 0-1.

[0052] Furthermore, based on the above embodiments, the architecture of the lightweight dual-branch capsule network encoder proposed in this invention is as follows: Figure 8 The diagram shows the architecture of a dual-branch capsule network applied to the classification of the MNIST dataset. It uses two parallel branches to perform feature extraction and dynamic routing, and finally aggregates them at the digital capsule layer to achieve classification.

[0053] like Figure 8 The encoder of the dual-branch capsule network shown consists of five parts: an input layer, a feature extraction layer, a primary capsule layer, a digital capsule layer, and an output layer. The input image first passes through two Conv-BN-ReLU layers for feature extraction, with a convolution kernel size of 3×3 and a stride of 1. Then, the resulting 28×28×32 feature maps are used as inputs to the two branches for multi-scale feature extraction of the image.

[0054] The feature extraction layer of branch one is a residual network consisting of 5 sets of Conv-BN-ReLU layers, with a convolution kernel size of 3×3 and strides of 2, 1, 1, 2 and 1 respectively. The feature extraction layer of branch two consists of an octave convolution and a residual network, containing a total of 10 Conv-BN-ReLU layers: one Initial Octave Conv-BN-ReLU layer, five Middle Octave Conv-BN-ReLU layers, one Depthwise Conv-BN-ReLU layer, one Middle Octave Conv-BN-ReLU layer, one Final Octave Conv-BN-ReLU layer, and one Conv-BN-Squash layer. The kernel size is 3×3, and the strides are 1, 1, 1, 2, 1, 1, 1, 2, 1, and 1, respectively. The value of α in the octave convolution is 0.125. Each branch forms 16 graphs, each containing 36 nodes, meaning each branch forms 576 primary capsules. The primary capsules have a dimension of 1×8. The primary capsules of each branch vote on the higher-level capsules using a multi-head attention graph pooling method. Each graph's 36 nodes generate a vote for a digital capsule, which is then used to form 10 1×16 digital capsules. The digital capsules from the two branches are then aggregated and averaged to obtain 10 1×16 digital prediction capsules. Finally, the modulus of these 10 digital capsules is calculated using Euclidean distance. The category corresponding to the digital capsule with the longer modulus is the classification prediction result for the image.

[0055] The decoder of the dual-branch capsule network is a neural network consisting of three fully connected layers. The number of neurons in the three fully connected layers are 512, 1024, and the number of pixels in the input image, respectively. Its input is a 10×16 digital capsule matrix output from the encoder, and its output is a reconstructed image of the same size as the input.

[0056] Furthermore, in order to verify the actual technical effect of the present invention, the following experiments were conducted: First, regarding the experimental datasets, the datasets used in this invention are commonly used internationally available image classification datasets in the field of computer vision, namely MNIST, FashionMNIST, CIFAR10, and SVH, such as... Figure 9 As shown, these four datasets all consist of images and corresponding labels, each containing both a training set and a test set, and all belong to a 10-class classification task.

[0057] MNIST and FashionMNIST are both grayscale images of size 28×28×1. MNIST contains handwritten digit images from 0 to 9, while FashionMNIST is a 10-class grayscale image dataset of clothing, with a higher classification task complexity than MNIST. CIFAR10 and SVHN are both color images of size 32×32×3. CIFAR10 contains 10-class real-world objects, while SVHN is Google's Street View house number dataset. This invention uses a 10-class dataset after cropping the original images, i.e., Format2. Some samples in this dataset contain multiple digits, and their corresponding labels are the digits located in the center of the image.

[0058] Secondly, regarding the basic experimental configuration, all experiments were conducted under the Linux operating system, with an Intel(R) Core(TM) i9-9880H CPU and an NVIDIA RTX3070 6G GPU, and 32GB of RAM and a 1TB solid-state drive. All experiments were based on the Python 3.7.6 programming language and used the Tensorflow 2.10.0 deep learning framework.

[0059] Third, regarding experimental parameters, the relevant parameters of this invention are set as follows: batch size is set to 50, epochs are set to 60, learning rate (lr) is 0.0001, learning rate decay exponent (lr_decay) is 0.95, the Adam optimizer is used, and the parameters of the marginal loss function are... Set to 0.5, m - With m + The values ​​are set to 0.1 and 0.9 respectively. In the loss function, the coefficient of the reconstruction loss function is 0.0005 for a single pixel. That is, for the MNIST and FashionMNIST datasets with a size of 28×28×1, the coefficient of the reconstruction loss function is... The coefficient is 0.392, while for the CIFAR10 and SVHN datasets of size 32×32×3, the coefficient of its reconstruction loss function is... It is 1.536.

[0060] Furthermore, the ablation experiment was mainly divided into two parts. The first part verified the performance of the capsule network by using three combinations of the two branches for feature extraction, as well as the advantages and disadvantages of various evaluation metrics such as network parameter quantity, number of floating point operations (FLOPs), and training loss value. The second part compared and verified the classification performance and other evaluation metrics of the capsule network using the graph pooling method with multi-head attention and the capsule network using the dynamic routing algorithm.

[0061] Table 1 shows the feature extraction layer of the dual-branch graph capsule network, which uses three combinations of two branches for feature extraction, thus achieving classification results on four datasets: FashionMNIST, MNIST, CIFAR10, and SVHN.

[0062] Table 1

[0063] As shown in Table 1, the test accuracy of using branch 2 alone is better than that of branch 1, but its parameter count and FLOPs are also higher. Using the parallel structure of branch 1 and branch 2, the test accuracy is higher than that of using either branch alone. Furthermore, the parameter count, FLOPs, and training time of the combined structure of branch 1 and branch 2 are all better than those of using either branch alone. The constructed model is more robust and has stronger generalization ability. In particular, on the more complex CIFAR10 and SVHN datasets with complex backgrounds and noise, its classification accuracy is significantly improved.

[0064] Table 2 shows the comparison results of graph pooling algorithms that incorporate fully connected classifiers, dynamic routing algorithms, and multi-head attention mechanisms. Here, FC indicates that the classifier adopts a fully connected layer structure, DR indicates the dynamic routing algorithm, and MHAGP indicates the graph pooling algorithm with multi-head attention mechanism.

[0065] Table 2

[0066] The two ablation experiments demonstrate that the parallel bi-branch feature extraction layer, consisting of a residual network in branch one and a residual octave convolution module in branch two, increases network width while enabling multi-scale feature extraction from images. This results in stronger image feature representation capabilities and a more robust model, achieving good classification performance on four existing datasets with fewer parameters and faster training speed. Secondly, the introduction of a residual structure in the bi-branch feature extraction layer mitigates the risk of a sharp drop in test accuracy caused by excessive network depth. Finally, a graph pooling algorithm with multi-head attention mechanism replaces the original capsule network's iterative dynamic routing algorithm between the primary capsule layer and the digital capsule layer, improving classification accuracy while reducing the number of parameters and increasing training speed.

[0067] Furthermore, the present invention also conducted comparative experiments. Table 3 shows the comparative experimental results of the proposed DB-GraCapsNet with the original capsule network, MLCN, and GraCapsNet on four datasets. Table 4 is a comparison table of the accuracy of different improved models on four datasets. As shown in Table 3, the proposed DB-GraCapsNet outperforms CapsNet, MLCN, and GraCapsNet in multiple evaluation metrics, including parameter count, FLOPs, and test accuracy. DB-GraCapsNet has only 28.93% of the parameters of CapsNet on the MNIST and FashionMNIST datasets, and only 41.46% of the parameters on the CIFAR10 and SVHN datasets, with test accuracies 0.14%, 4.6%, 13.32%, and 3.37% higher than CapsNet, respectively. This indicates that the parallel dual-branch structure of the proposed model can extract rich feature information from images, thereby constructing capsules with stronger feature representation capabilities. Furthermore, the use of a non-iterative multi-head attention graph pooling algorithm allows the model to achieve better classification accuracy and generalization performance with fewer parameters.

[0068] Figure 10 The figure shows the test accuracy curves of DB-GraCapsNet, the original capsule network, MLCN, and GraCapsNet on four datasets: MNIST, FashionMNIST, CIFAR10, and SVHN. The horizontal axis represents the training epoch, and the vertical axis represents the test accuracy.

[0069] Table 3

[0070] Depend on Figure 10 As shown in the test accuracy curves, the model proposed in this invention achieves higher test accuracy on all four datasets than the other three models. Particularly on the relatively complex FashionMNIST and CIFAR10 datasets, its test accuracy is 4.6% and 13.32% higher than the original capsule network, respectively; 1.54% and 9.59% higher than MLCN, respectively; and 1.27% and 3.56% higher than Graph CapsNet, respectively. Furthermore, the curves show that the network proposed in this invention converges quickly, requiring fewer iterations, indicating strong learning ability and generalization performance. Table 4 shows the comparative experimental results of the classification performance of the proposed model with several existing improved capsule network models on the four datasets.

[0071] As shown in Table 4, the experimental results demonstrate that the proposed model exhibits good robustness, outperforming previously proposed improved capsule network models in classification accuracy across all four datasets. Particularly noteworthy is the significant improvement in classification accuracy on the relatively complex CIFAR10 and SVHN datasets. This indicates that the proposed bi-branch graph capsule network model possesses strong learning ability and generalization capacity.

[0072] Table 4

[0073] Therefore, this invention proposes a lightweight bi-branch graph capsule network model for image classification. First, the single branch of the capsule network is optimized into a bi-branch structure: branch one uses a residual network, and branch two uses residual octave convolution. This achieves multi-scale feature extraction of images while preventing network degradation due to excessive depth. Second, multiple smaller 3×3 convolutional kernels replace the original 9×9 convolutional kernels of the capsule network, reducing the number of network parameters and computational cost while improving the model's non-linear expressive power. Finally, the dynamic routing algorithm is improved: the two branches iteratively update and independently generate their own high-level capsules using a multi-head attention graph pooling method. The high-level capsules are then aggregated at the digital capsule layer. Multiple ablation experiments and comparative experiments show that the proposed model achieves high classification accuracy on datasets MNIST, FashionMNIST, CIFAR10, and SVHN. Particularly on the more complex CIFAR10 and SVHN datasets, the classification performance is significantly improved, effectively enhancing image classification quality.

[0074] The foregoing has shown and described the basic principles, main features, and advantages of the present invention. Those skilled in the art should understand that the present invention is not limited to the above embodiments. The embodiments and descriptions in the specification are merely illustrative of the principles of the present invention. Various changes and modifications can be made to the present invention without departing from its spirit and scope. All such changes and modifications fall within the scope of the present invention as claimed, which is defined by the appended claims and their equivalents.

Claims

1. A lightweight dual-branch graph capsule network model applied to image classification, characterized in that, The method comprises the following steps: An input layer is used to receive an image to be classified; A double-branch feature extraction layer is connected with the input layer and is composed of two parallel branches, and is used to perform multi-scale feature extraction on the image to generate a feature map; A primary capsule layer is connected with the double-branch feature extraction layer, and is used to convert the feature map into a plurality of primary capsule vectors; A multi-head attention graph pooling module is connected with the primary capsule layer, and is used to perform consistency voting on the primary capsule vectors to generate a digital capsule vector; An output layer is connected with the multi-head attention graph pooling module, and is used to determine the classification result of the image according to the module length of the digital capsule vector.

2. The lightweight dual-branch graph capsule network model of claim 1, wherein, The double-branch feature extraction layer comprises: Branch one adopts a residual network with a skip connection to extract deep features of the image; 3. The lightweight dual-branch graph capsule network model of claim 2, wherein, The residual octuple convolution module comprises an initial octuple convolution, an intermediate octuple convolution, and a terminal octuple convolution; wherein the input of the intermediate octuple convolution comprises high-frequency feature maps and low-frequency feature maps , the output and is calculated by the following formula: where f denotes a convolution operation, W denotes a kernel weight of the convolution, upsample denotes an up-sampling operation, and pool denotes a pooling operation.

4. The lightweight dual-branch graph capsule network model of claim 1, wherein, Branch two adopts a residual octave convolution module to extract multi-scale features of the image by high-frequency and low-frequency mutual fusion. The feature maps of the primary capsule layer are divided into L groups, and each group has P feature maps i By the conversion matrix W 1i The graph head X is constructed as i The following formula is used: ， wherein, wherein , D in is the dimension of the capsules in the primary capsule layer, D out is the dimension of the capsules in the digital capsule layer, for each figure , which consists of K x K nodes, each node represents a capsule where i represents the i-th head, j represents the j-th capsule on the i-th head, and the values of the various dimensions of the same capsule are all from the feature maps of different channels but the same position; The attention matrix A for each graph is computed, whose elements A ab is computed by the following equation: ， wherein, , a, b represent the index values of the nodes in the graph, P m , P n represent the coordinates of the row and column corresponding to the two nodes with index values a, b in the graph, , P m , P n from (1, 1) to (K, K), is a hyperparameter; The multi-head attention graph pooling module performs the following steps: ， wherein, are trainable parameters, M is the number of digit capsule layers, i.e., the number of classes of the dataset to be classified, represents the importance of the L heads on the K 2 primary capsules to the M digit capsules. An attention matrix A and a conversion matrix W2 are used to calculate the attention coefficient matrix of each head, and the following formula is used: ， wherein, M digital capsules of dimension D are generated for each head; out of dimension D are generated for each head; The following formula is used to calculate the prediction result of the digital capsule for each head: ， wherein, wherein .

5. The lightweight dual-branch graph capsule network model of claim 4, wherein, The prediction results of the L heads are averaged, and are compressed through a squash function to obtain the final digital capsule vector ， The squash function is defined as:

6. The lightweight dual-branch graph capsule network model of claim 1, wherein, The model adopts a combination of an edge loss function and a reconstruction loss function as a total loss function; wherein the edge loss function L k is defined as: ， wherein, denotes whether class k exists or not, with value 1 if it exists and 0 otherwise, m + is the upper probability bound, with value 0.9, m - is the lower probability bound, with value 0.1, is the proportionality coefficient, which adjusts the weight of both, with value 0.

5.

7. The lightweight dual-branch graph capsule network model of claim 1, wherein, The squash function is used to compress the M high-level capsules in the digital capsule layer into a vector V which is direction-invariant and has a length of 0-1.

8. The lightweight dual-branch graph capsule network model of claim 1, wherein, The convolution kernel size used in the double-branch feature extraction layer and the primary capsule layer is 3*3, which replaces the 9*9 convolution kernel in the original capsule network, thereby reducing the network parameter quantity and the calculation amount.

9. An image classification method characterized by, The branch one and the branch two independently generate respective digital capsule prediction results through the multi-head attention graph pooling module, and the voting results of the two are combined and averaged in the digital capsule layer to output the final digital capsule vector. The method comprises the following steps: S1: receiving an image to be classified; S2: performing multi-scale feature extraction on the image through a double-branch feature extraction layer; S3: converting the feature map into a primary capsule vector; S4: performing voting on the primary capsule vector through a multi-head attention graph pooling module to generate a digital capsule vector; 10. An electronic device comprising a memory, a processor, and a computer program stored on the memory and executable on the processor, characterized in that, S5: determining the classification result of the image according to the module length of the digital capsule vector. The processor executes the computer program to realize the image classification method of claim 9. The processor executes the computer program to realize the image classification method of claim 9.

Images