Merge pull request #28 from GitHub-HongweiZhang/fix-deepfm

fix deepFM

Author: GitHub-HongweiZhang

prediction_flow/__init__.py

@@ -1 +1 @@
-__version__ = '0.1.3'
+__version__ = '0.1.4'

prediction_flow/pytorch/deepfm.py

@@ -134,22 +134,23 @@ def forward(self, x):
         for feature in self.features.number_features:
             embeddings.append(
                 self.embeddings[feature.name](
-                    x[feature.name].view(-1, 1)))
+                    x[feature.name].view(-1, 1)).unsqueeze(1))
         for feature in self.features.category_features:
             embeddings.append(
-                self.embeddings[feature.name](x[feature.name]))
+                self.embeddings[feature.name](x[feature.name]).unsqueeze(1))
         for feature in self.features.sequence_features:
             embeddings.append(
                 self._sequence_poolings[feature.name](
-                    self.embeddings[feature.name](x[feature.name])))
+                    self.embeddings[feature.name](x[feature.name])).unsqueeze(1))
 
         emb_concat = None
         if embeddings:
             emb_concat = torch.cat(embeddings, dim=1)
-
-        # fm
-        if self.fm:
-            final_layer_inputs.append(self.fm(emb_concat))
+            b, f, e = emb_concat.size()
+            # fm
+            if self.fm:
+                final_layer_inputs.append(self.fm(emb_concat))
+            emb_concat = emb_concat.view(b, f * e)
 
         # deep
         if self.mlp:

prediction_flow/pytorch/nn/fm.py

@@ -13,11 +13,14 @@
 class FM(nn.Module):
     """FM layer.
     """
-    def __init__(self):
+    def __init__(self, reduce_sum=True):
         super(FM, self).__init__()
+        self.reduce_sum = reduce_sum
 
     def forward(self, x):
-        sum_squared = torch.pow(torch.sum(x, dim=1), 2).unsqueeze(1)
-        squared_sum = torch.sum(torch.pow(x, 2), dim=1).unsqueeze(1)
-        second_order = 0.5 * (sum_squared - squared_sum)
-        return second_order
+        sum_squared = torch.pow(torch.sum(x, dim=1), 2)
+        squared_sum = torch.sum(torch.pow(x, 2), dim=1)
+        second_order = sum_squared - squared_sum
+        if self.reduce_sum:
+            output = torch.sum(second_order, dim=1, keepdim=True)
+        return 0.5 * output

prediction_flow/pytorch/nn/tests/test_fm.py

@@ -7,10 +7,12 @@
 def test_fm():
     fm = FM()
 
-    x = torch.as_tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
+    x = torch.as_tensor(
+        [[[1.0, 1.0, 1.0], [1.0, 2.0, 3.0]],
+         [[1.0, 1.0, 1.0], [4.0, 5.0, 6.0]]])
     actual = fm(x)
 
-    # 11.0 = 1 * 2 + 1 * 3 + 2 * 3
-    # 77.0 = 4 * 5 + 4 * 6 + 5 * 6
+    # 6.0 = 1 * 1 + 1 * 2 + 1 * 3
+    # 15.0 = 1 * 4 + 1 * 5 + 1 * 6
     np.testing.assert_array_almost_equal(
-        actual.numpy(), np.array([[11.0], [74.0]], dtype=np.float))
+        actual.numpy(), np.array([[6.0], [15.0]], dtype=np.float))