Anderson 视角

使用生成对抗网络创建合成伤口数据集

mm
将 Unite.AI 添加到您在 Google 上的首选来源

首次使用生成对抗网络创建合成伤口图像数据集,以解决医疗机器学习应用中缺乏多样化和可访问的伤口图像内容的问题。

该系统称为WG2AN,是巴腾工程与技术学院和AI健康公司eKare之间的合作,后者专门应用机器学习方法来测量和识别伤口。

生成对抗网络是在eKare提供的100-4000张标记的立体慢性伤口图像上训练的,包括匿名的压力、手术、淋巴血管事故、糖尿病和烧伤等各种类型的伤口图片。原始图像大小在1224×1224到2160×2160之间,都在可用光线下由医生拍摄。

为了适应模型训练架构中的潜在空间,图像被重缩放到512×512,并从背景中提取出来。为了研究数据集大小的影响,测试运行在100、250、500、1000、2000和4000张图像的批次上进行。

来源: https://ietresearch.onlinelibrary.wiley.com/doi/pdfdirect/10.1049/tje2.12033

来源: https://ietresearch.onlinelibrary.wiley.com/doi/pdfdirect/10.1049/tje2.12033

上面的图像显示了根据训练集大小和每次运行的epoch数,图像的细节和粒度逐渐增加。

WG2GAN的架构。来源: https://ietresearch.onlinelibrary.wiley.com/doi/pdfdirect/10.1049/tje2.12033

WG2GAN的架构。 来源: https://ietresearch.onlinelibrary.wiley.com/doi/pdfdirect/10.1049/tje2.12033

WG2GAN在PyTorch上运行,使用相对精简的消费级设置,配备8GB的VRAM和GTX 1080 GPU。训练时间在100-4000张图像的数据集大小范围内,跨越多个epoch,批次大小为64,作为准确性和性能之间的折衷。使用Adam优化器,在学习率为0.0002的条件下进行训练,直到损失为零。

上左,应用于伤口区域的分割;上中,实际伤口图像;上右,基于原始来源的合成伤口图像。下左,原始伤口;下右,WG2GAN生成的合成伤口图像。

上左,应用于伤口区域的分割;上中,实际伤口图像;上右,基于原始来源的合成伤口图像。下左,原始伤口;下右,WG2GAN生成的合成伤口图像。

在医疗数据集中,标注是一个不可避免的瓶颈。在本例中,研究人员使用了一种半自动标注系统,利用eKare早期的研究成果,该研究使用Play-Doh创建了现实世界的伤口模型,并为语义上下文着色。

eKare伤口模型

eKare伤口模型

研究人员注意到,在训练的初始阶段,当数据集非常多样化,权重被随机化时,模型需要很长时间(75个epoch)才能“稳定下来”。

当数据多样化时,GAN和编码器/解码器模型在早期阶段都难以获得泛化,如上图所示,这是WG2GAN的训练过程,从开始到损失为零的整个过程。

在训练过程中,必须小心确保模型不会过度专注于任何一个迭代或epoch的特征或特性,而是继续泛化到一个可用的平均损失,而不产生过度抽象原始数据的结果。在WG2GAN的例子中,这将冒着创建无边界、完全“虚构”的伤口的风险,这些伤口在太广泛的无关伤口类型中被拼接,而不是产生特定伤口类型内的准确变化范围。

控制机器学习数据集的范围

训练数据较少的模型可以更快地泛化,论文的研究人员认为,使用较小的数据集和较少的epoch可以获得更真实的图像:1000张图像的数据集,训练200个epoch。

虽然较小的数据集可能在较短的时间内获得高度真实的图像,但生成的图像范围和类型必然会受到限制。GAN和编码器/解码器训练过程中,输入数据的体积和多样性、生成图像的保真度以及生成图像的真实性之间存在着微妙的平衡——这些都是范围和权重的问题,并不仅限于医疗图像合成。

医疗数据集中的类别不平衡

一般来说,医疗机器学习不仅面临数据集的缺乏,还面临类别不平衡的问题,即某些特定疾病的数据在其宿主数据集中占比非常小,以至于可能被视为异常数据或在训练过程中被同化。

已经提出了几种方法来解决这个问题,例如欠采样或过采样。然而,问题通常是通过开发专门针对单个医疗问题的数据集来解决的。虽然这种方法在每个案例中都有效,但它却促进了医疗机器学习研究领域的“巴尔干化”文化,并可能减缓了该领域的整体进步。

机器学习作家,人类图像合成领域专家。曾任Metaphysic.ai研究内容负责人,直至其解散并并入DNEG的Brahma.ai。
Portfolio site:martinanderson.ai
Contact:martin@martinanderson.ai