1M parameters are enough? A lightweight CNN-based model for medical image segmentation

要約

畳み込みニューラル ネットワーク (CNN) と Transformer ベースのモデルは、高レベルの特徴を抽出し、画像の重要な側面をキャプチャできるため、医療画像のセグメンテーションに広く適用されています。
ただし、多くの場合、高精度の必要性と低い計算コストの要望の間にはトレードオフが存在します。
より高いパラメーターを持つモデルは、理論的にはより優れたパフォーマンスを達成できますが、計算がより複雑になり、メモリ使用量も増えるため、実装は現実的ではありません。
このペーパーでは、同等のパフォーマンスを維持できる、またはより優れたパフォーマンスを実現できる軽量の U-Net ベースのモデル、つまり U-Lite を探します。
私たちは、深さ方向分離可能畳み込みの原理に基づいて U-Lite を設計し、モデルが CNN の強みを活用し、計算パラメーターの数を大幅に削減できるようにします。
具体的には、モデルの受容野を拡大するために、エンコーダーとデコーダーの両方でカーネル 7×7 を使用した軸方向深さ方向畳み込みを提案します。
パフォーマンスをさらに向上させるために、ブランチの 1 つとしてボトルネックのフィルター 3×3 を備えた複数の Axial Dirated Depthwise Convolutions を使用します。
全体として、U-Lite に含まれるパラメータは 878K のみで、これは従来の U-Net より 35 分の 1、他の最新の Transformer ベースのモデルよりもはるかに少ないです。
提案されたモデルは、他の最先端のアーキテクチャと比較して、医療セグメンテーション タスクで優れたパフォーマンスを達成しながら、大量の計算の複雑さを削減します。
コードは https://github.com/duong-db/U-Lite から入手できます。

要約(オリジナル)

Convolutional neural networks (CNNs) and Transformer-based models are being widely applied in medical image segmentation thanks to their ability to extract high-level features and capture important aspects of the image. However, there is often a trade-off between the need for high accuracy and the desire for low computational cost. A model with higher parameters can theoretically achieve better performance but also result in more computational complexity and higher memory usage, and thus is not practical to implement. In this paper, we look for a lightweight U-Net-based model which can remain the same or even achieve better performance, namely U-Lite. We design U-Lite based on the principle of Depthwise Separable Convolution so that the model can both leverage the strength of CNNs and reduce a remarkable number of computing parameters. Specifically, we propose Axial Depthwise Convolutions with kernels 7×7 in both the encoder and decoder to enlarge the model receptive field. To further improve the performance, we use several Axial Dilated Depthwise Convolutions with filters 3×3 for the bottleneck as one of our branches. Overall, U-Lite contains only 878K parameters, 35 times less than the traditional U-Net, and much more times less than other modern Transformer-based models. The proposed model cuts down a large amount of computational complexity while attaining an impressive performance on medical segmentation tasks compared to other state-of-the-art architectures. The code will be available at: https://github.com/duong-db/U-Lite.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google