探索用于端到端十亿像素深度学习的共享内存架构
计算机视觉与模式识别
2023-04-25 v1 图像与视频处理
摘要
深度学习在医学影像领域取得了巨大进步，这得益于 GPU 的硬件进步。新模型开发的一个主要限制是训练期间 GPU 内存资源的饱和。在计算病理学领域尤其如此，其中的图像通常包含超过 10 亿像素。由于硬件限制，这些病理图像传统上被划分为小块以启用深度学习。在这项工作中，我们探索苹果公司最近发布的 M1 Ultra 系统级芯片（SoC）上的 GPU/CPU 共享内存架构是否可能提供一种解决方案。这些价格适中的系统（低于 5000 美元）提供对 128 GB 统一内存的访问（配备 M1 Ultra SoC 的 Mac Studio）。作为十亿像素深度学习的概念验证，我们在全切片图像（WSI）的十亿像素区域上从背景中识别组织。该模型是一个改进的 U-Net（4492 个参数），利用大核和高步幅。M1 Ultra SoC 能够直接在十亿像素图像（1600064000 像素，10.24 亿像素）上以批大小 1 训练该模型，使用超过 100 GB 统一内存进行处理，在 Tensorflow 2/Keras 下平均每批耗时 1 分 21 秒。正如预期，模型收敛，Dice 分数高达 0.989 0.005。至此训练历时 111 小时 24 分钟，共 4940 步。其他高 RAM 的 GPU 如 NVIDIA A100（最大商用可达 80 GB，约 15000 美元）尚未广泛可用（在 Amazon Web Services 上以 8 个为一组在选定区域预览，每小时 $40.96）。本研究是朝向使用通用网络架构实现 WSI 级端到端深度学习的有希望的一步。
引用
@article{arxiv.2304.12149,
title = {Exploring shared memory architectures for end-to-end gigapixel deep learning},
author = {Lucas W. Remedios and Leon Y. Cai and Samuel W. Remedios and Karthik Ramadass and Aravind Krishnan and Ruining Deng and Can Cui and Shunxing Bao and Lori A. Coburn and Yuankai Huo and Bennett A. Landman},
journal= {arXiv preprint arXiv:2304.12149},
year = {2023}
}