Add PretrainedEmbedding to support more usages of pretrained embeddings

benchmark/run_expid.py

@@ -74,10 +74,10 @@
     gc.collect()
 
     logging.info('******** Test evaluation ********')
-    test_gen = H5DataLoader(feature_map, stage='test', **params).make_iterator()
     test_result = {}
-    if test_gen:
-      test_result = model.evaluate(test_gen)
+    if params["test_data"]:
+        test_gen = H5DataLoader(feature_map, stage='test', **params).make_iterator()
+        test_result = model.evaluate(test_gen)
 
     result_filename = Path(args['config']).name.replace(".yaml", "") + '.csv'
     with open(result_filename, 'a+') as fw:

fuxictr/pytorch/layers/embeddings/__init__.py

@@ -1,2 +1,2 @@
 from .feature_embedding import *
-
+from .pretrained_embedding import *

fuxictr/pytorch/layers/embeddings/feature_embedding.py

@@ -17,10 +17,10 @@
 
 import torch
 from torch import nn
-import h5py
 import os
 import numpy as np
 from collections import OrderedDict
+from .pretrained_embedding import PretrainedEmbedding
 from fuxictr.pytorch.torch_utils import get_initializer
 from fuxictr.pytorch import layers
 
@@ -69,11 +69,11 @@ def __init__(self,
         for feature, feature_spec in self._feature_map.features.items():
             if self.is_required(feature):
                 if not (use_pretrain and use_sharing) and embedding_dim == 1:
-                    feat_emb_dim = 1 # in case for LR
+                    feat_dim = 1 # in case for LR
                     if feature_spec["type"] == "sequence":
                         self.feature_encoders[feature] = layers.MaskedSumPooling()
                 else:
-                    feat_emb_dim = feature_spec.get("embedding_dim", embedding_dim)
+                    feat_dim = feature_spec.get("embedding_dim", embedding_dim)
                     if feature_spec.get("feature_encoder", None):
                         self.feature_encoders[feature] = self.get_feature_encoder(feature_spec["feature_encoder"])
 
@@ -83,31 +83,24 @@ def __init__(self,
                     continue
 
                 if feature_spec["type"] == "numeric":
-                    self.embedding_layers[feature] = nn.Linear(1, feat_emb_dim, bias=False)
-                elif feature_spec["type"] == "categorical":
-                    padding_idx = feature_spec.get("padding_idx", None)
-                    embedding_matrix = nn.Embedding(feature_spec["vocab_size"], 
-                                                    feat_emb_dim, 
-                                                    padding_idx=padding_idx)
-                    if use_pretrain and "pretrained_emb" in feature_spec:
-                        embedding_matrix = self.load_pretrained_embedding(embedding_matrix,
-                                                                          feature_map, 
-                                                                          feature, 
-                                                                          freeze=feature_spec["freeze_emb"],
-                                                                          padding_idx=padding_idx)
-                    self.embedding_layers[feature] = embedding_matrix
-                elif feature_spec["type"] == "sequence":
-                    padding_idx = feature_spec.get("padding_idx", None)
-                    embedding_matrix = nn.Embedding(feature_spec["vocab_size"], 
-                                                    feat_emb_dim, 
-                                                    padding_idx=padding_idx)
+                    self.embedding_layers[feature] = nn.Linear(1, feat_dim, bias=False)
+                elif feature_spec["type"] in ["categorical", "sequence"]:
                     if use_pretrain and "pretrained_emb" in feature_spec:
-                        embedding_matrix = self.load_pretrained_embedding(embedding_matrix, 
-                                                                          feature_map, 
-                                                                          feature,
-                                                                          freeze=feature_spec["freeze_emb"],
-                                                                          padding_idx=padding_idx)
-                    self.embedding_layers[feature] = embedding_matrix
+                        pretrained_path = os.path.join(feature_map.data_dir, 
+                                                       feature_spec["pretrained_emb"])
+                        pretrain_dim = feature_spec.get("pretrain_dim", feat_dim)
+                        pretrain_usage = feature_spec.get("pretrain_usage", "init")
+                        self.embedding_layers[feature] = PretrainedEmbedding(feature,
+                                                                             feature_spec,
+                                                                             pretrained_path,
+                                                                             feat_dim,
+                                                                             pretrain_dim,
+                                                                             pretrain_usage)
+                    else:
+                        padding_idx = feature_spec.get("padding_idx", None)
+                        self.embedding_layers[feature] = nn.Embedding(feature_spec["vocab_size"], 
+                                                                      feat_dim, 
+                                                                      padding_idx=padding_idx)
         self.reset_parameters()
 
     def get_feature_encoder(self, encoder):
@@ -148,24 +141,6 @@ def is_required(self, feature):
         else:
             return True
 
-    def get_pretrained_embedding(self, pretrained_path, feature_name):
-        with h5py.File(pretrained_path, 'r') as hf:
-            embeddings = hf[feature_name][:]
-        return embeddings
-
-    def load_pretrained_embedding(self, embedding_matrix, feature_map, feature_name, freeze=False, padding_idx=None):
-        pretrained_path = os.path.join(feature_map.data_dir, feature_map.features[feature_name]["pretrained_emb"])
-        embeddings = self.get_pretrained_embedding(pretrained_path, feature_name)
-        if padding_idx is not None:
-            embeddings[padding_idx] = np.zeros(embeddings.shape[-1])
-        assert embeddings.shape[-1] == embedding_matrix.embedding_dim, \
-            "{}\'s embedding_dim is not correctly set to match its pretrained_emb shape".format(feature_name)
-        embeddings = torch.from_numpy(embeddings).float()
-        embedding_matrix.weight = torch.nn.Parameter(embeddings)
-        if freeze:
-            embedding_matrix.weight.requires_grad = False
-        return embedding_matrix
-
     def dict2tensor(self, embedding_dict, feature_list=[], feature_source=[], feature_type=[], flatten_emb=False):
         if type(feature_source) != list:
             feature_source = [feature_source]
@@ -214,7 +189,3 @@ def forward(self, inputs, feature_source=[], feature_type=[]):
                     embeddings = self.feature_encoders[feature](embeddings)
                 feature_emb_dict[feature] = embeddings
         return feature_emb_dict
-
-
-
-

fuxictr/pytorch/layers/embeddings/pretrained_embedding.py

@@ -0,0 +1,86 @@
+# =========================================================================
+# Copyright (C) 2023. Huawei Technologies Co., Ltd. All rights reserved.
+# 
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# =========================================================================
+
+
+import torch
+from torch import nn
+import h5py
+import os
+import numpy as np
+
+
+class PretrainedEmbedding(nn.Module):
+    def __init__(self,
+                 feature_name,
+                 feature_spec,
+                 pretrained_path,
+                 embedding_dim,
+                 pretrain_dim,
+                 pretrain_usage="init"):
+        """
+        Fusion pretrained embedding with ID embedding
+        :param: fusion_type: init/sum/concat
+        """
+        super().__init__()
+        assert pretrain_usage in ["init", "sum", "concat"]
+        self.pretrain_usage = pretrain_usage
+        padding_idx = feature_spec.get("padding_idx", None)
+        embedding_matrix = nn.Embedding(feature_spec["vocab_size"],
+                                        pretrain_dim,
+                                        padding_idx=padding_idx)
+        self.pretrain_embedding = self.load_pretrained_embedding(embedding_matrix,
+                                                                 pretrained_path,
+                                                                 feature_name,
+                                                                 freeze=feature_spec["freeze_emb"],
+                                                                 padding_idx=padding_idx)
+        if pretrain_usage != "init":
+            self.id_embedding = nn.Embedding(feature_spec["vocab_size"],
+                                             embedding_dim,
+                                             padding_idx=padding_idx)
+        if pretrain_usage == "sum" and embedding_dim != pretrain_dim:
+            self.proj_W = nn.Linear(pretrain_dim, embedding_dim)
+        else:
+            self.proj_W = None
+
+    def get_pretrained_embedding(self, pretrained_path, feature_name):
+        with h5py.File(pretrained_path, 'r') as hf:
+            embeddings = hf[feature_name][:]
+        return embeddings
+
+    def load_pretrained_embedding(self, embedding_matrix, pretrained_path, feature_name, freeze=False, padding_idx=None):
+        embeddings = self.get_pretrained_embedding(pretrained_path, feature_name)
+        if padding_idx is not None:
+            embeddings[padding_idx] = np.zeros(embeddings.shape[-1])
+        assert embeddings.shape[-1] == embedding_matrix.embedding_dim, \
+            "{}\'s pretrain_dim is not correct.".format(feature_name)
+        embeddings = torch.from_numpy(embeddings).float()
+        embedding_matrix.weight = torch.nn.Parameter(embeddings)
+        if freeze:
+            embedding_matrix.weight.requires_grad = False
+        return embedding_matrix
+
+    def forward(self, inputs):
+        feature_emb = self.pretrain_embedding(inputs)
+        if self.pretrain_usage != "init":
+            id_emb = self.id_embedding(inputs)
+            if self.pretrain_usage == "sum":
+                if self.proj_W is not None:
+                    feature_emb = self.proj_W(feature_emb) + id_emb
+                else:
+                    feature_emb += id_emb
+            if self.pretrain_usage == "concat":
+                feature_emb = torch.cat([feature_emb, id_emb], dim=-1)
+        return feature_emb

fuxictr/version.py

@@ -1 +1 @@
-__version__="2.0.4"
+__version__="2.1.0"

setup.py

@@ -5,9 +5,9 @@
 
 setuptools.setup(
     name="fuxictr",
-    version="2.0.4",
-    author="xue-pai",
-    author_email="xue-pai@users.noreply.github.com",
+    version="2.1.0",
+    author="fuxictr",
+    author_email="fuxictr@users.noreply.github.com",
     description="A configurable, tunable, and reproducible library for CTR prediction",
     long_description=long_description,
     long_description_content_type="text/markdown",