如何优化这个MaxPool2d实现

Question

我实现了一些MaxPool2d的代码（与PyTorch进行了比较，能够正确运行）。但是，在对MNIST数据集进行测试时，这个函数（updateOutput）需要很长时间才能完成。如何使用NumPy优化此代码？

class MaxPool2d(Module):
    def __init__(self, kernel_size):
        super(MaxPool2d, self).__init__()
        self.kernel_size = kernel_size
        self.gradInput = None

    def updateOutput(self, input):
        #print("MaxPool updateOutput")
        #start_time = time.time()
        kernel = self.kernel_size
        poolH = input.shape[2] // kernel
        poolW = input.shape[3] // kernel
        self.output = np.zeros((input.shape[0], 
                                input.shape[1], 
                                poolH,
                                poolW))
        self.index = np.zeros((input.shape[0],
                                    input.shape[1],
                                    poolH,
                                    poolW,
                                    2), 
                                    dtype='int32')

        for i in range(input.shape[0]):
            for j in range(input.shape[1]):
                for k in range(0, input.shape[2] - kernel+1, kernel):
                    for m in range(0, input.shape[3] - kernel+1, kernel):
                        M = input[i, j, k : k+kernel, m : m+kernel]
                        self.output[i, j, k // kernel, m // kernel] = M.max()
                        self.index[i, j, k // kernel, m // kernel] = np.array(np.unravel_index(M.argmax(), M.shape)) + np.array((k, m))

        #print(f"time: {time.time() - start_time:.3f}s")
        return self.output

输入形状 = (批量大小，输入通道数，高度，宽度)

输出形状 = (批量大小，输出通道数，高度 // 卷积核大小，宽度 // 卷积核大小)

- annaFerdsf

Answer 1

为了清晰起见，我删除了批处理大小和通道维度，简化了您的示例。 大部分时间花在了M.max()的计算上。我创建了基准函数update_output_b来使用全1常数数组进行此循环。

import time
import numpy as np

def timeit(cycles):
    def timed(func):
        def wrapper(*args, **kwargs):
            start_t = time.time()
            for _ in range(cycles):
                func(*args, **kwargs)
            t = (time.time() - start_t) / cycles
            print(f'{func.__name__} mean execution time: {t:.3f}s')

        return wrapper
    return timed

@timeit(100)
def update_output_b(input, kernel):
    ones = np.ones((kernel, kernel))

    pool_h = input.shape[0] // kernel
    pool_w = input.shape[1] // kernel
    output = np.zeros((pool_h, pool_w))

    for i in range(0, input.shape[0] - kernel + 1, kernel):
        for j in range(0, input.shape[1] - kernel + 1, kernel):
            output[i // kernel, j // kernel] = ones.max()

    return output

in_arr = np.random.rand(3001, 200)
update_output_b(in_arr, 3)

它的输出是update_output_b mean execution time: 0.277s，因为它没有完全使用numpy向量化操作。在可能的情况下，您应该始终优先使用本地的numpy函数而不是循环。

此外，使用输入数组的切片会减慢执行速度，因为访问连续的内存在大多数情况下更快。

@timeit(100)
def update_output_1(input, kernel):
    pool_h = input.shape[0] // kernel
    pool_w = input.shape[1] // kernel
    output = np.zeros((pool_h, pool_w))

    for i in range(0, input.shape[0] - kernel + 1, kernel):
        for j in range(0, input.shape[1] - kernel + 1, kernel):
            M = input[i : i + kernel, j : j + kernel]
            output[i // kernel, j // kernel] = M.max()

    return output

update_output_1(in_arr, 3)

代码返回update_output_1平均执行时间：0.332秒（与之前相比增加了55毫秒）

我添加了下面的向量化代码。它的速度快了约20倍（update_output_2平均执行时间：0.015秒），但可能远非最优。

@timeit(100)
def update_output_2(input, kernel):
    pool_h = input.shape[0] // kernel
    pool_w = input.shape[1] // kernel
    input_h = pool_h * kernel
    input_w = pool_w * kernel

    # crop input
    output = input[:input_h, :input_w]
    # calculate max along second axis
    output = output.reshape((-1, kernel))
    output = output.max(axis=1)
    # calculate max along first axis
    output = output.reshape((pool_h, kernel, pool_w))
    output = output.max(axis=1)

    return output

update_output_2(in_arr, 3)

它通过以下3个步骤生成输出：

• 将输入剪裁为可被内核整除的大小
• 沿第二轴计算最大值（减少第一轴中切片之间的偏移量）
• 沿第一轴计算最大值

编辑： 我已经添加了用于检索最大值索引的修改。但是，您应该检查索引算术，因为我只在随机数组上测试过它。

它在每个窗口沿第二轴计算output_indices，然后使用output_indices_selector选择第二轴上的最大值。

def update_output_3(input, kernel):
    pool_h = input.shape[0] // kernel
    pool_w = input.shape[1] // kernel
    input_h = pool_h * kernel
    input_w = pool_w * kernel

    # crop input
    output = input[:input_h, :input_w]

    # calculate max along second axis
    output_tmp = output.reshape((-1, kernel))
    output_indices = output_tmp.argmax(axis=1)
    output_indices += np.arange(output_indices.shape[0]) * kernel
    output_indices = np.unravel_index(output_indices, output.shape)
    output_tmp = output[output_indices]

    # calculate max along first axis
    output_tmp = output_tmp.reshape((pool_h, kernel, pool_w))
    output_indices_selector = (kernel * pool_w * np.arange(pool_h).reshape(pool_h, 1))
    output_indices_selector = output_indices_selector.repeat(pool_w, axis=1)
    output_indices_selector += pool_w * output_tmp.argmax(axis=1)
    output_indices_selector += np.arange(pool_w)
    output_indices_selector = output_indices_selector.flatten()

    output_indices = (output_indices[0][output_indices_selector],
                      output_indices[1][output_indices_selector])
    output = output[output_indices].reshape(pool_h, pool_w)

    return output, output_indices

- Jakub Gąsiewski

---

谢谢，你的代码很好用。但是我怎样才能保留最大元素的索引呢？我需要它们用于反向操作。 - annaFerdsf

你能告诉我如何更好地找到最大元素的索引吗？我尝试使用np.unravel_index，但是没有任何结果出现:( - annaFerdsf

你可以查看我发布的更新。这是一个相当晦涩的解决方案，所以最好的理解方法是在简单的示例上进行白板调试。 - Jakub Gąsiewski