gpt4 book ai didi

python - 如何优化这个 MaxPool2d 实现

转载 作者:行者123 更新时间:2023-12-04 09:38:36 25 4
gpt4 key购买 nike

我做了一些 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

输入形状 = (batch_size, n_input_channels, h, w)

输出形状 = (batch_size, n_output_channels, h//kern_size, w//kern_size)

最佳答案

为清楚起见,我通过删除批量大小和 channel 维度简化了您的示例。
大部分时间都花在了 M.max() 的计算上.我创建了基准函数 update_output_b使用常量数组执行此循环。

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 完全矢量化的操作。如果可能,您应该始终优先使用 native 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 mean execution time: 0.332s (与前一个相比+55ms)

我在下面添加了矢量化代码。它的工作速度提高了大约 20 倍( update_output_2 mean execution time: 0.015s ),但它可能远非最佳。
@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沿 ech 窗口中的第二个轴,然后使用 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

    关于python - 如何优化这个 MaxPool2d 实现,我们在Stack Overflow上找到一个类似的问题: https://stackoverflow.com/questions/62433704/

    25 4 0
    Copyright 2021 - 2024 cfsdn All Rights Reserved 蜀ICP备2022000587号
    广告合作:1813099741@qq.com 6ren.com