DeepPoly中卷积的纯矩阵实现
DeepPoly是一种证明神经网络在特定输入域上稳定的方法,对于普通的线性层可以达到完全的准确性,对于ReLU层使用线性边界近似。卷积层可以看成一种特殊的线性层,但是在实际实现中,为了性能考虑,需要额外处理。本文借用PyTorch的矩阵操作实现DeepPoly卷积。
线性层
我们先从相对简单的线性层出发,来解释一下DeepPoly的回传过程。
问题定义
对于线性层 →xm+1=Wm→xm+→bm ,已知
Wm,m−1l→xm−1+→bm,m−1l≤→xm≤Wm,m−1u→xm−1+→bm,m−1u求 →xm+1 的范围。
理论分析
→xm+1 可以表示成
Wm+1,m−1L→xm−1+→bm+1,m−1L≤→xm+1≤Wm+1,m−1U→xm−1+→bm+1,m−1U的形式,与输入域的形式类似,因此,只需要求出参数 Wi+1,i−1l,→bi+1,i−1l,Wi+1,i−1u,bi+1,i−1u 就可以合并两层。通过代入,可以得到参数的计算公式:
Wm+1,m−1L[x,z]=∑yWm[x,y](Wm,m−1L[y,z][Wm[x,y]>0]+Wm,m−1U[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中的
λ→xm−1≤→xm≤A→xm−1+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,Xm−1uvw 之间的关系:
Xm+1ijk=∑uvwWm+1,mijkuvw∑xyzWm,m−1uvwxyzXm−1x,j+v+y,k+w+z同时我们知道, Xm+1uvw,Xm−1uvw 之间的形式应该为:
Xm+1ijk=∑abcWm+1,m−1ijkabcXm−1a,j+b,k+c和线性层不同,这里的 Xm−1a,j+b,k+c 与 Xm−1x,j+v+y,k+w+z 并不是一一对应的,而是多对一的关系,因此无法直接得到新参数,不过可以通过对应关系,得到
Wm+1,m−1ijkabc=∑{Wm+1,mijkuvwWm,m−1uvwxyz|∀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层,和线性层类似,只需要对每个像素取 λ 即可。