如何通过设计模式简化PyTorch Lightning模型类的重复加载逻辑？

你可以用Mixin类或者类装饰器来复用这段load方法的逻辑，避免每个模型类重复编写代码：

方案一：使用Mixin类 #

Mixin是Python中复用类级代码的常用方式，你可以定义一个包含通用load方法的Mixin类，让所有模型类继承它：

from pytorch_lightning import LightningModule

class LightningModelLoadMixin:
    @classpath  # 保留Hydra需要的装饰器
    @classmethod
    def load(cls, **kwargs):
        if "checkpoint_file" in kwargs:
            # 调用当前类的load_from_checkpoint方法，适配不同模型类
            return cls.load_from_checkpoint(kwargs["checkpoint_file"])
        else:
            return cls(**kwargs)

# 模型类继承Mixin和LightningModule
class MyModel(LightningModelLoadMixin, LightningModule):
    def __init__(self, **kwargs):
        super().__init__()
        # 你的模型初始化逻辑

class AnotherModel(LightningModelLoadMixin, LightningModule):
    def __init__(self, **kwargs):
        super().__init__()
        # 另一个模型的初始化逻辑

所有继承了LightningModelLoadMixin的模型类都会自动拥有load方法，无需重复实现。

方案二：使用类装饰器 #

如果你不想修改模型类的继承链，可以用装饰器给每个模型类动态添加load方法：

from pytorch_lightning import LightningModule

def add_load_method(cls):
    @classpath  # 保留Hydra需要的装饰器
    @classmethod
    def load(cls_inner, **kwargs):
        if "checkpoint_file" in kwargs:
            return cls_inner.load_from_checkpoint(kwargs["checkpoint_file"])
        else:
            return cls_inner(**kwargs)
    cls.load = load
    return cls

# 用装饰器装饰模型类
@add_load_method
class MyModel(LightningModule):
    def __init__(self, **kwargs):
        super().__init__()
        # 你的模型初始化逻辑

@add_load_method
class AnotherModel(LightningModule):
    def __init__(self, **kwargs):
        super().__init__()
        # 另一个模型的初始化逻辑

注意事项 #

• 两种方案都用cls（或cls_inner）代替了硬编码的类名，确保调用的是当前类的方法，适配不同模型类的需求。
• 保留了@classpath装饰器，满足Hydra对统一实例化接口的要求。

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