tzrec/models/dbmtl.py [51:80]:
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
        self._task_tower_cfgs = self._model_config.task_towers
        self.init_input()
        self.group_name = self.embedding_group.group_names()[0]
        feature_in = self.embedding_group.group_total_dim(self.group_name)

        self.bottom_mlp = None
        if self._model_config.HasField("bottom_mlp"):
            self.bottom_mlp = MLP(
                feature_in, **config_to_kwargs(self._model_config.bottom_mlp)
            )
            feature_in = self.bottom_mlp.output_dim()

        self.mmoe = None
        if self._model_config.HasField("expert_mlp"):
            self.mmoe = MMoEModule(
                in_features=feature_in,
                expert_mlp=config_to_kwargs(self._model_config.expert_mlp),
                num_expert=self._model_config.num_expert,
                num_task=len(self._task_tower_cfgs),
                gate_mlp=config_to_kwargs(self._model_config.gate_mlp)
                if self._model_config.HasField("gate_mlp")
                else None,
            )
            feature_in = self.mmoe.output_dim()

        self.task_mlps = nn.ModuleDict()
        for task_tower_cfg in self._task_tower_cfgs:
            if task_tower_cfg.HasField("mlp"):
                tower_mlp = MLP(feature_in, **config_to_kwargs(task_tower_cfg.mlp))
                self.task_mlps[task_tower_cfg.tower_name] = tower_mlp
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -



tzrec/models/dc2vr.py [46:75]:
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
        self._task_tower_cfgs = self._model_config.task_towers
        self.init_input()
        self.group_name = self.embedding_group.group_names()[0]
        feature_in = self.embedding_group.group_total_dim(self.group_name)

        self.bottom_mlp = None
        if self._model_config.HasField("bottom_mlp"):
            self.bottom_mlp = MLP(
                feature_in, **config_to_kwargs(self._model_config.bottom_mlp)
            )
            feature_in = self.bottom_mlp.output_dim()

        self.mmoe = None
        if self._model_config.HasField("expert_mlp"):
            self.mmoe = MMoEModule(
                in_features=feature_in,
                expert_mlp=config_to_kwargs(self._model_config.expert_mlp),
                num_expert=self._model_config.num_expert,
                num_task=len(self._task_tower_cfgs),
                gate_mlp=config_to_kwargs(self._model_config.gate_mlp)
                if self._model_config.HasField("gate_mlp")
                else None,
            )
            feature_in = self.mmoe.output_dim()

        self.task_mlps = nn.ModuleDict()
        for task_tower_cfg in self._task_tower_cfgs:
            if task_tower_cfg.HasField("mlp"):
                tower_mlp = MLP(feature_in, **config_to_kwargs(task_tower_cfg.mlp))
                self.task_mlps[task_tower_cfg.tower_name] = tower_mlp
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -