Module重构讨论

[xiezipeng-ML]

关于解决libai中参数层层传递的问题讨论，主要思路是让内部Module直接获取参数，不通过外部传递：
简单写了一个demo，可以直接放到libai下跑。

创建一个ModuleBase基类：

from omegaconf import DictConfig
import oneflow as flow

from libai.config import LazyCall, configurable
from libai.models import build_model


cfg = dict(
    in_dim = 1,
    out_dim = 2,
    act = "gelu_tanh"
)
cfg = DictConfig(cfg)

class ModuleBase(flow.nn.Module):
    def __init__(self, cfg=None):
        super().__init__()
        self.cfg = cfg

class MLP(ModuleBase):
    def __init__(
        self, 
        in_dim,
        out_dim,
        cfg=None
    ):
        super().__init__(cfg)
        self.a = in_dim
        self.b = out_dim
        self.act = cfg.act

class Transformer(ModuleBase):
    def __init__(
        self,
        in_dim,
        out_dim,
        cfg=None
    ):
        super().__init__(cfg)
        self.mlp = MLP(
            in_dim,
            out_dim,
            cfg=cfg
        )

class BertModel(ModuleBase):
    @configurable
    def __init__(
        self, 
        in_dim,
        out_dim,
        cfg=None
    ):
        super().__init__(cfg)
        self.transformer = Transformer(
            in_dim,
            out_dim,
            cfg=cfg
        )
    
    @classmethod
    def from_config(cls, cfg):
        return {
            "in_dim": cfg.in_dim,
            "out_dim": cfg.out_dim,
            "cfg": cfg,
        }

bert_model = LazyCall(BertModel)(cfg=cfg)
bert = build_model(bert_model)

print(bert.transformer.mlp.act)    # output: gelu_tanh

ModuleBase的方案的代价是需要每个layer和model继承，然后多出一个cfg parameter。但是现在暂时感觉ModuleBase的存在用处不大（还需要讨论），所以下面是不用ModuleBase的demo，这个方案的代价只是多出一个cfg parameter。

from omegaconf import DictConfig
import oneflow as flow

from libai.config import LazyCall, configurable
from libai.models import build_model


cfg = dict(
    in_dim = 1,
    out_dim = 2,
    act = "gelu_tanh"
)
cfg = DictConfig(cfg)


class MLP(flow.nn.Module):
    def __init__(
        self, 
        in_dim,
        out_dim,
        cfg=None
    ):
        super().__init__()
        self.a = in_dim
        self.b = out_dim
        self.act = cfg.act


class Transformer(flow.nn.Module):
    def __init__(
        self,
        in_dim,
        out_dim,
        cfg=None
    ):
        super().__init__()
        self.mlp = MLP(
            in_dim,
            out_dim,
            cfg=cfg
        )


class BertModel(flow.nn.Module):
    @configurable
    def __init__(
        self, 
        in_dim,
        out_dim,
        cfg=None
    ):
        super().__init__()
        self.transformer = Transformer(
            in_dim,
            out_dim,
            cfg=cfg
        )
    
    @classmethod
    def from_config(cls, cfg):
        return {
            "in_dim": cfg.in_dim,
            "out_dim": cfg.out_dim,
            "cfg": cfg,
        }


bert_model = LazyCall(BertModel)(cfg=cfg)
bert = build_model(bert_model)

print(bert.transformer.mlp.act)