def fuse_bn(W, bn, eps=1e-5):
scale = bn.gamma.data() / np.sqrt(bn.running_var.data() + eps)
return (W * scale.reshape(-1, 1, 1, 1),
bn.beta.data() - bn.running_mean.data() * scale)
def fuse_block(blk):
c = blk.conv3.weight.shape[0]
W1 = np.zeros((c, c, 3, 3))
W1[:, :, 1, 1] = blk.conv1.weight.data().reshape(c, c)
I = np.zeros((c, c, 3, 3))
I[:, :, 1, 1] = np.eye(c)
W3, b3 = fuse_bn(blk.conv3.weight.data(), blk.bn3)
W1, b1 = fuse_bn(W1, blk.bn1)
W0, b0 = fuse_bn(I, blk.bn0)
fused = nn.Conv2D(c, kernel_size=3, padding=1)
fused.initialize()
fused(np.zeros((1, c, 3, 3))) # materialize the deferred parameters
fused.weight.data()[:] = W3 + W1 + W0
fused.bias.data()[:] = b3 + b1 + b0
return fused