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

要約

畳み込みニューラルネットワーク（CNN）とトランスフォーマベースのモデルは、高レベルの特徴を抽出し、画像の重要な側面を捉えることができるため、医療画像のセグメンテーションに広く適用されている。しかし、高精度の必要性と低計算コストへの欲求の間には、しばしばトレードオフの関係がある。より高いパラメータを持つモデルは、理論的にはより高い性能を達成できるが、計算の複雑さとメモリ使用量の増加をもたらすため、実装は現実的ではない。本論文では、U-Netに基づく軽量なモデル、すなわちU-Liteについて検討する。U-Liteは深さ方向に分離可能な畳み込み（Depthwise Separable Convolution）の原理に基づいて設計されており、CNNの強みを生かすと同時に、計算パラメーターの数を大幅に減らすことができる。具体的には、モデルの受容野を拡大するために、エンコーダとデコーダの両方で7×7のカーネルを持つ軸方向の深さ方向の畳み込みを提案する。さらに性能を向上させるため、ボトルネックとなる部分には、3×3のフィルターを持つ複数の軸方向希釈深度方向畳み込みを分岐の1つとして使用する。全体として、U-Liteは878Kのパラメータしか含まず、従来のU-Netより35倍少なく、他の最新の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, DeepL