Domain independent post-processing with graph U-nets: applications to electrical impedance tomographic imaging^⋆

As the name suggests, graph networks act on graph input data. Historically, graph convolutional networks have been applied to citation networks, social networks, or knowledge graphs. However, whenever you have a computational mesh, you inherently have a graph representing the connections between the elements or nodes in your computational mesh. In particular, for irregular meshes commonly associated with FEM there are two natural options for the associated graph: the mesh elements, or the mesh nodes (figure 2). The adjacency matrix ${ \mathcal A }$ is a sparse matrix listing the edges (connections) between the graph nodes. The basic form of the sparse adjacency matrix is ${{ \mathcal A }}_{{ij}}=1$ if there is an edge connected nodes n_i and n_j . Self-loops are recorded separately.

Zoom In Zoom Out Reset image size

As opposed to convolutional layers which require regular n-dimensional regular input data, graph convolutional layers are designed to work on simple, homogeneous graph-type data. While the list of graph convolutional layers is ever expanding, we used the layer proposed by Kipf and Welling (2016) designed to be analogous to the CNN setting:

$$\begin{eqnarray}&&{H}^{\left(i+1\right)}={\tilde{D}}^{-\tfrac{1}{2}}\left({ \mathcal A }+I\right){\tilde{D}}^{-\tfrac{1}{2}}{H}^{\left(i\right)}{W}^{(i)},\end{eqnarray} \tag{ 1 }$$

where the inputs to the layer are the graph's feature matrix ${H}^{\left(i\right)}\in {{\mathbb{R}}}^{N\,\times \,f\left(i\right)}$ and adjacency matrix ${ \mathcal A }\in {{\mathbb{R}}}^{N\times N}$, and the output is a new feature matrix ${H}^{\left(i+1\right)}\in {{\mathbb{R}}}^{N\,\times \,f\left(i+1\right)}$ for the graph with the same structure (Kipf and Welling 2016). Self loops, or edges between a graph node and itself, are represented by the identity matrix I, of the same size as ${ \mathcal A }$, and are added to the adjacency matrix. Then, that sum is multiplied on the left and right by ${\tilde{D}}^{-\tfrac{1}{2}}$ to account for the number of edges each node has. The diagonal matrix $\tilde{D}$, which does not contain trainable parameters and is only determined by ${ \mathcal A }$, is defined by ${\tilde{D}}_{{ii}}=1+{\sum }_{j}{{ \mathcal A }}_{{ij}}$. Finally, multiplication of the scaled adjacency matrix ${\tilde{D}}^{-\tfrac{1}{2}}\left({ \mathcal A }+I\right){\tilde{D}}^{-\tfrac{1}{2}}$ by the input feature matrix H⁽ⁱ⁾ aggregates information within local neighborhoods and multiplication by the weight matrix $W\in {{\mathbb{R}}}^{{f}^{\left(i\right)}\times {f}^{\left(i+1\right)}}$ takes linear combinations of the aggregated features to form the output features.

Bias parameters $b\in {{\mathbb{R}}}^{{f}^{\left(i+1\right)}}$ can be included in a graph convolution by adding them to the output feature vector of each node (each row of the output feature matrix). The weight matrix W (and optional bias vector) are the trainable parameters that are optimized during training. One significant difference between convolutional layers and graph convolutional layers is in how they aggregate information within a pixel or node's neighborhood. Convolutional layers learn linear aggregation functions via the kernel parameters while graph convolutions aggregate information according to a fixed linear function, ${\tilde{D}}^{-\tfrac{1}{2}}\left({ \mathcal A }+I\right){\tilde{D}}^{-\tfrac{1}{2}}$, determined by the graph's adjacency matrix ${ \mathcal A }$. The non-learned aggregation function of such a graph convolution has raised questions of the learning capacity of graph convolutional layers (Hoang and Maehara 2019). Despite those concerns, graph convolutional layers have been used successfully for a variety of graph and node classification tasks (Kipf and Welling 2016, Zhang et al 2019). GCNs have also been used for model-based learning directly on irregular mesh data (Herzberg et al 2021). In the down-sampling path, graph convolutional layers are used like in the original graph U-net (Gao and Ji 2019) and in previous work on images represented as graphs (Herzberg et al 2021).

After graph convolutions, pooling layers are used to move down and up the U-net. Several node-dropping, hierarchical pooling layers (layers that gradually coarsen a graph by removing nodes) for GNNs have been proposed including self-attention graph pooling (Lee et al 2019), adaptive structure aware pooling (Ranjan et al 2019), and gPool (Gao and Ji 2019). The gPool layer selects graph nodes to preserve by first learning a projection of nodes' feature vectors. For each node, its feature vector is projected by a vector $p\in {{\mathbb{R}}}^{\left(i\right)}$ to return a scalar score. Then, the nodes with the top projection scores are preserved. The adjacency matrix is also sliced to preserve only the rows and columns for the preserved nodes. The parameters in the projection vector p are optimized during training.

One significant difference between the gPool layer and the max-pooling layer used in CNNs is that the gPool layer performs global node selection while the max-pooling layer performs local selection. That is, the max-pooling layer only considers features within a subregion or window of the input when preserving pixel data, while the gPool layer considers the projection scores from the entire graph when deciding which nodes to preserve. Self-attention graph pooling, adaptive structure aware pooling, and other node-dropping, hierarchical pooling layers also consider the whole graph when selecting which nodes to preserve, which gives way to the possibility that entire regions of a graph could be discarded when pooling (Liu et al 2022); something that is not possible with CNN-based max-pooling. Therefore a new graph pooling layer, the k-means cluster (kMC) max pool layer, was developed with local pooling in mind.

To create local windows in the input graph, the k-means++ algorithm ⁴ (Arthur and Vassilvitskii 2007) was used to cluster the graph nodes. The inputs to the k-means++ algorithm include the locations of the N graph nodes and the number of clusters N_c desired. The output of the algorithm was the cluster assignments $c\in {\left\{1,\ldots ,{N}_{c}\right\}}^{N}$ and the locations of the clusters. Locations of the input nodes are determined from the FE mesh that the input graph was representing. That is, for data on the mesh nodes, the locations of those nodes are used, and, for data on the elements, the element centroid locations can be used. The locations of the output clusters are taken as the centroid of the graph nodes in the cluster. As the k-means++ algorithm is stochastic, it was repeated multiple times, and the cluster assignments with the minimum within-cluster spacing selected.

With clusters defined, max-pooling is performed within each cluster along each feature of the input graph. Therefore, the output of the pooling operation is a set of N_c nodes at the cluster locations and with the maximal features from the input nodes within each cluster. The adjacency matrix of the output graph is determined by the cluster assignments as well. If any edge connected two input nodes assigned to separate clusters, an edge was drawn between the output nodes representing those clusters. Figure 3 (top) depicts the structure of the input and output graphs of the kMC max pooling operation.

Zoom In Zoom Out Reset image size

In encoder-decoder type GNNs that have a down-sampling path of node-dropping, hierarchical pooling layers and a symmetrical up-sampling path, the up-sampling often uses the gUnpooling layer (Gao and Ji 2019, Liu et al 2022). Nodes and edges that were removed in the associated pooling layer are restored in the gUnpooling layer. The features of the restored nodes are set to zero. There are no trainable parameters in the gUnpool layer. The gUnpool layers could not be used here as pairing it with the k-means cluster max pool layer, would result in the output graph having all nodes' features set to zero because all output nodes are restored. Like the gUnpool layer, a new clone cluster unpool layer, shown in figure 3 (bottom) was designed to restore the pre-pooled graph structure. Instead of returning the previously removed nodes with features set to 0, the nodes are restored with feature vectors equal to the node representing the cluster to which the nodes belong. That is, each output node is restored as a clone/copy of the input node representing the cluster including the output node. The proposed layer was developed to act like a nearest-neighbor up-sampling layer, or a transpose convolutional layer with equal stride and kernel size and parameters fixed to 1.

For each graph, the adjacency matrix and downsampled clustering assignments can be computed offline prior to the training or processing through the network. The resulting graph U-net Λ_θ takes as inputs $\left({x}_{{\rm{in}}},{ \mathcal A },{\bf{c}}\right)$ where ${\bf{c}}=\left({c}_{1},\ldots ,{c}_{{N}_{p}}\right)$ are the cluster assignments for the N_p pooling (and unpooling) layers in the network. The post-processed image $\hat{x}$ is then the output of

$$\begin{eqnarray}&&\hat{x}={{\rm{\Lambda }}}_{\theta }\left({x}_{\mathrm{in}},{ \mathcal A },{\bf{c}}\right).\end{eqnarray} \tag{ 2 }$$

Note that each cluster assignment vector, c_j , is of a different length, and c represents the list of vectors. Figure 1 provides a diagram of the proposed graph U-net with N_p =3 pooling layers. The k-means cluster max pooling and clone cluster unpooling layers used in this graph U-net allow the cluster assignments to be mesh specific and computed prior to training or predicting. In the experiments conducted, no problems were noticed by training or prediction samples having different meshes and cluster assignments. Computing cluster assignments can take several minutes depending on the number of input nodes, clusters, and repetitions, thus computing them ahead of time keeps both training and prediction fast. In addition, the number of clusters used at each pooling layer can be tuned similarly to how the kernel size of a CNN max pooling layer can be adjusted. Overall, the new proposed graph pooling layer was a fast and flexible alternative to the existing graph pooling layers, and it behaves more like the traditional max pooling layer used in convolutional U-nets. The proposed clone cluster unpooling layer works naturally with the cluster-based pooling in the graph U-net architecture. Alternatively, regional downsampling of the graph and adjacency matrix could be performed via mesh coarsening as in Suk et al (2023).

Training data consists of $\left\{{{\rm{x}}}_{\text{}\mathrm{recon}}^{{\rm{i}}},{x}_{\text{}{true}}^{i}\right\}$ pairs where ${{\rm{x}}}_{\text{}\mathrm{recon}}^{{\rm{i}}}$ is produced via the user's chosen solution method on measurement data yⁱ . The computational mesh associated to ${{\rm{x}}}_{\text{}\mathrm{recon}}^{{\rm{i}}}$ leads to an adjacency matrix ${{ \mathcal A }}^{i}$ and corresponding cluster assignments cⁱ . Then, the network is trained using ${{\rm{x}}}_{\text{}\mathrm{recon}}^{{\rm{i}}}$, Aⁱ , and cⁱ as inputs and ${x}_{\text{}{true}}^{i}$ as the desired outputs. A loss function involving the network outputs ${\hat{x}}^{i}$ and the known true images ${x}_{\text{}{true}}^{i}$ can be used. Either MSE or an ℓ¹ loss function are natural choices. One could also weight the loss function differently for each sample or for each graph node. Once the training of the network is complete, the parameters $\hat{\theta }$ are saved for use in the online prediction stage.

Note that training data could have all different computational meshes, all the same, or any mix thereof. In this work, a single reconstruction mesh was used per network trained to demonstrate a simple case. Note that this does not restrict later passing reconstructions on different meshes through the trained network (i.e. a network trained on reconstructions from mesh^j , adjacency matrix ${{ \mathcal A }}^{j}$, with clusters c^j is not restricted only to mesh^j as the adjacency matrix and clusters are inputs to the network). Several test cases are presented in section 3 demonstrating the flexibility of the networks to input data coming from different domain shapes, experimental setups, and even higher dimensional data (3D when trained on 2D).

In the online prediction stage, an updated reconstruction $\hat{x}$ is estimated from measurement data y by first computing an initial reconstruction x_recon from y and passing x_recon, its adjacency matrix ${ \mathcal A }$, and corresponding cluster assignments c through the trained network using (2) where the trained parameters $\hat{\theta }$, which were saved during the offline stage, are used in the network.

An alternative to the graph U-net presented in this work was to interpolate image data defined on a FE mesh to a pixel grid so that a typical convolutional U-net can be used. Thus, the reconstruction x_in was computed, as before, on a FE mesh from measurement data y, then interpolated to a pixel grid ${\tilde{x}}_{\mathrm{in}}=f\left({{\rm{x}}}_{\text{}\mathrm{in}}\right)$. Next, a post-processing CNN ${{\rm{\Lambda }}}_{\theta }^{{\rm{CNN}}}$ is used to estimate the reconstruction on the pixel grid: $\tilde{x}={{\rm{\Lambda }}}_{\theta }\left({\tilde{x}}_{\mathrm{in}}\right)$. If desired, the network output can then be interpolated back to the FE computational mesh: $\hat{x}={f}^{\dagger }\left(\tilde{x}\right).$

In this setting, the CNN networks are trained using a set of initial reconstruction and truth image pairs, where both have been interpolated from the computational mesh to the pixel grid. It may be desirable to utilize a CNN to post-process images, as opposed to a GNN, as there is precedent for using CNNs on image data. Still, the loss of fine detail in refined portions of the mesh and the errors induced by interpolating to and from a pixel grid could be prohibitive or time intensive in certain imaging applications.

Six experimental tank datasets with conductive agar targets, taken on three different EIT machines will be used to evaluate the graph U-net method presented in section 2.1.1. The first set, denoted KIT4, ⁵ comes from the 16 electrode KIT4 system at the University of Eastern Finland (Kourunen et al 2008, Hamilton et al 2019). Conductive agar targets (pink 0.323 S m⁻¹, white 0.061 S m⁻¹) were placed in a saline bath, with measured conductivity 0.135 S m⁻¹, in a chest shaped tank with perimeter 1.02 m and 16 evenly spaced electrodes of width 20 mm. The computational reconstruction mesh used for the KIT4 datasets contained 3984 elements. Three experiments were performed to mimic heart (pink) and lung (white) imaging (figure 4-left). Sample KIT4-S.1 shows a 'healthy' setup with two low conductivity targets (lungs) and one high conductivity target (heart). Sample KIT4-S.2 has a cut in the 'lung' target on the viewer's left while KIT4-S3 replaces the missing portion with a more conductive agar (e.g. possibly a pleural effusion). For each setup, 16 adjacent current patterns were applied with an amplitude of 3 mA and current frequency of 10 kHz, and measurements were recorded on all electrodes. The regularization parameters for TV were selected as λ = 5 × 10⁻⁵ and γ = 10⁻¹⁴ based on testing from simulated datasets.

Zoom In Zoom Out Reset image size

Next, archival data from the 32 electrode ACT3 system (Cook et al 1994, Isaacson et al 2004) was used. In Sample ACT3-S.1 ((figure 4-right), a circular heart target (0.750 S m⁻¹) and two lung targets (0.240 S m⁻¹) were placed in a saline bath (0.424 S m⁻¹) in a tank of radius 0.15 m. Trigonometric current patterns with maximum amplitude 0.2 mA and frequency 28.8 kHz were applied on the 32 equally spaced electrodes of width 25 mm.

Lastly, data from the ACT5 system (Rajabi Shishvan et al 2021), was used for extension testing to 3D data. Samples ACT5-S.1 and ACT5-S.2 (figure 5) were collected on plexiglas box with interior dimensions 17.0 cm x 25.5cm x 17.0 cm with 32 electrodes of size 8 cm x 8 cm (Hamilton et al 2022). The top of the tank is removable and has small holes allowing for filling and resealing between experiments. Spherical agar targets of measured conductivity 0.290 S m⁻¹ were placed in tap water measuring 0.024 S m⁻¹. Optimal current patterns for the saline only tank were obtained and used for ACT5-S.1 and S.2. ⁶

Zoom In Zoom Out Reset image size

Simulated data using the chest-shaped domain corresponding the the KIT4 data, with 16 electrodes and adjacent current pattern injection, was used to generate the training data. The measurement mesh contained 10 274 elements, and the reconstruction mesh was the same as the one used for the experimental KIT4 datasets (3984 elements). The simulated true conductivity distributions had an equal chance of either three or four randomly placed elliptical targets that were not allowed to overlap or touch the boundary. For each ellipse, the major axis was between [0.03 and 0.07] meters and the minor axis was [50–90]% of the value of the major axis. Each target had an equal chance of having a constant conductivity in the range of $\left[0.04,0.07\right]$ S m⁻¹ or $\left[0.25,0.35\right]$ S m⁻¹, while the background conductivity values were constant in the range of $\left[0.11,0.17\right]$ S m⁻¹. The measured voltage data was simulated using all 16 possible adjacent current patterns with an amplitude of 3 mA and 0.5% relative noise added to the voltages prior to reconstruction with TV. This corresponds to the simulated data having an SNR of about 51 dB, well below the SNR of the experimental data (KIT4 65.52 dB (Nissinen et al 2011), and ACT3/ACT5 96 dB).

Eight separate U-nets were trained, four based on GNNs and four on CNNs. The four graph U-nets are named GNN-TVx and the four convolutional U-nets are named CNN-TVx where the 'x' represents the TV method iterate used as input to the network. That is GNN-TV2 and CNN-TV2 are a graph U-net and classic U-net, respectively, that use the second iteration σ₂ of the TV method as input. The networks were trained using 5000 training samples and 500 validation samples. The number of training samples used was greater than the number used for the model-based methods in Herzberg et al (2021) since the networks have many more trainable parameters. The number of samples was in line with other implementations of the U-net architecture for EIT (Hamilton et al 2019) and did not result in severe or harmful over-fitting of the networks trained here. More testing could be done to determine if even fewer samples could be used for training. The ADAM optimizer (Kingma and Ba 2015) with an initial learning rate of 5 × 10⁻⁴ and mini batches of 32 samples were used to optimize the parameters. Other learning rates and batch sizes were also tested and resulted in similar minimum loss values and training times. Training of each network was stopped if the validation loss, MSE, failed to decrease over the course of 50 epochs and the parameters (weights and biases) that resulted in the lowest validation loss were saved.

The loss plots for the GNNs and CNNs using different iterates as input are shown in figure 6. As expected, the U-nets using later iterates as input achieved lower minimum validation losses than the networks using TV1 inputs. There was a greater difference in the minimum validation loss values between graph U-nets with different inputs compared to classic U-nets with different inputs. Still, for both types, there was not a large difference between the U-nets using the third and fourth iterates as inputs. The CNNs reached lower raw minimum validation loss values but they cannot be compared directly to the graph U-nets because the loss was computed over different domains and discretizations.

Zoom In Zoom Out Reset image size

The shapes of the loss curves are different between the types of U-nets but consistent across input iteration. The graph U-nets required 130-280 epochs to reach a minimum validation loss while the convolutional U-nets required only 20–30 epochs. After reaching the minimum validation loss, the graph U-nets validation loss values leveled out in later epochs, while the CNN validation loss values slightly increased. For both network types, the training loss values continued to decrease. All of these characteristics were consistent across repetitions of training independent networks, and more research is needed to determine why the differences exist or what the effects are in the final reconstructions. In addition, determining if the network weights resulting in the lowest validation loss produce the 'best' reconstructions is also the topic of future research as metrics other than MSE are also critical in assessing EIT reconstruction quality.

As there is no universally accepted metric for assessing the quality of EIT reconstructions, the metrics listed below, along with visual inspection, will be used collectively to assess reconstruction quality:

Dynamic Range: $\mathrm{DR}=\tfrac{\max (\hat{\sigma })-\min (\hat{\sigma })}{\max ({\sigma }_{\mathrm{true}})-\min ({\sigma }_{\mathrm{true}})}\times 100 \% .$

Total Variation Ratio: $\mathrm{TVR}=\tfrac{\sum \left|{\bf{L}}\hat{\sigma }\right|}{\sum \left|{\bf{L}}{\sigma }_{\mathrm{true}}\right|}\times 100 \%$.

Relative l₂ Voltage Error: ${\mathrm{RE}}_{V}^{{l}_{2}}=\tfrac{{\parallel U(\hat{\sigma })-V\parallel }_{2}}{{\parallel V\parallel }_{2}}$

where 100% is a perfect score for DR and TVR while 0 is ideal for relative voltage error.

Note that these metrics above do not account for the varying sizes of the mesh elements. If desired, the metrics can be scaled relative to element size as well and the networks even trained based on such weighted metrics. Further work is needed to determine if the weighted or unweighted metrics and/or using weighted loss functions during training are more correlated with visually high quality reconstructions. Such work, while interesting is left for future studies. Here, only unweighted metrics were used and reported. Additionally, region of interest (ROI) metrics may be of higher interest than global image metrics. Where appropriate, e.g. lung imaging, ROI metrics are also presented and represent the mean reconstructed conductivity value in the indicated region.