Title

Revisiting Global Statistics Aggregation for Improving Image Restoration
代码：https://github.com/megvii-research/tlsc
文章：http://arxiv.org/abs/2112.04491
Megvii

即插即用

Summary

本文提出了一个方法，Test-time Local Statistics Converter(TLSC)，不需要额外训练和finetune，就可以规避按patch和整图跑网络时的空间统计信息不一致问题。

Problem Statement

首先明确一下，什么是空间统计信息（spatial statistics）：记录空间上下文的信息，如HINet中用的IN，如SENet的Global Average Pooling。
项目中，我们经常会训练时分patch训或者random crop训，在测试时却是整图输入网络。按patch跑网络和整图跑网络，它们的空间统计信息是不一样的。（本文说自己是第一个指出这种不一致的，之前被广泛忽略。）

为了避免这种不一致，现在主要有两种做法：

• 训练和测试的时候都用patch
  这种做法，测试时需要做patch融合，容易导致边界伪影或者边界的不和谐
• 训练和测试的时候都用整图
  这种做法会导致效果的下降。（作者列了张表证明这种性能下降，我猜测是因为random crop带来了数据多样性，更利于效果）

Method(s)

作者从空间统计信息出发，解决这种batch和整图跑网络的不一致问题。

上图（b）展示了训练和测试的时候都用整图的做法，上图（c）展示了训练和测试的时候都用patch的做法。这两者的弊端前文都已经提到。
上图(d)展示了本文的做法，整图输入，其核心思想是，对于空间统计信息，从局部窗口获取，而不从整图获取。

关于计算量的说明：
如果我们从高宽为的窗口获取局部空间统计信息，理论上计算复杂度是#card=math&code=O%28HWK%7Bh%7DK%7Bw%7D%29&id=mccfJ)
但实际应用中，我们可以将其视作submatrix sum问题，使用prefix sum trick来解决，这样复杂度就变成了#card=math&code=O%28HW%29&id=kmikr)，和整图算的计算量一样。

具体使用：

1. 在SE Block中使用TLSC
   在应用average pooling时，由局部窗口获取结果，而不是global average pooling
2. 在IN中使用TLSC
   这个相对明了一点，因为IN中是要统计mean和std的。本来是基于HW去算，现在变成了基于算

留个印象，之后我会去看下代码。

Evaluation

能改善patch和整图跑网络的不一致问题。这种不一致越严重，改善效果越明显。

对于窗口的大小，和训练时的尺寸密切相关。但这里有个很有趣的现象，最佳的窗口尺寸可能要比训练时的patch尺寸要大一些。
作者给了个解释很合理：测试时，越大的窗口尺寸意味着信息越多，也意味着和训练时的gap越大，这是一种权衡，适量地大一些是个比较好的选择。

Criticism

作者提了一个很直观很符合理解但一直没人注意到的点子。这个点子朴素但实用，我觉得很多地方都可以尝试一下，试一试的成本也很小，我很喜欢这篇的idea。
他还发现了一个现象和我的认知一致：即测试时的尺寸比训练时的尺寸大一丢丢效果可能比尺寸完全保持一致好。

为了保证足够的篇幅，文章不够干练，好多没用的话。

Code

code理解上还有些问题，需要去补功课

local avg pool

class AvgPool2d(nn.Module):
    def __init__(self, kernel_size=None, base_size=None, auto_pad=True, fast_imp=False):
        super().__init__()
        self.kernel_size = kernel_size
        self.base_size = base_size
        self.auto_pad = auto_pad
        # only used for fast implementation
        self.fast_imp = fast_imp
        self.rs = [5,4,3,2,1]
        self.max_r1 = self.rs[0]
        self.max_r2 = self.rs[0]
    def extra_repr(self) -> str:
        return 'kernel_size={}, base_size={}, stride={}, fast_imp={}'.format(
            self.kernel_size, self.base_size, self.kernel_size, self.fast_imp
        )
    def forward(self, x):
        if self.kernel_size is None and self.base_size:
            if isinstance(self.base_size, int):
                self.base_size = (self.base_size, self.base_size)
            self.kernel_size = list(self.base_size)
            self.kernel_size[0] = x.shape[2]*self.base_size[0]//train_size[-2]
            self.kernel_size[1] = x.shape[3]*self.base_size[1]//train_size[-1]
            # only used for fast implementation
            self.max_r1 = max(1, self.rs[0]*x.shape[2]//train_size[-2])
            self.max_r2 = max(1, self.rs[0]*x.shape[3]//train_size[-1])
        if self.fast_imp:   # Non-equivalent implementation but faster
            h, w = x.shape[2:]
            if self.kernel_size[0]>=h and self.kernel_size[1]>=w:
                out = F.adaptive_avg_pool2d(x,1)
            else:
                r1 = [r for r in self.rs if h%r==0][0]
                r2 = [r for r in self.rs if w%r==0][0]
                # reduction_constraint
                r1 = min(self.max_r1, r1)
                r2 = min(self.max_r2, r2)
                s = x[:,:,::r1, ::r2].cumsum(dim=-1).cumsum(dim=-2)
                n, c, h, w = s.shape
                k1, k2 = min(h-1, self.kernel_size[0]//r1), min(w-1, self.kernel_size[1]//r2)
                out = (s[:,:,:-k1,:-k2]-s[:,:,:-k1,k2:]-s[:,:,k1:,:-k2]+s[:,:,k1:,k2:])/(k1*k2)
                out = torch.nn.functional.interpolate(out, scale_factor=(r1,r2))
        else:
            n, c, h, w = x.shape
            s = x.cumsum(dim=-1).cumsum_(dim=-2)
            s = torch.nn.functional.pad(s, (1,0,1,0)) # pad 0 for convenience
            k1, k2 = min(h, self.kernel_size[0]), min(w, self.kernel_size[1])
            s1, s2, s3, s4 = s[:,:,:-k1,:-k2],s[:,:,:-k1,k2:], s[:,:,k1:,:-k2], s[:,:,k1:,k2:]
            out = s4+s1-s2-s3
            out = out / (k1*k2)
        if self.auto_pad:
            n, c, h, w = x.shape
            _h, _w = out.shape[2:]
            # print(x.shape, self.kernel_size)
            pad2d = ((w - _w)//2, (w - _w + 1)//2, (h - _h) // 2, (h - _h + 1) // 2)
            out = torch.nn.functional.pad(out, pad2d, mode='replicate')
        return out

local instance norm

class LocalInstanceNorm2d(nn.Module):
    def __init__(self, num_features, eps=1e-5, momentum=0.1,
                 affine=False, track_running_stats=False):
        super().__init__()
        assert not track_running_stats
        self.affine = affine
        if self.affine:
            self.weight = nn.Parameter(torch.ones(num_features))
            self.bias = nn.Parameter(torch.zeros(num_features))
        else:
            self.register_parameter('weight', None)
            self.register_parameter('bias', None)
        self.avgpool = AvgPool2d()
        self.eps = eps
    def forward(self, input):
        mean_x = self.avgpool(input) # E(x)
        mean_xx = self.avgpool(torch.mul(input, input)) # E(x^2)
        mean_x2 = torch.mul(mean_x, mean_x) # (E(x))^2
        var_x = mean_xx - mean_x2 # Var(x) = E(x^2) - (E(x))^2
        mean = mean_x
        var = var_x
        input = (input - mean) / (torch.sqrt(var + self.eps))
        if self.affine:
            input = input * self.weight.view(1,-1, 1, 1) + self.bias.view(1,-1, 1, 1)
        return input