如何向PyTorch LambdaLR调度器传递自定义函数的参数

你可以通过以下几种方式给自定义的func1传递额外参数（比如base），让它适配LambdaLR只接受单参数（epoch）函数的要求：

方法1：用lambda表达式包装 #

直接用lambda把额外参数固定，返回一个仅接收epoch的函数：

from torch import optim

def func1(epoch, base=10):
    return epoch / base

optimizer = optim.SGD(model.parameters(), lr=0.1)
# 这里指定base为20，替换默认值
scheduler = optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=lambda epoch: func1(epoch, base=20))

方法2：使用functools.partial偏置参数 #

利用functools.partial来预先绑定base参数，生成适配要求的函数：

from torch import optim
from functools import partial

def func1(epoch, base=10):
    return epoch / base

optimizer = optim.SGD(model.parameters(), lr=0.1)
# 绑定base参数为15
lr_lambda_func = partial(func1, base=15)
scheduler = optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=lr_lambda_func)

方法3：用闭包生成适配函数 #

写一个嵌套函数，外层接收base参数，内层返回仅处理epoch的逻辑：

from torch import optim

def make_lr_lambda(base=10):
    def lr_lambda(epoch):
        return epoch / base
    return lr_lambda

optimizer = optim.SGD(model.parameters(), lr=0.1)
# 生成指定base的函数
scheduler = optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=make_lr_lambda(base=5))

这几种方法都能让你在保留自定义函数扩展性的同时，顺利传递额外参数给LambdaLR。

内容的提问来源于stack exchange，提问作者Eypros