目录
用于 3D 超声平面定位的热启动和主动终止智能体[MICCAI 2019] | DDQN-AT | 特征点对齐、强化学习、CNN、RNN、SP 定位 |
摘要
1.为RL框架配备了一个特征点感知对齐模块,使用一个特征点感知对齐模型来利用地图集中关于标准平面的先验信息,减少复杂的宫内环境和不同的胎儿姿势造成的大搜索空间,为智能体动作提供热启动和强大的空间边界,从而确保其有效性。
2.使用CNN来检测超声体积中的解剖标志,并将它们注册到一个平面特定的图谱中,CNN在由对齐模块产生的有界环境中搜索目标sp。
因此,图集的平面配置为RL智能体操作提供了强大的空间边界和有效的初始化。
3.采用自适应的基于RNN的终止模块,在最优交互时动态停止RL代理。
基于递归神经网络(RNN)的策略来主动终止智能体的交互过程,而不是被动地和经验地终止agent推理。基于RNN的策略可以自适应地找到最优的终止点,从而同时提高了定位系统的精度和效率。
方法
为RL框架配备了(1)一个热启动的特征点感知对齐模块,以确保其有效性,
(2)采用基于递归神经网络的策略来主动终止交互过程,以提高其准确性和效率
平面定位的深度强化学习框架
类似于超声医生操作(action)探针(agent)扫描器官(environment),同时可视化屏幕上的中间平面(state),agent与environment交互,以学习能够最大化累积奖励(reward)的最佳策略,直到获得SP(最终state)。
设笛卡尔坐标系中的一个平面表示为cos(α)x+cos(β)y+cos(φ)z+d=0,其中n=(cos(α),cos(β),cos(φ))为法线,d为平面到体积中心原点的距离。
当智能体与环境交互时,系统将获得最优的平面参数。
动作Action
动作定义为对平面参数的增量调整。与[1]类似,完整的动作空间定义为8个动作,
即A=。
给定一个动作,平面参数被相应地修改(例如)。
对每次迭代只调整一个平面参数,其他参数保持不变。
角度调整的步长为,
每次迭代的距离步长ad设置为0.5voxel
智能体执行动作后,平面参数相应更新为
每个有效的动作根据规则r=sgn(D(Pt−1,Pg)−D(Pt,Pg))获得其标量奖励r,其中D计算从预测平面Pt到地面真实Pg的欧氏距离。r朝向首选目标。∈{+1,0,−1}指示智能体是否正在移动
在Q-learning之后,Deep Q-Network[6](DQN)可以通过深度网络学习状态动作值函数Q(s,a),
作为动作选择策略。
状态State
根据当前平面参数从体中采样体素重建的二维超声图像。
由于重建后的图像大小可能会改变,我们将图像填充到一个正方形上,并将其大小调整为224×224。
此外,我们将前两次迭代得到的两幅图像与当前平面连接起来,丰富了类似于[31]的状态信息。
奖励Reward
在本研究中,奖励被定义为agent是否接近或远离目标
sgn(·)为符号函数。
Pt t时刻预测平面的平面参数(Pt,t时刻 plane)
Pg 目标真实平面(Pg,ground plane)
欧几里德距离
R∈{−1,0,+1},
其中+1和−1分别表示正运动和负运动,0表示没有调整。
智能体Agent
采用Q-learning的[32]作为SP定位的解决方案。
与现有的使用深度神经网络直接估计[31]值[31]的工作不同,
dueling learning[33]是用来鼓励agent学习哪些状态应该更加权,哪些状态在选择适当的行动时是冗余的。
具体来说,Q-value函数分别分解为一个状态值函数和一个与状态相关的动作优势函数。
agent神经网络的体系结构。蓝、绿、黄色的矩形分别代表卷积层、全局平均池化层、全连通层。
如图所示,deep duel neural network以状态作为输入,输出动作。
使用预先训练过的VGG [34]作为卷积主干。
#models提供了常用网络结构,及其预训练模型,可以通过简单调用来读取网络结构和预训练模型
from torchvision.models import vgg13_bn
#autograd自动微分
from torch.autograd import Variable
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.feature = vgg13_bn(pretrained=False).features
self.linear1 = nn.Linear(512*7*7, 512)
self.relu = nn.ReLU(inplace=True)
self.drop = nn.Dropout(0.5)
self.fc1_adv = nn.Linear(in_features=512, out_features=128)
self.fc1_val = nn.Linear(in_features=512, out_features=128)
self.fc2_adv = nn.Linear(in_features=128, out_features=8)
self.fc2_val = nn.Linear(in_features=128, out_features=1)
def forward(self, x):
x = self.feature(x)
x = x.view(x.size(0), -1)
x = self.drop(self.relu(self.linear1(x)))
adv = self.drop(self.relu(self.fc1_adv(x)))
val = self.drop(self.relu(self.fc1_val(x)))
adv = self.fc2_adv(adv)
val = self.fc2_val(val).expand(x.size(0), 8)
x = val + adv - adv.mean(1).unsqueeze(1).expand(x.size(0), 8)
return x
每层特征数中的特征数分别为64、128、256、512、512个。
为了缓解梯度消失的问题,在神经网络的每个卷积层之后添加批归一化层[35]。
将提取出的高级特征输入到两个独立的全连接层流中,以估计状态值和与状态相关的动作优势值。全连接层的隐藏单元在状态值估计流中为512、128、1,
在与状态相关的动作值估计流中为512、128、8。这两个流的输出被融合,以预测最终的q值
回放缓存Replay Buffer
内存容器,它存储agent的转换,以执行学习过程的经验重放。
元素转换通常用一个向量表示,
给定优先级的replay buffer存储agent-环境交互中每个步骤t的状态s、动作a、奖励r和下一步s0的状态的数据序列,这可以消除数据相关性,提高采样效率。
本研究采用优先级回放缓存[36]来提高学习效率。
具有高误差的第i个序列元素将优先从buffer中采样,其采样概率可以计算为,
p = 0.6控制使用多少优先级
δ = 0.05被设置为调整误差ei。
采用了重要性抽样权值来纠正优先级重放中数据分布变化引起的偏差。
损失Loss
在agent达到最终状态之前,存在着许多状态和行动的组合。
在经典的Q-learning(Danain,1992)中,我们没有将所有的state,action,(Q-values)存储在表中,
而是选择了像深度Q network(DQN)(Mnih等人,2015)这样的策略,
其中CNN被用来模拟state,action和Q-values之间的关系。
具体来说,我们采用了双DQN(DDQN)(Van Hasselt et al.,2016)方法,以减轻DQN引起的向上偏差,更好地稳定训练过程。
将Q-learning函数Q(s,a;w)分解为两个独立的估计量:
状态值函数
状态相关的动作优势函数,
W= 。
其中s是代理的输入状态,a是动作,wc,wv,ws分别表示卷积层和两个全连通层的参数,
agent的Q值函数的计算方法为:
|A| = 8表示动作空间的大小
为了提高DQN对三维超声中噪声环境E的鲁棒性,最终选择了Double DQN(DDQN)[11]作为平面定位的深度智能体。
DDQN的损失函数被定义为:
γ∈[0,1]是衡量未来奖励的折扣因子,
s,a是当前的状态和动作,ˆs和ˆa是下一步的状态和行动。
M是经验重放存储器,以避免频繁的数据采样。w和˜w为当前和目标网络参数。
Q network,Q(ω)用于获取当前步骤的Q-values和下一步中的动作a0。
表示用来估计下一步的Q-values的目标Q网络。
ω和˜ω是Q network和目标Q network的网络参数。
利用− greedy 算法选择action,以平衡训练过程中的探索和开发(Mnih et al.,2015)。
通常,Q-values将由全连接层输出,并指导agent的操作来更新平面参数。
热启动的特征点感知平面对准
图像质量低、数据量大、胎儿姿势多样,在三维超声中定位SP是非常具有挑战性的。
此外,由于嘈杂的三维超声环境,在[9]中使用的随机状态初始化经常无法定位SPs。
因此,提出了一个landmark感知对齐模块,作为通过解剖先验知识的搜索过程的专用warm start。
为了确保RL 智能体与嘈杂的3D 超声环境的有效交互,提出了一个特征点感知平面对齐模块来利用解剖先验,并为智能体提供一个热启动。
(优化中warm start一般是采用“相关或简化问题的最优解(至少是可行解)”来作为原问题的初始值,因此可以更快收敛。
此外,在非线性优化问题中可能存在很多局部极小值,良好的初值猜测可以有效的避免陷入局部极小,因此warm start为优化问题提供了非常好的开端。
例如,采用上一次优化的结果作为下一次优化的初值。)
(图像中的landmark译为特征点,能够以点的形式描述出某个物体的形状和部件。得到特征点后,可以得到简单的几何形状信息或者是在其周围的提取强特征进行进一步的利用,如分类或回归等)
具体来说,我们首先用定制的3D U-net[8]检测了胎儿大脑的三个标志,即胼胝体膝、胼胝体压部和小脑蚓部,如图(a)三个红点所示。
然后利用这些特征点
将测试体积与地图集 对齐。(atlas 包含参考特征点和标准平面参数的)
(landmark配准类似于人脸关键点的对齐,人脸对齐可以简单的理解为就是为了把比如歪着的人脸、尺度不正常的人脸、在图像中位置偏移的人脸,都矫正成比较标准的人脸
图出自人脸检测-人脸对齐-人脸识别原理及方法_陈 洪 伟的博客-CSDN博客_人脸比对原理)
RigidRegistration刚性配准 影像对准结果
图出自PythonITK3D刚性影像对准RigidRegistration置中初始化教程与范例 _ 小媛啾
与[5,7]将共同的解剖模型应用于各种标准平面的解剖模型不同,建议为每个平面选择特定的图谱进行改进定位精度。
图3 100个超声体积的特征点(左)与特定地点的地图集空间(中间)对齐,
为RL智能体操作提供了强大的空间边界(右)。
红、绿、蓝点表示图1(a).所示的标志
图1 胎儿大脑平面的三维图像。(a)蓝线表示经丘脑(TT)和经小脑(TC)平面位置。红点(从左到右)显示三个标志:胼胝体膝、胼胝体压部、小脑蚓部。(b)(c)两卷的平面示例说明了巨大的搜索空间和较大的胎儿姿势变化。(彩色图形在线版)
图5 胎儿脑标志物在对齐前后分布的三维可视化。不同的颜色点代表不同的类别地标。
最后,将图谱的标准平面映射到测试超声体积上,并作为我们的RL 智能体的热启动。
一种标准平面P的图集选择公式如下:
其中,i,j是体积指数。niP是P在体积i中的法线,Θ计算法线之间的夹角,diP是从体积i中的平面P到原点的距离。Tij是从体积j到体积i的变换矩阵,由基于特征点标注的刚性配准决定。
图3显示了我们对100超声体积的特征点对齐的效果。精确的对齐保证了RL代理的初始点的有效性,从而导致快速和改进的平面定位。
过程
landmark感知对齐
1.使用预先训练好的CNN来检测超声体积中的解剖标志(Landmarks)
2.将每个超声体积作为一个初始化的模板(Atlas),采用迭代最近点算法对剩余的超声体积进行基于landmark的刚性配准
3.根据每个模板(Atlas)的平面 与 标准平面真实值 之间测量的平均平面误差,
选择误差最小的作为最终模板(Atlas)。
warm start
1.与直接回归不同,基于landmark检测和匹配,将测试体积与模板atlas对齐
2.通过由landmark计算的 变换矩阵 将 体积 映射到 模板空间,得到有界搜索环境
3.初始平面:
训练时,训练DDQN的起始平面在±25◦的角度范围和±10 mm的距离范围内,围绕真实平面随机初始化。该范围由基于模板的平均平面定位误差决定。
测试时,将模板的标准平面映射到测试体积上,作为初始平面。
4.CNN在有界环境中搜索目标标准切面(sp),landmark的高斯映射作为标准平面生成
基于递归神经网络的主动终止
为了确保RL 智能体之间的有效交互,提出了一个基于RNN的主动终止(AT)模块来告诉智能体何时停止。
通常,没有明确定义的标准来终止RL学习的迭代推理。对终止状态的低估和高估往往会降低最终的定位。
现有的工作利用一个预定义的最大阶跃,一个较低的Q值[1]或Q值[3]的振荡作为终止的指标。
第一个如果设置为较大的数字,会浪费大量的计算资源,而后两个并不一定会得到最优的结果。
图4 训练数据集上8个候选动作(黄色)和ADI(蓝色)的平均q值。绿点表示ADI最大的最优终止步长。
如图4所示,角度和距离改善最大的最优终止步长(ADI)既不是最大的步长,也不是Q值最小的步长。
这促使我们提出一种新的策略来积极地学习最佳步骤。具体来说,考虑到迭代推理的序列特征,
如图2所示,我们用递归神经网络表示了q值序列与最优步长之间的映射。
每个状态的8个候选动作的q值作为我们的RNN的输入,
然后RNN学习输出ADI最高的最优终止步骤,即最显著的角度和距离改善。
用智能体对训练量的推理结果来训练RNN模型。
在测试过程中,根据RNN输出终止了智能体的迭代动作,得到了最终的平面参数。
有了这种主动终止机制,智能体可以在不进行过度迭代的情况下进行有效的推断。
实验
Dataset Setting
验证了解决方案,在超声体积定位两个标准平面(TT和TC)的胎儿大脑。
建立了一个来自430名健康孕妇志愿者的430个产前超声体积的数据集。
胎龄范围为19-31周,比[4,7]要宽得多。
数据集的平均体积大小是270×207×235,统一的体素大小是0.5×0.5×0.5 mm3。
有5年经验的超声医生为所有体积提供地标和标准平面的手工注释。
将数据集随机分成330体积和100体积进行训练和测试
(各向同性就是各个方向体素相同,如文件的体素空间(spacing)为 1mm*1mm*1mm
各向异性就是各个方向体素不同, 如文件的体素空间(spacing)为 1mm*1mm*5mm)
经当地机构审查委员会批准,所有超声都是匿名的,并由专家使用带有集成3D探头的迈瑞DC-9超声系统获得。
train是训练集,
val是训练过程中的测试集,是为了在边训练边看到训练的结果,及时判断学习状态。来验证是否过拟合、以及用来调节训练参数等。与train没有交集,对最终训练出的模型没有贡献。
test就是训练模型结束后,用于评价模型结果的测试集。
只有train就可以训练,val不是必须的,比例也可以设置很小。
test对于model训练也不是必须的,但是一般都要预留一些用来检测,通常推荐比例是8:1:1
但现在很多模型都不用validation了。
现在模型中防止过拟合的机制已经比较完善了,Dropout\BN等做的很好了。
而且很多时候都用原来的模型进行fine tune,也比从头开始更难过拟合。
所以大家一般都定一个训练迭代次数,直接取最后的模型来测试。
参数
PyTorch
RL
Adam优化器[40]对整个框架进行了训练,
对DDQN进行了 epoch =100(约4天)的训练、
将损失函数(公式2)中的discount factor(折扣因子) γ设为0.9。(γ∈[0,1])
目标Q network每2000次迭代复制一次Q network的参数。
# =============== define training
# batch size, INT
self.batch_size = 4
# target net weight update term, INT
# 目标Q network每2000次迭代复制一次当前Q network的参数
self.target_step_counter = 2000
# learning rate, FLOAT
self.lr = 5e-5
# weight decay, FLOAT
self.weight_decay = 1e-4
# reward decay γ, FLOAT
self.gamma = 0.95
# memory capacity, INT 有优先级的Replay-buffffer
self.memory_capacity = 15000
# =============== define default
# gpu id
self.gpu_id = 0
# total epoch
self.num_epoch = 100
# max steps
self.max_step = 75
RNN(主动终止)
对于训练RNN变体(原始RNN和LSTM),优化器是Adam与L1回归损失,批大小= 100,隐藏大小= 64和epoch=200(约15 min)。L2损失用于训练。采用迭代最近点算法进行测试用例与图集之间的刚性配准。
3D U-net(检测胎儿脑landmarks)
Adam优化器,batch size = 1, learning rate = 0.001, moment is 0.5, epoch = 40
受GPU内存限制,超声体积调整为0.4倍用于训练。
特征点的高斯地图作为标准切面生成。
用L2-范数回归损失训练了一个定制的3D U-net [38](检测了胎儿大脑的三个标志,即胼胝体膝、胼胝体压部和小脑蚓部,如图(a)三个红点),所示表示为:
其中N = 3为landmark数,Hi、ˆHi分别表示第i个预测的landmark热图和标准landmark热图。
基于验证集选择了超参数,并在保留的测试集上使用几个指标来评估我们的方法的性能。
对每个具有不同大小的超参数进行了模型训练,并评估了在验证数据集上的性能。
选择了具有最佳验证性能的超参数值作为训练阶段的默认设置。
在本研究中,搜索了三个高影响的超参数,包括Replay Buffer的大小、γ和。
评定标准
使用了三个标准来评估平面定位
空间相似性:
1.两个平面之间的二面角(Ang)
np,ng表示预测平面和目标平面的法线
2.两个平面到原点的欧氏距离差(Dis)
dp,dg表示与体积的距离原点到预测平面,原点到真实平面
(Ang和Dis是基于平面采样函数,即cos(α)x + cos(β)y + cos(γ)z = d,
有效体素大小为0.5 mm3/voxel)
内容相似性:
3.峰值结构相似度(SSIM)。
并将迭代t中的ADI定义为距离起始平面的距离和角度的累积变化之和,如下
-RL 智能体(DDQN-at)在两种标准平面的定位化上具有良好的性能,优于基于回归的方法(Regress)、基于注册的方法(AtlasRegist)和它们的组合(RegistRegress)。
这可以归因于主动的交互作用沿着轨迹向最优平面搜索的智能体的搜索过程。
-在ddqn上使用特征点感知空间对齐模块作为热启动时,它们在标准平面定位上的性能明显优于不对齐的方法(DDQN-nA)。
此外,所提出的空间对齐模块也可以部署在回归模型中,并导致明显的改进(RegistRegress)。
-所提出的主动终止可以用更少的智能体迭代来实现更好的定位。
与其他终止策略相比,如最大步长(DDQN-maxS)和最小q值(DDQN-minQ),我们的基于AT的方法通常提供更好的定位性能。
其中,DDQN-AT(LSTM)的学习效果最好,因为它对Q-value sequence的学习能力更强。更重要的是,在配备了AT模块后,我们的RL-智能体平均需要13个步骤来定位标准平面,而没有使用AT模块需要100个步骤。考虑到这样的迭代步骤花费最大的计算,AT模块肯定会提高RL智能体的效率。
标准切面(左)和预测(右)平面
图5。TC(左)和TT(右)的结果
左,主动终止步骤(红色虚线)与绿点的最优步骤比较,标准切面三维可视化和预测平面(右)。
在图5中,可视化了DDQN-AT(LSTM)对TC和TT平面定位的两个测试结果。
从图像内容和空间关系来看,对于这两个任务,都准确地捕获了平面,这非常接近标准切面值。
主动终止策略还提供了从Q-value sequence中学习的能力,并达到最佳的终止步骤(绿点),以实现大角度和距离的改进(ADI)。
定量和定性分析
所提出的RL代理(DDQN-AT)在两种标准平面的定位上具有良好的性能,优于基于回归的方法(回归)、基于注册的方法(AtlasRegist)和它们的组合(注册回归)。这可以归因于主动的交互作用沿着轨迹向最优平面搜索的代理的搜索过程。
其次,从表中可以清楚地看出,当在ddqn上使用地标感知空间对齐模块作为热启动时,它们在标准平面定位上的性能明显优于不对齐的方法(DDQN-nA)。此外,所提出的空间对齐模块也可以部署在回归模型中,并导致明显的改进(注册回归)。
最后,所提出的主动终止可以用更少的推理迭代来实现更好的定位。与其他终止策略相比,如最大步长(DDQN-maxS)和最小q值(DDQN-minQ),我们的基于AT的方法通常提供更好的定位性能。其中,DDQN-AT(LSTM)的学习效果最好,因为它对q值序列的学习能力更强。更重要的是,在配备了AT模块后,我们的RL-agent平均需要13个步骤来定位标准平面,而没有使用AT模块需要100个步骤。考虑到这样的迭代步骤花费最大的计算,AT模块肯定会提高RL代理的效率。
在图5中,我们可视化了DDQN-AT(LSTM)对TC和TT平面定位的两个测试结果。从图像内容和空间关系来看,对于这两个任务,我们的方法都准确地捕获了平面,这非常接近真实值。
我们的主动终止策略还提供了从q值序列中学习的能力,并达到最佳的终止步骤(绿点),以实现大角度和距离的改进(ADI)。
博文相关链接
识别、提取三维超声中标准平面的总结+论文+代码合集_luemeon的博客-CSDN博客