UNET Segmentation

U-Net is an architecture for semantic segmentation. It consists of a contracting path (Up to down) and an expanding path (Down to up). During the contraction, the spatial information is reduced while feature information is increased. The contracting path follows the typical architecture of a convolutional network. It consists of the repeated application of two 3x3 convolutions (unpadded convolutions), each followed by a rectified linear unit (ReLU) and a 2x2 max pooling operation with stride 2 for downsampling. At each downsampling step, we double the number of feature channels. Every step in the expansive path consists of an upsampling of the feature map followed by a 2x2 convolution (“up-convolution”) that halves the number of feature channels, a concatenation with the correspondingly cropped feature map from the contracting path, and two 3x3 convolutions, each followed by a ReLU. The cropping is necessary due to the loss of border pixels in every convolution. At the final layer, a 1x1 convolution is used to map each 64-component feature vector to the desired number of classes. In total the network has 23 convolutional layers.

Channel-wise Cross Attention

Channel-wise Cross Attention is a module for semantic segmentation used in the UCTransNet architecture. It is used to fuse features of inconsistent semantics between the Channel Transformer and U-Net decoder. It guides the channel and information filtration of the Transformer features and eliminates the ambiguity with the decoder features. Mathematically, we take the $i$-th level Transformer output $\mathbf{O\_{i}} \in \mathbb{R}^{C×H×W}$ and i-th level decoder feature map $\mathbf{D\_{i}} \in \mathbb{R}^{C×H×W}$ as the inputs of Channel-wise Cross Attention. Spatial squeeze is performed by a global average pooling (GAP) layer, producing vector $\mathcal{G}\left(\mathbf{X}\right) \in \mathbb{R}^{C×1×1}$ with its $k$th channel $\mathcal{G}\left(\mathbf{X}\right) = \frac{1}{H×W}\sum^{H}\_{i=1}\sum^{W}\_{j=1}\mathbf{X}^{k}\left(i, j\right)$. We use this operation to embed the global spatial information and then generate the attention mask: $$ \mathbf{M}\_{i} = \mathbf{L}\_{1} \cdot \mathcal{G}\left(\mathbf{O\_{i}}\right) + \mathbf{L}\_{2} \cdot \mathcal{G}\left(\mathbf{D}\_{i}\right) $$ where $\mathbf{L}\_{1} \in \mathbb{R}^{C×C}$ and $\mathbf{L}\_{2} \in \mathbb{R}^{C×C}$ and being weights of two Linear layers and the ReLU operator $\delta\left(\cdot\right)$. This operation in the equation above encodes the channel-wise dependencies. Following ECA-Net which empirically showed avoiding dimensionality reduction is important for learning channel attention, the authors use a single Linear layer and sigmoid function to build the channel attention map. The resultant vector is used to recalibrate or excite $\mathbf{O\_{i}}$ to $\mathbf{\bar{O}\_{i}} = \sigma\left(\mathbf{M\_{i}}\right) \cdot \mathbf{O\_{i}}$, where the activation $\sigma\left(\mathbf{M\_{i}}\right)$ indicates the importance of each channel. Finally, the masked $\mathbf{\bar{O}}\_{i}$ is concatenated with the up-sampled features of the $i$-th level decoder.

UCTransNet: Rethinking the Skip Connections in U-Net from a Channel-wise Perspective with Transformer

기법

Channel-wise Cross Attention