python - 如何优化这个 NumPy 代码？

Question

以下代码是我的 Python 代码中的瓶颈:

def get_payoff(self,  actual, predicted):
    if abs(actual - 1.0) < 1e-5:  # if actual == 1
        if predicted < 0.5:
            return self.fn_payoff * (0.5 - predicted)
        elif predicted > 0.5:
            return self.tp_payoff * (predicted - 0.5)
        else:
            return 0
    else:
        if predicted < 0.5:
            return self.tn_payoff * (0.5 - predicted)
        elif predicted > 0.5:
            return self.fp_payoff * (predicted - 0.5)
        else:
            return 0

def get_total_payoff(self):
    total_payoff = 0
    for target_element, prediction_element in zip(np.nditer(self.target), np.nditer(predictions)):
        total_payoff += self.get_payoff(target_element, prediction_element)

fn_payoff、tp_payoff、tn_payoff 和 fp_payoff 均为 float 。 self.target 和 self.predictions 都是 numpy ndarray。

我认为有某种方法可以用某种 numpy 向量化替换 get_total_payoff 中的 for 循环，但我不知道如何处理 if/then 语句以正确进行向量化。

Answer 1

向量化根据条件使用不同表达式的函数的关键是使用 np.choose 。另外，在您的情况下， predict-0.5 和 0.5-predict 可以替换为 abs(predict-0.5)，再加上对predict==0.5 的情况(我猜测特殊处理是为了正确处理 NaN)。

import numpy as np

class A(object):
    def __init__(self):
        self.fn_payoff = 222.
        self.tn_payoff = 444.
        self.fp_payoff = 777.
        self.tp_payoff = 888.
        self.target = np.array([ 0.3, 1., 2. ])
        self.predictions = np.array([ 0.4, 0.5, 1.7 ])

    def get_payoff(self,  actual, predicted):
        if abs(actual - 1.0) < 1e-5:  # if actual == 1
            if predicted < 0.5:
                return self.fn_payoff * (0.5 - predicted)
            elif predicted > 0.5:
                return self.tp_payoff * (predicted - 0.5)
            else:
                return 0
        else:
            if predicted < 0.5:
                return self.tn_payoff * (0.5 - predicted)
            elif predicted > 0.5:
                return self.fp_payoff * (predicted - 0.5)
            else:
                return 0

    def get_total_payoff(self):
        total_payoff = 0
        for target_element, prediction_element in zip(np.nditer(self.target), np.nditer(self.predictions)):
            total_payoff += self.get_payoff(target_element, prediction_element)
        return total_payoff

    def get_total_payoff_VECTORIZED(self):
        actual_mask = np.abs(self.target - 1) < 1e-5
        predict_mask = self.predictions < 0.5
        payoff_n = np.choose(actual_mask, [ self.tn_payoff, self.fn_payoff ])
        payoff_p = np.choose(actual_mask, [ self.fp_payoff, self.tp_payoff ])
        payoff = np.choose(predict_mask, [ payoff_p, payoff_n ]) * abs(self.predictions-0.5)
        payoff[self.predictions==0.5] = 0
        return payoff.sum()

a = A()
print a.get_total_payoff()
=> 976.8
print a.get_total_payoff_VECTORIZED()
=> 976.8