Add super/norm layers in xcore

This commit is contained in:
D-X-Y
2021-05-06 16:38:58 +08:00
parent ff5e544240
commit 4c14c7b85b
6 changed files with 392 additions and 13 deletions

View File

@@ -10,21 +10,26 @@ __all__ = ["get_model"]
from xlayers.super_core import SuperSequential
from xlayers.super_core import SuperSimpleNorm
from xlayers.super_core import SuperLeakyReLU
from xlayers.super_core import SuperLinear
from xlayers.super_core import super_name2norm
from xlayers.super_core import super_name2activation
def get_model(config: Dict[Text, Any], **kwargs):
model_type = config.get("model_type", "simple_mlp")
if model_type == "simple_mlp":
act_cls = super_name2activation[kwargs["act_cls"]]
norm_cls = super_name2norm[kwargs["norm_cls"]]
mean, std = kwargs.get("mean", None), kwargs.get("std", None)
hidden_dim1 = kwargs.get("hidden_dim1", 200)
hidden_dim2 = kwargs.get("hidden_dim2", 100)
model = SuperSequential(
SuperSimpleNorm(kwargs["mean"], kwargs["std"]),
SuperLinear(kwargs["input_dim"], 200),
SuperLeakyReLU(),
SuperLinear(200, 100),
SuperLeakyReLU(),
SuperLinear(100, kwargs["output_dim"]),
norm_cls(mean=mean, std=std),
SuperLinear(kwargs["input_dim"], hidden_dim1),
act_cls(),
SuperLinear(hidden_dim1, hidden_dim2),
act_cls(),
SuperLinear(hidden_dim2, kwargs["output_dim"]),
)
else:
raise TypeError("Unkonwn model type: {:}".format(model_type))