在 numpy 中随机索引多维 arrays

Question

How to change this code avoid using python for-loop, but numpy function

h, w = 2, 2
im = np.random.randint(255, size=(h, w, 3))
index = np.random.randint(3, size=(h, w))
number = np.random.randint(255, size=(h, w))
for i in range(h):
    for j in range(w):
        im[i, j, index[i, j]] += number[i, j]

Answer 1

The code below does the same as your original code, but avoids using for loops. Why you'd want to do this is unclear however, since I would consider this solution decidedly worse than the original.

from timeit import timeit
import numpy as np

h, w = 20, 20
im = np.random.randint(255, size=(h, w, 3))


def increase_random(x):
    result = np.copy(x)
    result[np.random.randint(3)] += np.random.randint(255)
    return result


def loops():
    index = np.random.randint(3, size=(h, w))
    number = np.random.randint(255, size=(h, w))
    for i in range(h):
        for j in range(w):
            im[i, j, index[i, j]] += number[i, j]


def vectorized():
    irv(im)


irv = np.vectorize(increase_random, signature='(n)->(n)')

print(timeit(vectorized, number=10))
print(timeit(loops, number=10))

I've added some time measurements to show that, in this case, vectorization does nothing to improve performance. On my machine, the loops code is about 25x faster.

However, if the operation you're performing is simpler or complex but easier to optimize, it would likely benefit from vectorization. It just so happens that your example is not very likely to benefit from it, while the loop is fairly small and efficient.

Answer 2

With an example as small as yours you will be having a hard time speeding up the loopy code. numpy pays for the overhead it incurs only from a certain problem size upwards:

The plot shows execution time of original loopy code ( OP ) vs vectorized code ( pp ) over total number of pixels wxh .

It was generated using:

from simple_benchmark import BenchmarkBuilder, MultiArgument
import numpy as np
from scipy.misc import face

B = BenchmarkBuilder()

@B.add_function()
def OP(im,index,number):
    im = im.copy()
    h,w,_ = im.shape
    for i in range(h):
        for j in range(w):
            im[i, j, index[i, j]] += number[i, j]
    return im

@B.add_function()
def pp(im,index,number):
    im = im.copy()
    h,w,_ = im.shape
    h,w = np.ogrid[:h,:w]
    im[h,w,index] += number
    return im

@B.add_arguments('#pixels')
def argument_provider():
    im = face()
    h,w,_ = im.shape
    mh,mw = h//2,w//2
    for exp in range(-8,1):
        fr = 2.**exp
        dh,dw = int(fr*mh),int(fr*mw)
        index = np.random.randint(3, size=(2*dh, 2*dw))
        number = np.random.randint(255, size=(2*dh, 2*dw),dtype=im.dtype)
        yield 4*dh*dw,MultiArgument([im[mh-dh:mh+dh,mw-dw:mw+dw],index,number])

r = B.run()
r.plot()

import pylab
pylab.savefig('randomchannel.png')