Repository: NVIDIA-Merlin/NVTabular

Issue:
[FEA] Move or remove get_emb_sz method on Categorify op
<a href="https://github.com/alecgunny"><img src="https://avatars0.githubusercontent.com/u/14932242?v=4" align="left" width="96" height="96" hspace="10"></img></a> **Issue by [alecgunny](https://github.com/alecgunny)**
_Wednesday May 13, 2020 at 16:53 GMT_
_Originally opened as https://github.com/rapidsai/recsys/issues/121_

----

It's not entirely clear to me why `Categorify.get_emb_sz` and its associated helper methods are included as methods on the `Categorify` object, especially when nothing in the code actually uses attributes from the object itself (with the exception of the `self.embed_sz`, which gets set rather than read from):
```python
    def get_emb_sz(self, encoders, cat_names):
        work_in = {}
        for key in encoders.keys():
            work_in[key] = encoders[key] + 1
        # sorted key required to ensure same sort occurs for all values
        ret_list = [
            (n, self.def_emb_sz(work_in, n))
            for n in sorted(cat_names, key=lambda entry: entry.split("_")[0])
        ]
        return ret_list

    def emb_sz_rule(self, n_cat: int) -> int:
        return min(16, round(1.6 * n_cat ** 0.56))

    def def_emb_sz(self, classes, n, sz_dict=None):
        """Pick an embedding size for `n` depending on `classes` if not given in `sz_dict`.
        """
        sz_dict = sz_dict if sz_dict else {}
        n_cat = classes[n]
        sz = sz_dict.get(n, int(self.emb_sz_rule(n_cat)))  # rule of thumb
        self.embed_sz[n] = sz
        return n_cat, sz
```
While I'm personally of the opinion that it's not our library's job to be providing rules of thumb for building deep learning models, just to provide the data for whatever rules the user wants to use, if we're intent on having this function somewhere it seems like it would be better served as a standalone function of a single encoder or as a property of the encoder itself:
```python
emb_szs = [get_embed_sz(proc.stats["categories"][column]) for column in proc.columns_ctx["categorical"]["base"]]

# or
emb_szs = [proc.stats["categories"][column].embed_sz[1] for column in proc.columns_ctx["categorical"]["base"]]
```



Discussion:
Note that solutions like

```python
emb_szs = [get_embed_sz(proc.stats["categories"][column]) for column in proc.columns_ctx["categorical"]["base"]]
```

will break compatibility with `torch_dataloader` (and possibly other modules too), which sort columns (see for example [this line](https://github.com/NVIDIA/NVTabular/blob/2e5f90bc47c96292484f210a39016918038a6911/nvtabular/torch_dataloader.py#L128) in `create_tensors`).

@benfred @alecgunny is there any agreed upon convention for column ordering? I'm wondering why `torch_dataloader` reorders columns, and whether this is done elsewhere, too.

If reordering is necessary, then I'll create a new issue to order the `proc.stats["categories"]` dictionary (among others), since as of CPython 3.6 (and more officially Python 3.7 in general) dictionaries are ordered.
Just realized that this same logic is present in `get_emb_sz` itself. Is this sorting convention necessary or is it something we can get rid of? Sorting according to this key seems odd and error-prone:

```
sorted(cat_names, key=lambda entry: entry.split("_")[0])
```
@rdipietro I think this is kind of a distinct issue, though certainly one worth noting. My point with the snippet
```
emb_szs = [get_embed_sz(proc.stats["categories"][column]) for column in proc.columns_ctx["categorical"]["base"]]
```
was simply to show what the call signature (i.e. `get_embed_sz(...)`) looks like. How you choose to iterate through it in a list is up to you (though obviously a good implementation will match how our dataloader orders things, and makes more critical exposing the means by which our dataloader does that ordering), but won't change the underlying function call.

While I'm thinking about this, I think it's probably worth discussing the choice of the syntax `embed_sz`. I know it matches what FastAI uses, but as we begin to expand beyond it as the target use case, I think abandoning some of its more idiosyncratic usage of abbreviations is worth doing for the sake of clarity.

Resolving patch:
diff --git a/examples/criteo_benchmark.py b/examples/criteo_benchmark.py
--- a/examples/criteo_benchmark.py
+++ b/examples/criteo_benchmark.py
@@ -40,7 +40,7 @@ def parse_args():
 
 from nvtabular import Workflow
 from nvtabular.io import GPUDatasetIterator, device_mem_size
-from nvtabular.ops import Categorify, LogOp, Normalize, ZeroFill
+from nvtabular.ops import Categorify, LogOp, Normalize, ZeroFill, get_embedding_size
 from nvtabular.torch_dataloader import DLCollator, DLDataLoader, FileItrDataset
 
 
@@ -144,9 +144,7 @@ def parse_args():
 
 embeddings = [
     x[1]
-    for x in proc.df_ops["Categorify"].get_emb_sz(
-        proc.stats["categories"], proc.columns_ctx["categorical"]["base"]
-    )
+    for x in get_embedding_size(proc.stats["categories"], proc.columns_ctx["categorical"]["base"])
 ]
 print("Creating Iterators for dataloader")
 start = time()
diff --git a/nvtabular/ops.py b/nvtabular/ops.py
--- a/nvtabular/ops.py
+++ b/nvtabular/ops.py
@@ -1154,7 +1154,6 @@ class Categorify(DFOperator):
         Replaces the transformed column with the original input
         if set Yes
     cat_names :
-    embed_sz :
     """
 
     default_in = CAT
@@ -1171,7 +1170,6 @@ def __init__(
         preprocessing=True,
         replace=True,
         cat_names=None,
-        embed_sz=None,
         out_path=None,
         split_out=None,
         na_sentinel=None,
@@ -1186,7 +1184,6 @@ def __init__(
         self.gpu_mem_util_limit = gpu_mem_util_limit
         self.gpu_mem_trans_use = gpu_mem_trans_use
         self.cat_names = cat_names if cat_names else []
-        self.embed_sz = embed_sz if embed_sz else {}
         self.out_path = out_path or "./"
         self.split_out = split_out
         self.na_sentinel = na_sentinel or 0
@@ -1242,30 +1239,32 @@ def op_logic(self, gdf: cudf.DataFrame, target_columns: list, stats_context={}):
                 new_gdf[new_col] = new_gdf[new_col].astype(self.dtype, copy=False)
         return new_gdf
 
-    def get_emb_sz(self, encoders, cat_names):
-        # sorted key required to ensure same sort occurs for all values
-        ret_list = [
-            (n, self.def_emb_sz(encoders, n))
-            for n in sorted(cat_names, key=lambda entry: entry.split("_")[0])
-        ]
-        return ret_list
 
-    def emb_sz_rule(self, n_cat: int) -> int:
-        return min(16, round(1.6 * n_cat ** 0.56))
+def get_embedding_order(cat_names):
+    """ Returns a consistent sorder order for categorical variables
 
-    def def_emb_sz(self, classes, n, sz_dict=None):
-        """Pick an embedding size for `n` depending on `classes` if not given in `sz_dict`.
-        """
-        sz_dict = sz_dict if sz_dict else {}
-        n_cat = classes[n]
-        sz = sz_dict.get(n, int(self.emb_sz_rule(n_cat)))  # rule of thumb
-        self.embed_sz[n] = sz
-        return n_cat, sz
-
-    def get_embeddings(self, encoders, cat_names):
-        embeddings = {}
-        for col in sorted(cat_names):
-            path = encoders[col]
-            num_rows, _, _ = cudf.io.read_parquet_metadata(path)
-            embeddings[col] = (num_rows, self.emb_sz_rule(num_rows))
-        return embeddings
+    Parameters
+    -----------
+    cat_names : list of str
+        names of the categorical columns
+    """
+    return sorted(cat_names)
+
+
+def get_embedding_size(encoders, cat_names):
+    """ Returns a suggested size of embeddings based off cardinality of encoding categorical
+    variables
+    Parameters
+    -----------
+    encoders : dict
+        The encoding statistics of the categorical variables (ie. from workflow.stats["categories"])
+    cat_names : list of str
+        names of the categorical columns
+    """
+    # sorted key required to ensure same sort occurs for all values
+    ret_list = [(n, _emb_sz_rule(encoders[n])) for n in get_embedding_order(cat_names)]
+    return ret_list
+
+
+def _emb_sz_rule(n_cat: int) -> int:
+    return n_cat, int(min(16, round(1.6 * n_cat ** 0.56)))
diff --git a/nvtabular/torch_dataloader.py b/nvtabular/torch_dataloader.py
--- a/nvtabular/torch_dataloader.py
+++ b/nvtabular/torch_dataloader.py
@@ -21,6 +21,7 @@
 from torch.utils.dlpack import from_dlpack
 
 from nvtabular.io import GPUFileIterator
+from nvtabular.ops import get_embedding_order
 
 
 class FileItrDataset(torch.utils.data.IterableDataset):
@@ -124,11 +125,7 @@ def create_tensors_plain(gdf, cat_cols=None, cont_cols=None, label_cols=None):
 
 
 def combine_tensors(cats, conts, label):
-    cats_list = (
-        [cats[x] for x in sorted(cats.keys(), key=lambda entry: entry.split("_")[0])]
-        if cats
-        else None
-    )
+    cats_list = [cats[x] for x in get_embedding_order(cats.keys())] if cats else None
     conts_list = [conts[x] for x in sorted(conts.keys())] if conts else None
     label_list = [label[x] for x in sorted(label.keys())] if label else None
 
@@ -152,12 +149,10 @@ def _one_df(
         _to_tensor(gdf_label, torch.float32, label, to_cpu=False)
 
 
-def get_final_cols(preproc):
+def _get_final_cols(preproc):
     if "cols" not in preproc.columns_ctx["final"]:
         preproc.create_final_cols()
-    cat_names = sorted(
-        preproc.columns_ctx["final"]["cols"]["categorical"], key=lambda entry: entry.split("_")[0],
-    )
+    cat_names = get_embedding_order(preproc.columns_ctx["final"]["cols"]["categorical"])
     cont_names = sorted(preproc.columns_ctx["final"]["cols"]["continuous"])
     label_name = sorted(preproc.columns_ctx["final"]["cols"]["label"])
     return cat_names, cont_names, label_name
@@ -178,7 +173,7 @@ def process_one_df(
         gdf = preproc.apply_ops(gdf)
 
     if preproc:
-        cat_names, cont_names, label_names = get_final_cols(preproc)
+        cat_names, cont_names, label_names = _get_final_cols(preproc)
 
     _one_df(
         gdf,
@@ -228,8 +223,7 @@ def __iter__(self):
         for _ in range(self.num_chunks):
             chunk = buff.get()
             yield from TensorItr(
-                chunk,
-                batch_size=self.batch_size,
+                chunk, batch_size=self.batch_size,
             )
 
     def load_chunk(self, out):


Test patch: