DeepPoly中卷积的纯矩阵实现

DeepPoly是一种证明神经网络在特定输入域上稳定的方法,对于普通的线性层可以达到完全的准确性,对于ReLU层使用线性边界近似。卷积层可以看成一种特殊的线性层,但是在实际实现中,为了性能考虑,需要额外处理。本文借用PyTorch的矩阵操作实现DeepPoly卷积。

线性层

我们先从相对简单的线性层出发,来解释一下DeepPoly的回传过程。

问题定义

对于线性层 xm+1=Wmxm+bm ,已知

Wm,m1lxm1+bm,m1lxmWm,m1uxm1+bm,m1u

xm+1 的范围。

理论分析

xm+1 可以表示成

Wm+1,m1Lxm1+bm+1,m1Lxm+1Wm+1,m1Uxm1+bm+1,m1U

的形式,与输入域的形式类似,因此,只需要求出参数 Wi+1,i1l,bi+1,i1l,Wi+1,i1u,bi+1,i1u 就可以合并两层。通过代入,可以得到参数的计算公式:

Wm+1,m1L[x,z]=yWm[x,y](Wm,m1L[y,z][Wm[x,y]>0]+Wm,m1U[y,z][Wm[x,y]<0])
Wm+1,m1U[x,z]=yWm[x,y](Wm,m1U[y,z][Wm[x,y]>0]+Wm,m1L[y,z][Wm[x,y]<0])

直观理解上,就是如果该层的权值为正,下限就取上一层的下限,上限就取上一层的上限;如果为负,就反之。

具体实现

在具体实现上,为了批量化,需要使用gather函数来根据正负选取权值:

def backsubsitute(self, lowerW: Tensor, lowerB: Tensor, upperW: Tensor, upperB: Tensor):
    new_sz = lowerW.shape[0]
    pre_sz, base_sz = self.lowerW.shape
    W = torch.stack([self.lowerW, self.upperW], dim=1) # [p,2,q]
    W = W.expand((new_sz, -1, -1, -1)) # [o,p,2,q]
    B = torch.stack([self.lowerB, self.upperB], dim=1) # [p,2]
    B = B.expand((new_sz, -1, -1)) # [o,p,2]

    idx_neg = (lowerW < 0).long() # [o,p]
    idxW = idx_neg.unsqueeze(-1).unsqueeze(-1).expand((-1, -1, -1, base_sz)) # [o,p,1,q]
    idxB = idx_neg.unsqueeze(-1) # [o,p,1]
    lW = torch.gather(W, dim=2, index=idxW).squeeze(2) # [o,p,1,q] -> [o,p,q]
    lB = torch.gather(B, dim=2, index=idxB) # [o,p,1]
    lowerW = lowerW.unsqueeze(1) # [o,p] -> [o,1,p]
    newLowerW = torch.bmm(lowerW, lW).squeeze(1) # [o,1,p]x[o,p,q]=[o,1,q] -> [o,q]
    newLowerB = torch.bmm(lowerW, lB).view(-1) # [o,1,p]x[o,p,1]=[o,1,1] -> [o]
    newLowerB += lowerB

    idx_pos = (upperW > 0).long()
    idxW = idx_pos.unsqueeze(-1).unsqueeze(-1).expand((-1, -1, -1, base_sz))
    idxB = idx_pos.unsqueeze(-1)
    uW = torch.gather(W, dim=2, index=idxW).squeeze(2)
    uB = torch.gather(B, dim=2, index=idxB)
    upperW = upperW.unsqueeze(1)
    newUpperW = torch.bmm(upperW, uW).squeeze(1)
    newUpperB = torch.bmm(upperW, uB).view(-1)
    newUpperB += upperB
    return self.preLayer.backsubsitute(newLowerW, newLowerB, newUpperW, newUpperB)

ReLU层也是用线性上下限证明的,因此可以直接复用,区别在于层本身的上下限要改成DeepPoly中的

λxm1xmAxm1+B

卷积层

DeepPoly中卷积层的主要矛盾在于,普通的卷积是共享权值的,但是对于DeepPoly的卷积ReLU层,每个像素的 λ 不一定相等,导致每一层的卷积核不一定相等,因此需要分别存储处理才能达到最高的准确度。

常见的卷积操作为

Xm+1uvw=xyzWm+1,muxyzXmx,v+y,w+z

权值只有在通道上是分离的,而在长宽上都是共享的。在这里,我们就需要把长宽也分离出来:

Xm+1uvw=xyzWm+1,muvwxyzXmx,v+y,w+z

代入 Xmxyz ,得到 Xm+1uvw,Xm1uvw 之间的关系:

Xm+1ijk=uvwWm+1,mijkuvwxyzWm,m1uvwxyzXm1x,j+v+y,k+w+z

同时我们知道, Xm+1uvw,Xm1uvw​ 之间的形式应该为:

Xm+1ijk=abcWm+1,m1ijkabcXm1a,j+b,k+c

和线性层不同,这里的 Xm1a,j+b,k+c​​​ 与 Xm1x,j+v+y,k+w+z​​​ 并不是一一对应的,而是多对一的关系,因此无法直接得到新参数,不过可以通过对应关系,得到

Wm+1,m1ijkabc={Wm+1,mijkuvwWm,m1uvwxyz|a=x,b=v+y,c=w+z}

在代码实现中,为了避免计算每组 (a,b,c) 对应的 (x,y,z,v,w) ,我们从 (i,j,k,u,v,w) 的角度循环,把每次的结果累加到 a=x,b=v+y,c=w+z 的位置。这种方法的缺点在于,循环 v,w 具有依赖性,因此我们只批量化了循环 i,j,k,u 。好在 v,w 一般是很小的数,因此对于性能的影响并不大。

def backsubsitute_conv(self, lowerW: Tensor, lowerB: Tensor, upperW: Tensor, upperB: Tensor,     stride, padding):
    if self.preLayer is None:
        return self.getBoundUnderWB(
            lowerW, lowerB, upperW, upperB, stride, padding)
    out_ch_1, in_ch_1, k_h_1, k_w_1 = self.lowerW.shape
    out_ch_2, out_h_2, out_w_2, in_ch_2, k_h_2, k_w_2 = lowerW.shape
    s1 = 2
    p1 = 1
    s2 = stride
    p2 = padding
    s = s1 * s2
    p = p1 + s1*p2
    k_h_merge = k_h_1 + s1*k_h_2 - s1
    k_w_merge = k_w_1 + s1*k_w_2 - s1
    assert(out_ch_1 == in_ch_2)
    newLowerW = torch.zeros((out_ch_2, out_h_2, out_w_2,
        in_ch_1, k_h_merge, k_w_merge))
    newUpperW = torch.zeros((out_ch_2, out_h_2, out_w_2,
        in_ch_1, k_h_merge, k_w_merge))
    newLowerB = torch.zeros((out_ch_2, out_h_2, out_w_2))
    newUpperB = torch.zeros((out_ch_2, out_h_2, out_w_2))

    W1 = torch.stack([self.lowerW, self.upperW], dim=1)
    B1 = torch.stack([self.lowerB, self.upperB], dim=1)
    i, j, k, u = torch.meshgrid(
        torch.arange(out_ch_2),
        torch.arange(out_h_2),
        torch.arange(out_w_2),
        torch.arange(in_ch_2)
        )
    i, j, k = i.flatten(), j.flatten(), k.flatten()
    u = u.flatten()
    for v in range(k_h_2):
        for w in range(k_w_2):
            coef = lowerW[:,:,:,:,v,w] # [i,j,k,u]
            idx_neg = (coef < 0).long().flatten()
            lW = W1[u, idx_neg].reshape((out_ch_2, out_h_2, out_w_2, in_ch_2, in_ch_1, k_h_1, k_w_1))
            lB = B1[u, idx_neg].reshape((out_ch_2, out_h_2, out_w_2, in_ch_2))
            coef_ex = coef.unsqueeze(4).unsqueeze(5).unsqueeze(6)
            newLowerW[:,:,:,:,s1*v:s1*v+k_h_1, s1*w:s1*w+k_w_1] += (coef_ex * lW).sum(3)
            newLowerB += (coef * lB).sum(3)

            coef = upperW[:,:,:,:,v,w]
            idx_pos = (coef > 0).long().flatten()
            uW = W1[u, idx_pos].reshape((out_ch_2, out_h_2, out_w_2, in_ch_2, in_ch_1, k_h_1, k_w_1))
            uB = B1[u, idx_pos].reshape((out_ch_2, out_h_2, out_w_2, in_ch_2))
            coef_ex = coef.unsqueeze(4).unsqueeze(5).unsqueeze(6)
            newUpperW[:,:,:,:,s1*v:s1*v+k_h_1, s1*w:s1*w+k_w_1] += (coef_ex * uW).sum(3)
            newUpperB += (coef * uB).sum(3)

    newLowerB += lowerB
    newUpperB += upperB
    return self.preLayer.backsubsitute(newLowerW, newLowerB, newUpperW, newUpperB, s, p)

具体代码中,我们还增加了步长和填充的影响。对于卷积ReLU层,和线性层类似,只需要对每个像素取 λ 即可。