fix deepfm

apache · Jul 18, 2020 · fa151e5 · fa151e5
1 parent 380358c
commit fa151e5
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/submarine-sdk/pysubmarine/submarine/ml/pytorch/model/ctr/deepfm.py b/submarine-sdk/pysubmarine/submarine/ml/pytorch/model/ctr/deepfm.py
@@ -33,9 +33,9 @@ class _DeepFM(nn.Module):
     def __init__(self, num_fields, num_features, embedding_dim, out_features, hidden_units,
                  dropout_rates, **kwargs):
         super().__init__()
-        self.field_linear = FeatureLinear(num_features=num_features,
+        self.feature_linear = FeatureLinear(num_features=num_features,
                                         out_features=out_features)
-        self.field_embedding = FeatureEmbedding(num_features=num_features,
+        self.feature_embedding = FeatureEmbedding(num_features=num_features,
                                               embedding_dim=embedding_dim)
         self.pairwise_interaction = PairwiseInteraction()
         self.dnn = DNN(in_features=num_fields * embedding_dim,
@@ -48,8 +48,8 @@ def forward(self, feature_idx, feature_value):
         :param feature_idx: torch.LongTensor (batch_size, num_fields)
         :param feature_value: torch.LongTensor (batch_size, num_fields)
         """
-        emb = self.field_embedding(feature_idx, feature_value)  # (batch_size, num_fields, embedding_dim)
-        linear_logit = self.field_linear(feature_idx, feature_value)
+        emb = self.feature_embedding(feature_idx, feature_value)  # (batch_size, num_fields, embedding_dim)
+        linear_logit = self.feature_linear(feature_idx, feature_value)
         fm_logit = self.pairwise_interaction(emb)
         deep_logit = self.dnn(torch.flatten(emb, start_dim=1))