Pytorch forward 函数在computer vision中的加速

最近在研究EDSR代码（项目地址：https://github.com/sksq96/pytorch-summary）的时候看到了forward_chop function，该参数的help写的是

parser.add_argument('--chop', action='store_true',
                    help='use memory-efficient forward')

于是对其怎样加速进行了学习，将思路和大家分享一下。
这部分代码黏贴如下：

首先是开始部分的代码，我们假设最开始的输入是 tuple(tensor(1, 3, 678, 1020)) (里面的是size， 表示 1 x 3 x 678 x 1020)

def forward_chop(self, *args, shave=10, min_size=160000):
    n_GPUs = min(self.args.n_GPUs, 4)
    h, w = args[0].size()[-2:]
    top, left = slice(0, h//2 + shave), slice(0, w//2 + shave)
    bottom, right = slice(h - h//2 - shave, h), slice(w - w//2 - shave, w)
    x_chops = [torch.cat([
        a[..., top, left],
        a[..., top, right],
        a[..., bottom, left],
        a[..., bottom, right]
    ]) for a in args]

这部分将输入的一张图片拆分成了左上、左下、右上、右下四个部分并按照batch dim进行了拼接。
这样得到 x_chops = [tensor(4, 3, 349, 520)]，这里长宽各增加了一个shave * 2
接下来是对图片的大小进行判断：

	y_chops = []
       if h * w < 4 * min_size:
           for i in range(0, 4, n_GPUs):
               x = [x_chop[i:(i + n_GPUs)] for x_chop in x_chops]
               y = P.data_parallel(self.model, *x, device_ids=range(n_GPUs))
               if not isinstance(y, list):
                   y = [y]
               if not y_chops:
                   y_chops = [[c for c in _y.chunk(n_GPUs, dim=0)] for _y in y]
               else:
                   for y_chop, _y in zip(y_chops, y):
                       y_chop.extend(_y.chunk(n_GPUs, dim=0))
       else:
           for x in zip(*x_chops):
               p = tuple(_x.unsqueeze(0) for _x in x)
               y = self.forward_chop(*p, shave=shave, min_size=min_size)
               if not isinstance(y, list):
                   y = [y]
               if not y_chops:
                   y_chops = [[_y] for _y in y]
               else:
                   for y_chop, _y in zip(y_chops, y):
                       y_chop.append(_y)

这里由于 678 x 1020 - 4 * 160000 = 51560 > 0，所以走的 else语句
这里得到 p = (1, 3, 349, 520) 然后再作为输入调用forward_chop。前面还是一样的，只不过在这里走的是if语句。注意此时if语句中的x 的大小是和n_GPUs 有关的，也就是说如果你的n_GPUs == 1，那么你 x 的size为(1, 3, 184, 270)， 如果是2，就是(2, 3, 184, 270)，这样也就将多个batch 作为并行输入到你的GPU device 上加快速度。这里由于我的是1个GPU，整个for循环运行完之后可以得到y_chops = [[tensor(1, 3, 386, 540), tensor(1, 3, 386, 540), tensor(1, 3, 386, 540), tensor(1, 3, 386, 540)]]。
接下来就是恢复原大小了：

	h *= self.args.scale
    w *= self.args.scale
    top, left = slice(0, h//2), slice(0, w//2)
    bottom, right = slice(h - h//2, h), slice(w - w//2, w)
    bottom_r, right_r = slice(h//2 - h, None), slice(w//2 - w, None)
    b, c = y_chops[0][0].size()[:-2]
    y = [y_chop[0].new(b, c, h, w) for y_chop in y_chops]
    for y_chop, _y in zip(y_chops, y):
        _y[..., top, left] = y_chop[0][..., top, left]
        _y[..., top, right] = y_chop[1][..., top, right_r]
        _y[..., bottom, left] = y_chop[2][..., bottom_r, left]
        _y[..., bottom, right] = y_chop[3][..., bottom_r, right_r]
    if len(y) == 1:
        y = y[0]
    return y

可以看到 通过 _y 将拆分的四个部分再按位置拼接回原来的大小，最后得到一个tensor(1, 3, 698, 1040)和原输入大小相同的tensor。
总结一下，在这里通过将一张大的图片分成四个部分并行做输入进行测试，再将分别测试得到的输出拼接成原大小得到原图直接做forward的结果，相当于是减少了长和宽增加了batch_size。

以上是个人观点，如有不对还请大佬指出。