如何自动将PyTorch（Lightning）类初始化函数的所有输入参数记录到MLFlow？

这个痛点我太懂了！手动把每个参数塞进mlflow.log_params()不仅麻烦，还很容易漏写或者写错参数名，尤其是参数多的时候。好在我们可以用Python的inspect模块来自动搞定这件事，下面给你几个实用的方案：

方案一：在__init__里直接自动提取参数 #

这种方法适合单个类使用，直接在初始化函数里加一段代码就能自动收集所有参数：

import inspect
import mlflow
from typing import Optional, Union

class YourLightningModel:
    def __init__(self, model_name: str, num_classes: int, device: str = 'cuda:0', learning_rate: float = 5e-5,
                 do_layer_freeze: bool = True, extra_class_layers: Optional[Union[int, list]] = None,
                 fine_tune_dropout_rate: float = 0):
        # 先完成你的初始化逻辑（如果需要把参数绑定到self的话）
        self.model_name = model_name
        self.num_classes = num_classes
        self.device = device
        self.learning_rate = learning_rate
        self.do_layer_freeze = do_layer_freeze
        self.extra_class_layers = extra_class_layers
        self.fine_tune_dropout_rate = fine_tune_dropout_rate
        
        # 自动提取所有初始化参数
        sig = inspect.signature(self.__init__)
        param_dict = {}
        for param_name in sig.parameters:
            if param_name != 'self':  # 跳过self参数
                param_dict[param_name] = getattr(self, param_name)
        
        # 一键记录到MLFlow
        mlflow.log_params(param_dict)

如果不想把所有参数都绑定到self，也可以直接从局部变量里提取：

def __init__(self, model_name: str, num_classes: int, device: str = 'cuda:0', learning_rate: float = 5e-5,
             do_layer_freeze: bool = True, extra_class_layers: Optional[Union[int, list]] = None,
             fine_tune_dropout_rate: float = 0):
    # 直接从局部变量获取参数（刚传入时的变量都在这里）
    local_vars = locals().copy()
    sig = inspect.signature(self.__init__)
    # 只保留__init__定义的参数，排除self
    param_dict = {name: local_vars[name] for name in sig.parameters if name != 'self'}
    
    mlflow.log_params(param_dict)
    
    # 继续你的初始化逻辑
    ...

方案二：写一个通用装饰器（推荐） #

如果你有多个PyTorch Lightning类都需要自动记录参数，写一个装饰器可以避免重复代码，复用性拉满：

import inspect
import mlflow
from functools import wraps

def log_init_params(func):
    @wraps(func)
    def wrapper(self, *args, **kwargs):
        # 先执行原初始化函数
        func(self, *args, **kwargs)
        
        # 自动绑定所有传入的参数（包括默认值）
        sig = inspect.signature(func)
        bound_args = sig.bind(self, *args, **kwargs)
        bound_args.apply_defaults()
        
        # 过滤掉self，整理成参数字典
        param_dict = {k: v for k, v in bound_args.arguments.items() if k != 'self'}
        
        # 记录到MLFlow
        mlflow.log_params(param_dict)
    return wrapper

然后只需要在你的类的__init__上加上这个装饰器就行：

class YourLightningModel:
    @log_init_params
    def __init__(self, model_name: str, num_classes: int, device: str = 'cuda:0', learning_rate: float = 5e-5,
                 do_layer_freeze: bool = True, extra_class_layers: Optional[Union[int, list]] = None,
                 fine_tune_dropout_rate: float = 0):
        # 你的初始化逻辑该怎么写就怎么写，不用管参数记录
        ...

这个装饰器的好处是：不管你是用位置参数还是关键字参数传值，甚至参数有默认值没传的情况，它都能正确捕获所有参数的最终值，非常灵活。

补充说明 #

为什么Scikit-learn可以直接用autolog，而PyTorch Lightning不行？主要是因为Scikit-learn的模型有统一的API规范，MLFlow可以很方便地自动识别参数；而PyTorch Lightning的类结构自由度很高，没有统一的参数暴露方式，所以需要我们手动用inspect来提取，但上面的方法完全可以达到和autolog一样的效果。

备注：内容来源于stack exchange，提问作者illan