Holistically-Nested Edge Detection(HED)论文阅读笔记

Holistically-Nested Edge Detection(HED)阅读笔记

HED是一种基于深监督模式的二至边缘检测模型，以VGG模型作为基础结构，利用多层级的特征融合方式，实现二值边缘检测任务

HED网络示意图

​ 在卷积层后面侧边插入一个输出层side-output层，在side-output层上进行deep supervision，使得结果向着边缘检测方向进行，同时随着side-output层向后大小的越来越小，将receptive field(https://zhuanlan.zhihu.com/p/28492837感受野，表示网络内部的不同位置的神经元队员图像的感受范围的大小)变大，最后通过一个weight-fusion layer得到多尺度下的输出，用pytorch实现如下

class HED(nn.Module):
    def __init__(self):
        super(HED, self).__init__()
        # conv1
        self.conv1 = nn.Sequential(
            nn.Conv2d(3, 64, 3, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(64, 64, 3, padding=1),
            nn.ReLU(inplace=True),
        )

        # conv2
        self.conv2 = nn.Sequential(
            nn.MaxPool2d(2, stride=2, ceil_mode=True),  # 1/2
            nn.Conv2d(64, 128, 3, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(128, 128, 3, padding=1),
            nn.ReLU(inplace=True),
        )

        # conv3
        self.conv3 = nn.Sequential(
            nn.MaxPool2d(2, stride=2, ceil_mode=True),  # 1/4
            nn.Conv2d(128, 256, 3, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(256, 256, 3, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(256, 256, 3, padding=1),
            nn.ReLU(inplace=True),
        )

        # conv4
        self.conv4 = nn.Sequential(
            nn.MaxPool2d(2, stride=2, ceil_mode=True),  # 1/8
            nn.Conv2d(256, 512, 3, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(512, 512, 3, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(512, 512, 3, padding=1),
            nn.ReLU(inplace=True),
        )

        # conv5
        self.conv5 = nn.Sequential(
            nn.MaxPool2d(2, stride=2, ceil_mode=True),  # 1/16
            nn.Conv2d(512, 512, 3, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(512, 512, 3, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(512, 512, 3, padding=1),
            nn.ReLU(inplace=True),
        )

        '''
            在卷积层每一层后面加上一层side-output层，在side-output层上进行deep supervision
        '''
        self.dsn1 = nn.Conv2d(64, 1, 1)
        self.dsn2 = nn.Conv2d(128, 1, 1)
        self.dsn3 = nn.Conv2d(256, 1, 1)
        self.dsn4 = nn.Conv2d(512, 1, 1)
        self.dsn5 = nn.Conv2d(512, 1, 1)
        self.fuse = nn.Conv2d(5, 1, 1)

    def forward(self, x):
        h = x.size(2)
        w = x.size(3)

        conv1 = self.conv1(x)
        conv2 = self.conv2(conv1)
        conv3 = self.conv3(conv2)
        conv4 = self.conv4(conv3)
        conv5 = self.conv5(conv4)

        # 将每个卷积层的输出结果进行双线性上采样来扩展输出
        d1 = self.dsn1(conv1)
        d2 = F.upsample_bilinear(self.dsn2(conv2), size=(h, w))
        d3 = F.upsample_bilinear(self.dsn3(conv3), size=(h, w))
        d4 = F.upsample_bilinear(self.dsn4(conv4), size=(h, w))
        d5 = F.upsample_bilinear(self.dsn5(conv5), size=(h, w))

        # 把每一个尺度上得到的相同大小的image融合到一起，得到最终输出的image，就是具有边缘检测效果的image
        fuse = self.fuse(torch.cat((d1, d2, d3, d4, d5), 1))

        d1 = F.sigmoid(d1)
        d2 = F.sigmoid(d2)
        d3 = F.sigmoid(d3)
        d4 = F.sigmoid(d4)
        d5 = F.sigmoid(d5)
        fuse = F.sigmoid(fuse)

        return d1, d2, d3, d4, d5, fuse

损失函数的设计

​ 作者重置了权重的损失函数，之所以重置权重是因为在图像中边缘像素只占整个图像像素的很少一部分，如果把图像中边缘像素点看成正样本，非边缘像素点看成负样本，则正负样本出现了极度不均衡的情况，作者根据这个问题，采用正负样本不均衡比率来加权两类像素点，对于训练标签中的边缘像素点加大器分类损失的权重，对非边缘像素点减少其分类损失的权重
$$\begin{array}{rl}{\ell }_{\text{side }}^{(m)}\left(\mathbf{W},{\mathbf{w}}^{(m)}\right)=-\beta & \sum _{j\in Y_{+}}\log \mathrm{Pr}\left(y_j=1|X;\mathbf{W},{\mathbf{w}}^{(m)}\right)\\ -(1-\beta )& \sum _{j\in Y_{-}}\log \mathrm{Pr}\left(y_j=0|X;\mathbf{W},{\mathbf{w}}^{(m)}\right)\end{array}$$
其中
$$\beta =Y_{+}/\left(Y_{+}+Y_{-}\right)$$
是样本不均衡比率。Y+代表标签图像中边缘像素点的个数。Y_代表标签图像中非边缘像素点的个数
$$P\left(y_j=1|X;W,w^m\right)=\sigma \left(a_j^{(m)}\right)\in [0,1]$$
代表的是在第m个阶段中位于像素点j处的预测值，激活函数为sigmoid函数，
$$l_{side}^{(m)}$$
表示第m个阶段的损失函数，最后的数据融合层的损失函数如下所示
$$L_{fuse}(W,w,h)=\mathrm{Dist}\left(\mathrm{Y},{\hat{Y}}_{fuse}\right)$$

$${\hat{Y}}_{fuse}=\sigma \left(\sum _{m=1}^M{h}_m{\hat{A}}_{side}^m\right)$$

hm为第m个阶段输出特征的融合权重，融合方式采用的是特征叠加求和，最终在网络训练优化算法使用的是随机梯度下降算法，模型自我优化的过程实质是使得下式取得最小值的过程
$$(\mathbf{W},\mathbf{w},\mathbf{h}{)}^{\star }=\mathrm{argmin}\left({\mathcal{L}}_{\mathrm{s}\mathrm{i}\mathrm{d}\mathrm{e}}(\mathbf{W},\mathbf{w})+{\mathcal{L}}_{\mathrm{f}\mathrm{u}\mathrm{s}\mathrm{e}}(\mathbf{W},\mathbf{w},\mathbf{h})\right)$$