Python：当函数做为参数时的技巧

我们之前在《Python技法3: 匿名函数、回调函数、高阶函数》中提到，可以通过lambda表达式来为函数设置默认参数，从而修改函数的参数个数:

import math
def distance(p1, p2):
    x1, y1 = p1
    x2, y2 = p2
    return math.hypot(x2 - x1, y2 - y1)
points = [(1, 2), (3, 4), (5, 6), (7, 8)]
pt = (4, 3)
points.sort(key=lambda p: distance(p, pt))
print(points)
# [(3, 4), (1, 2), (5, 6), (7, 8)]

下面我们在深度学习项目情境中展示下该手法的应用。

我们这是一个联邦学习项目，有多个client客户端，每个client中都有机器学习模型。我们现在有一份现有的祖传代码不能改动，该祖传代码中client类的初始化函数中需要传入模型类的初始化函数和优化器类的初始化函数（注意，不是模型的对象和优化器的对象），然后在client类的构造函数中完成模型对象的初始化和优化器对象的初始化。

client类的构造函数部分如下所示如下：

class Client(FederatedTrainingDevice):
    def __init__(self, model_fn, optimizer_fn, batch_size=128, train_frac=0.8):
        self.model = model_fn().to(device)
        self.optimizer = optimizer_fn(self.model.parameters())
        ...

但是，这样的代码就会面临一个问题，model和optimizer的初始化是需要超参数的，超参数如何传进去？这时，就可以用我们前面所说的lambda技巧来解决了：

client = Client(lambda: ConvNet(input_size=input_sz, num_classes=num_cls), lambda x : torch.optim.SGD(x, lr=0.1, momentum=0.9)

这一还有一个注意的地方，就是lambda函数是可以无参数的，故可以用于完成我们这里model_fn函数的替代。