oneflow.one_embedding

OneFlow one_embedding operations.

class oneflow.one_embedding.MultiTableEmbedding(name, embedding_dim, dtype, key_type, tables, store_options, default_initializer=None)

MultiTableEmbedding represent multi Embedding tables with same embedding_dim, dtype, and key_type.

Parameters

  • name (str) – The name of Embedding

  • embedding_dim (int) – the size of each embedding vector

  • dtype (flow.dtype) – the data type of embeddings

  • key_type (flow.dtype) – the data type of feature ids

  • tables (list) – list of table param which can be made by flow.one_embedding.make_table_options

  • store_options (dict) – store option of Embedding

  • default_initializer (dict, optional) – if tables param is None, use default_initializer to initialize table. Defaults to None.

For example:

>>> import oneflow as flow
>>> import numpy as np
>>> import oneflow.nn as nn
>>> # a simple example with 3 table
>>> table_size_array = [39884407, 39043, 17289]
>>> vocab_size = sum(table_size_array)
>>> num_tables = len(table_size_array)
>>> embedding_size = 128
>>> scales = np.sqrt(1 / np.array(table_size_array))
>>> tables = [
>>>     flow.one_embedding.make_table_options(
>>>         flow.one_embedding.make_uniform_initializer(low=-scale, high=scale)
>>>     )
>>>     for scale in scales
>>> ]
>>> store_options = flow.one_embedding.make_cached_ssd_store_options(
>>>     cache_budget_mb=8192, persistent_path="/your_path_to_ssd", capacity=vocab_size,
>>> )
>>> embedding = flow.one_embedding.MultiTableEmbedding(
>>>     name="my_embedding",
>>>     embedding_dim=embedding_size,
>>>     dtype=flow.float,
>>>     key_type=flow.int64,
>>>     tables=tables,
>>>     store_options=store_options,
>>> )
>>> embedding.to("cuda")
>>> mlp = flow.nn.FusedMLP(
>>>     in_features=embedding_size * num_tables,
>>>     hidden_features=[512, 256, 128],
>>>     out_features=1,
>>>     skip_final_activation=True,
>>> )
>>> mlp.to("cuda")
>>>
>>> class TrainGraph(flow.nn.Graph):
>>>     def __init__(self,):
>>>         super().__init__()
>>>         self.embedding_lookup = embedding
>>>         self.mlp = mlp
>>>         self.add_optimizer(
>>>             flow.optim.SGD(self.embedding_lookup.parameters(), lr=0.1, momentum=0.0)
>>>         )
>>>         self.add_optimizer(
>>>             flow.optim.SGD(self.mlp.parameters(), lr=0.1, momentum=0.0)
>>>         )
>>>     def build(self, ids):
>>>         embedding = self.embedding_lookup(ids)
>>>         loss = self.mlp(flow.reshape(embedding, (-1, num_tables * embedding_size)))
>>>         loss = loss.sum()
>>>         loss.backward()
>>>         return loss
>>> ids = np.random.randint(0, 1000, (100, num_tables), dtype=np.int64)
>>> ids_tensor = flow.tensor(ids, requires_grad=False).to("cuda")
>>> graph = TrainGraph()
>>> loss = graph(ids_tensor)
>>> print(loss)
load_snapshot(snapshot_name)

load snapshot

Parameters

snapshot_name (str) – the snapshot_name, snapshot will be load from your_configed_persistent_path

For example:

>>> import oneflow as flow
>>> # use embedding create by flow.one_embedding.MultiTableEmbedding
>>> embedding.load_snapshot("my_snapshot1")
>>> # load a snapshot named "my_snapshot1" from your_configed_persistent_path
save_snapshot(snapshot_name)

save snapshot

Parameters

snapshot_name (str) – the snapshot_name, snapshot will be saved in the snapshots dir under your_configed_persistent_path

For example:

>>> import oneflow as flow
>>> # use embedding create by flow.one_embedding.MultiTableEmbedding
>>> embedding.save_snapshot("my_snapshot1")
>>> # a snapshot named "my_snapshot1" have been saved in the "snapshots" dir under your_configed_persistent_path
>>> # which can be reload by flow.one_embedding.load_snapshot
oneflow.one_embedding.MultiTableEmbedding.forward(self, ids, table_ids=None)

Embedding lookup operation

Parameters

  • ids (flow.tensor) – the feature ids

  • table_ids (flow.tensor, optional) – the table_id of each id, must be same shape as ids. There is no need to pass table_ids, if has config only one table or the ids has shape (batch_size, num_tables), and each column’s id belongs to the column_id th table, otherwise, you should pass the tensor_ids.

Returns

the result of embedding lookup

Return type

flow.tensor

class oneflow.one_embedding.MultiTableMultiColumnEmbedding(name, embedding_dim, dtype, key_type, tables, store_options, default_initializer=None)

MultiTableMultiColumnEmbedding represent multi Embedding tables with multi embedding_dim, same dtype, and key_type.

Parameters

  • name (str) – The name of Embedding

  • embedding_dim (list) – list of the size of each embedding vector

  • dtype (flow.dtype) – the data type of embeddings

  • key_type (flow.dtype) – the data type of feature ids

  • tables (list) – list of table param which can be made by flow.one_embedding.make_table_options

  • store_options (dict) – store option of Embedding

  • default_initializer (dict, optional) – if tables param is None, use default_initializer to initialize table. Defaults to None.

For example:

>>> import oneflow as flow
>>> import numpy as np
>>> import oneflow.nn as nn
>>> # a simple example with 3 table, every table has two column, the first column embedding_size is 10 and the second is 1.
>>> # every table's first column initialize with uniform(-1/sqrt(table_size), 1/sqrt(table_size)), second column initialize with normal(0, 1/sqrt(table_size))
>>> table_size_array = [39884407, 39043, 17289]
>>> vocab_size = sum(table_size_array)
>>> num_tables = len(table_size_array)
>>> embedding_size_list = [10, 1]
>>> scales = np.sqrt(1 / np.array(table_size_array))
>>> tables = [
>>>     flow.one_embedding.make_table_options(
>>>       [flow.one_embedding.make_column_options(
>>>         flow.one_embedding.make_uniform_initializer(low=-scale, high=scale)),
>>>        flow.one_embedding.make_column_options(
>>>         flow.one_embedding.make_normal_initializer(mean=0, std=scale))]
>>>     )
>>>     for scale in scales
>>> ]
>>> store_options = flow.one_embedding.make_cached_ssd_store_options(
>>>     cache_budget_mb=8192, persistent_path="/your_path_to_ssd", capacity=vocab_size,
>>> )
>>> embedding = flow.one_embedding.MultiTableMultiColumnEmbedding(
>>>     name="my_embedding",
>>>     embedding_dim=embedding_size_list,
>>>     dtype=flow.float,
>>>     key_type=flow.int64,
>>>     tables=tables,
>>>     store_options=store_options,
>>> )
>>> embedding.to("cuda")
>>> mlp = flow.nn.FusedMLP(
>>>     in_features=sum(embedding_size_list) * num_tables,
>>>     hidden_features=[512, 256, 128],
>>>     out_features=1,
>>>     skip_final_activation=True,
>>> )
>>> mlp.to("cuda")
>>>
>>> class TrainGraph(flow.nn.Graph):
>>>     def __init__(self,):
>>>         super().__init__()
>>>         self.embedding_lookup = embedding
>>>         self.mlp = mlp
>>>         self.add_optimizer(
>>>             flow.optim.SGD(self.embedding_lookup.parameters(), lr=0.1, momentum=0.0)
>>>         )
>>>         self.add_optimizer(
>>>             flow.optim.SGD(self.mlp.parameters(), lr=0.1, momentum=0.0)
>>>         )
>>>     def build(self, ids):
>>>         embedding = self.embedding_lookup(ids)
>>>         loss = self.mlp(flow.reshape(embedding, (-1, num_tables * sum(embedding_size_list))))
>>>         loss = loss.sum()
>>>         loss.backward()
>>>         return loss
>>> ids = np.random.randint(0, 1000, (100, num_tables), dtype=np.int64)
>>> ids_tensor = flow.tensor(ids, requires_grad=False).to("cuda")
>>> graph = TrainGraph()
>>> loss = graph(ids_tensor)
>>> print(loss)
load_snapshot(snapshot_name)

load snapshot

Parameters

snapshot_name (str) – the snapshot_name, snapshot will be load from your_configed_persistent_path

For example:

>>> import oneflow as flow
>>> # use embedding create by flow.one_embedding.MultiTableEmbedding
>>> embedding.load_snapshot("my_snapshot1")
>>> # load a snapshot named "my_snapshot1" from your_configed_persistent_path
save_snapshot(snapshot_name)

save snapshot

Parameters

snapshot_name (str) – the snapshot_name, snapshot will be saved in the snapshots dir under your_configed_persistent_path

For example:

>>> import oneflow as flow
>>> # use embedding create by flow.one_embedding.MultiTableEmbedding
>>> embedding.save_snapshot("my_snapshot1")
>>> # a snapshot named "my_snapshot1" have been saved in the "snapshots" dir under your_configed_persistent_path
>>> # which can be reload by flow.one_embedding.load_snapshot
oneflow.one_embedding.MultiTableMultiColumnEmbedding.forward(self, ids, table_ids=None)

Embedding lookup operation

Parameters

  • ids (flow.tensor) – the feature ids

  • table_ids (flow.tensor, optional) – the table_id of each id, must be same shape as ids. There is no need to pass table_ids, if has config only one table or the ids has shape (batch_size, num_tables), and each column’s id belongs to the column_id th table, otherwise, you should pass the tensor_ids.

Returns

the result of embedding lookup

Return type

flow.tensor

oneflow.one_embedding.make_device_mem_store_options(persistent_path, capacity, size_factor=1, physical_block_size=512)

make GPU only store_options param of MultiTableEmbedding

Parameters

  • persistent_path (str, list) – persistent storage path of Embedding. If passed a str, current rank Embedding will be saved in path/rank_id-num_ranks path. If passed a list, the list length must equals num_ranks, each elem of list represent the path of rank_id Embedding.

  • capacity (int) – total capacity of Embedding

  • size_factor (int, optional) – store size factor of embedding_dim, if SGD update, and momentum = 0, should be 1, if momentum > 0, it should be 2. if Adam, should be 3. Defaults to 1.

  • physical_block_size (int, optional) – physical_block_size should be sector size. Defaults to 512.

Returns

GPU only store_options param of MultiTableEmbedding

Return type

dict

See also oneflow.one_embedding.make_cached_ssd_store_options()

oneflow.one_embedding.make_cached_ssd_store_options(cache_budget_mb, persistent_path, capacity=None, size_factor=1, physical_block_size=512, host_cache_budget_mb=0)

make SSD use GPU and host as cache store_options param of MultiTableEmbedding. If cache_budget_mb > 0 and host_cache_budget_mb > 0, use GPU and host memory as multi-level cache.

Parameters

  • cache_budget_mb (int) – the MB budget of per GPU as cache.

  • persistent_path (str, list) – persistent storage path of Embedding, must use fast SSD because of frequently random disk access during training. If passed a str, current rank Embedding will be saved in path/rank_id-num_ranks path. If passed a list, the list length must equals num_ranks, each elem of list represent the path of rank_id Embedding.

  • capacity (int) – total capacity of Embedding

  • size_factor (int, optional) – store size factor of embedding_dim, if SGD update, and momentum = 0, should be 1, if momentum > 0, it should be 2. if Adam, should be 3. Defaults to 1.

  • physical_block_size (int, optional) – physical_block_size should be sector size. Defaults to 512.

  • host_cache_budget_mb (int) – the MB budget of host memory as cache per rank. Defaults to 0.

Returns

SSD use GPU and host as cache store_options param of MultiTableEmbedding

Return type

dict

For example:

>>> import oneflow as flow
>>> store_options = flow.one_embedding.make_cached_ssd_store_options(
>>>     cache_budget_mb=8192, persistent_path="/your_path_to_ssd", capacity=vocab_size,
>>> )
>>> # pass the store_options to the "store_options" param of flow.one_embedding.MultiTableEmbedding
>>> # ...
oneflow.one_embedding.make_cached_host_mem_store_options(cache_budget_mb, persistent_path, capacity, size_factor=1, physical_block_size=512)

make host use GPU as cache store_options param of MultiTableEmbedding

Parameters

  • cache_budget_mb (int) – the MB budget of per GPU as cache.

  • persistent_path (str, list) – persistent storage path of Embedding. If passed a str, current rank Embedding will be saved in path/rank_id-num_ranks path. If passed a list, the list length must equals num_ranks, each elem of list represent the path of rank_id Embedding.

  • capacity (int) – total capacity of Embedding

  • size_factor (int, optional) – store size factor of embedding_dim, if SGD update, and momentum = 0, should be 1, if momentum > 0, it should be 2. if Adam, should be 3. Defaults to 1.

  • physical_block_size (int, optional) – physical_block_size should be sector size. Defaults to 512.

Returns

host use GPU as cache store_options param of MultiTableEmbedding

Return type

dict

See also oneflow.one_embedding.make_cached_ssd_store_options()

oneflow.one_embedding.make_uniform_initializer(low, high)

make uniform initializer param of make_table_options

Parameters

  • low (float) – A python scalar. Lower bound of the range of random values to generate.

  • high (float) – A python scalar. Upper bound of the range of random values to generate.

Returns

initializer param of make_table_options

Return type

dict

For example:

>>> import oneflow as flow
>>> initializer = flow.one_embedding.make_uniform_initializer(low=-scale, high=scale)
>>> # pass the initializer to flow.one_embedding.make_table_options
>>> # ...
oneflow.one_embedding.make_normal_initializer(mean, std)

make normal initializer param of make_table_options

Parameters

  • mean (float) – A python scalar. Mean of the random values to generate.

  • std (float) – A python scalar. Standard deviation of the random values to generate.

Returns

initializer param of make_table_options

Return type

dict

For example:

>>> import oneflow as flow
>>> initializer = flow.one_embedding.make_normal_initializer(mean=0, std=0.01)
>>> # pass the initializer to flow.one_embedding.make_table_options
>>> # ...
oneflow.one_embedding.make_table_options(param)

make table param of Embedding tables

Parameters

param (dict or list) – param can be initializer or list of column_option. initializer can be made by make_uniform_initializer or make_normal_initializer, column options can be made by make_column_options

Returns

table param of Embedding tables

Return type

dict

For example:

>>> import oneflow as flow
>>> initializer = flow.one_embedding.make_uniform_initializer(low=-scale, high=scale)
>>> table1 = flow.one_embedding.make_table_options(initializer)
>>> table2 = flow.one_embedding.make_table_options(initializer)
>>> tables = [table1, table2]
>>> # pass the tables to the "tables" param of flow.one_embedding.MultiTableEmbedding or flow.one_embedding.MultiTableMultiColumnEmbedding
>>> # ...
oneflow.one_embedding.make_table(param)

alias of oneflow.one_embedding.make_table_options

See also oneflow.one_embedding.make_table_options()

Copyright 2020 The OneFlow Authors. 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

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.

class oneflow.one_embedding.Ftrl(params: Union[Iterator[oneflow.nn.Parameter], List[Dict]], lr: float = 0.001, weight_decay: float = 0.0, lr_power: float = - 0.5, initial_accumulator_value: float = 0.1, lambda1: float = 0.0, lambda2: float = 0.0, beta: float = 0.0)

FTRL Optimizer.

The formula is:

\[ \begin{align}\begin{aligned}& accumlator_{i+1} = accumlator_{i} + grad * grad\\& sigma = (accumulator_{i+1}^{lr\_power} - accumulator_{i}^{lr\_power}) / learning\_rate\\& z_{i+1} = z_{i} + grad - sigma * param_{i}\\\begin{split}\text{} param_{i+1} = \begin{cases} 0 & \text{ if } |z_{i+1}| < \lambda_1 \\ -(\frac{\beta+accumlator_{i+1}^{lr\_power}}{learning\_rate} + \lambda_2)*(z_{i+1} - sign(z_{i+1})*\lambda_1) & \text{ otherwise } \\ \end{cases}\end{split}\end{aligned}\end{align} \]

Example 1:

# Assume net is a custom model.
ftrl = flow.one_embedding.FTRL(net.parameters(), lr=1e-3)

for epoch in range(epochs):
    # Read data, Compute the loss and so on.
    # ...
    loss.backward()
    ftrl.step()
    ftrl.zero_grad()
Parameters

  • params (iterable) – iterable of parameters to optimize or dicts defining parameter groups

  • lr (float, optional) – learning rate. Defaults to 1e-3.

  • weight_decay (float, optional) – weight decay (L2 penalty). Defaults to 0.0.

  • lr_power (float, optional) – learning rate decrease factor. Defaults to -0.5.

  • initial_accumulator_value (float, optional) – The initial value of accumlator. Defaults to 0.1.

  • lambda1 (float, optional) – L1 regularization strength. Defaults to 0.0.

  • lambda2 (float, optional) – L2 regularization strength. Defaults to 0.0.

  • beta (float, optional) – The value of beta. Defaults to 0.0.

step(closure: Optional[Callable] = None)

Performs a single optimization step.

Parameters

closure (callable, optional) – A closure that reevaluates the model and returns the loss.

property support_sparse

Whether the Optimizer support sparse update.

oneflow.one_embedding.make_persistent_table_reader(paths, snapshot_name, key_type, value_type, storage_dim, physical_block_size=512)

Creates a reader for reading persistent table.

Parameters

  • paths (list) – paths of tables to read

  • snapshot_name (str) – name of the snapshot to read

  • key_type (flow.dtype) – the data type of key

  • value_type (flow.dtype) – the data type of value

  • storage_dim (int) – number of elements in each value

  • physical_block_size (int, optional) – physical_block_size should be sector size. Defaults to 512

oneflow.one_embedding.make_persistent_table_writer(paths, snapshot_name, key_type, value_type, storage_dim, physical_block_size=512)

Creates a writer for writing persistent table.

Parameters

  • paths (list) – paths of tables to write

  • snapshot_name (str) – name of the snapshot to write

  • key_type (flow.dtype) – the data type of key

  • value_type (flow.dtype) – the data type of value

  • storage_dim (int) – number of elements in each value

  • physical_block_size (int, optional) – physical_block_size should be sector size. Defaults to 512