python partial、map的使用

背景

看mmdetection源码的时候，发现经过FPN的多层输出都会通过multi_apply(func, *args, **kwargs)计算结果，而对FPN的每一层的操作是通过func实现的。这让我好奇multi_apply究竟做了什么？

def multi_apply(func, *args, **kwargs):
    """Apply function to a list of arguments.

    Note:
        This function applies the ``func`` to multiple inputs and
        map the multiple outputs of the ``func`` into different
        list. Each list contains the same type of outputs corresponding
        to different inputs.

    Args:
        func (Function): A function that will be applied to a list of
            arguments

    Returns:
        tuple(list): A tuple containing multiple list, each list contains \
            a kind of returned results by the function
    """
    pfunc = partial(func, **kwargs) if kwargs else func
    map_results = map(pfunc, *args)
    return tuple(map(list, zip(*map_results)))

看注释这个函数返回的是个tuple，里面包含多个list，每个list是func函数的其中一个返回值的list。

partial

partial是Python模块functools中定义的一个函数，用于在固定某些参数的情况下执行某函数。
例如我们写了一个计算数值幂的函数

def power(a, b):
	print(a**b)

当前需求仅需要计算某个数的平方值，并且要大量的使用，这个时候我们当然可以使用

power(3, b=2)
power(4, b=2)

来计算，但是这样代码是冗余的，不够优雅，这个时候可以利用partial来解决

def power(a, b):
	print(a**b)

if __name__ == '__main__':
	func = partial(power, b=2)
	func(3)
	func(4)

out:
9
16
partial的作用可以理解为重新创造一个固定某些参数的函数以方便使用。

map

map函数是python内建函数，其作用是，将func函数依次作用于Iterable的每个元素上，返回一个新的Iterable对象

results = map(func, [3,4])
print(list(results))

out:
[9,16]

那么multi_apply()的作用也就显而易见了，其实就是注释解释的那样，将输入func中计算得到的值分别保存到几个list中，并整体存在一个tuple中

def compute(a, b):
	return a+b, a-b, a*b

if __name__ == '__main__':
	compute_func = partial(compute, b=3)
	results_list = map(compute_func, [4,5,6,7])
	print(tuple(map(list, zip(*results_list))))

([7, 8, 9, 10], [1, 2, 3, 4], [12, 15, 18, 21])