Source code for nvtabular.loader.tensorflow

# Copyright (c) 2021, NVIDIA CORPORATION.
# 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,
# See the License for the specific language governing permissions and
# limitations under the License.
import contextlib
import logging
import os

import dask.dataframe as dd
import numpy as np

from merlin.core.dispatch import HAS_GPU
from merlin.schema import Tags
from nvtabular.loader.backend import DataLoader
from nvtabular.loader.tf_utils import configure_tensorflow, get_dataset_schema_from_feature_columns

from_dlpack = configure_tensorflow()
LOG = logging.getLogger("nvtabular")
# tf import must happen after config to restrict memory use
import tensorflow as tf  # noqa

# noqa
    from import Dataset

    nvt_dataset_class = Dataset
except ImportError:
    nvt_dataset_class = None
# pylint has issues with TF array ops, so disable checks until fixed:
# pylint: disable=no-value-for-parameter,unexpected-keyword-arg,redundant-keyword-arg

dd_engine = {
    "parquet": dd.read_parquet,
    "csv": dd.read_csv,
    "df": dd.DataFrame,

def _validate_dataset(paths_or_dataset, batch_size, buffer_size, engine, device, reader_kwargs):
    # TODO: put this in parent class and allow
    # torch dataset to leverage as well?

    # if a dataset was passed, just return it
    if hasattr(paths_or_dataset, "schema"):
        return paths_or_dataset

    # otherwise initialize a dataset
    # from paths or glob pattern
    if isinstance(paths_or_dataset, str):
        files =
        parent, file = os.path.split(paths_or_dataset)
        _is_empty_msg = f"Couldn't find file pattern {file} in directory {parent}"
        # TODO: some checking around attribute
        # error here?
        files = list(paths_or_dataset)
        _is_empty_msg = "paths_or_dataset list must contain at least one filename"

    assert isinstance(files, list)
    if len(files) == 0:
        raise ValueError(_is_empty_msg)

    if not engine:
        # default engine is parquet
        engine = "parquet"

    cpu = device and "cpu" in device

    if nvt_dataset_class:
        return nvt_dataset_class(files, engine=engine, cpu=cpu)
            "NVTabular Dataset class not detected, reverting to Dask Dataframe."
            "Expect slower iteration speeds."
    return dd_engine[engine](files)

def _validate_schema(feature_columns, cat_names, cont_names, schema=None):
    _uses_feature_columns = feature_columns is not None
    _uses_explicit_schema = (cat_names is not None) or (cont_names is not None)

    cat_tag_names = schema.select_by_tag([Tags.CATEGORICAL]).column_names if schema else []
    cont_tag_names = schema.select_by_tag([Tags.CONTINUOUS]).column_names if schema else []
    _uses_dataset_schema = cat_tag_names or cont_tag_names

    if _uses_feature_columns and _uses_explicit_schema:
        raise ValueError(
            "Passed `feature_column`s and explicit column names, must be one or the other"
    elif _uses_feature_columns:
        return get_dataset_schema_from_feature_columns(feature_columns)
    elif _uses_explicit_schema:
        cat_names = cat_names or []
        cont_names = cont_names or []
        return cat_names, cont_names
    elif _uses_dataset_schema:
        cat_tag_names = cat_tag_names or []
        cont_tag_names = cont_tag_names or []
        return cat_tag_names, cont_tag_names
        raise ValueError(
            "Must either pass a list of TensorFlow `feature_column`s "
            "or explicit `cat_name` and `cont_name` column name lists."

def _get_schema(dataset):
    if hasattr(dataset, "schema"):
        return dataset.schema
    return None

[docs]class KerasSequenceLoader(tf.keras.utils.Sequence, DataLoader): """ Infinite generator used to asynchronously iterate through CSV or Parquet dataframes on GPU by leveraging an NVTabular `Dataset`. Applies preprocessing via NVTabular `Workflow` objects and outputs tabular dictionaries of TensorFlow Tensors via `dlpack <>`_. Useful for training tabular models built in Keras and trained via ` <>`_. The data loading scheme is implemented by loading, preprocessing, and batching data in an asynchronous thread. The amount of randomness in shuffling is controlled by the `buffer_size` and `parts_per_chunk` kwargs. At load time, sub-chunks of data with size controlled by `buffer_size` are loaded from random partitions in the dataset, and `parts_per_chunk` of them are concatenated into a single chunk, shuffled, and split into batches. This means that each chunk has `buffer_size*parts_per_chunk` rows, and due to the asynchronous nature of the dataloader that means there are, including the batch being processed by your network, `3*buffer_size*parts_per_chunk` rows of data in GPU memory at any given time. This means that for a fixed memory budget, using more `parts_per_chunk` will come at the expense of smaller `buffer_size`, increasing the number of reads and reducing throughput. The goal should be to maximize the total amount of memory utilized at once without going OOM and with the fewest number of reads to meet your epoch-level randomness needs. An important thing to note is that TensorFlow's default behavior is to claim all GPU memory for itself at initialziation time, which leaves none for NVTabular to load or preprocess data. As such, we attempt to configure TensorFlow to restrict its memory allocation on a given GPU using the environment variables `TF_MEMORY_ALLOCATION` and `TF_VISIBLE_DEVICE`. If `TF_MEMORY_ALLOCATION < 1`, it will be assumed that this refers to a fraction of free GPU memory on the given device. Otherwise, it will refer to an explicit allocation amount in MB. `TF_VISIBLE_DEVICE` should be an integer GPU index. Iterator output is of the form `(dict(features), list(labels))`, where each element of the features dict is a `feature_name: feature_tensor` and each elemtn of the labels list is a tensor, and all tensors are of shape `(batch_size, 1)`. Note that this means vectorized continuous and multi-hot categorical features are not currently supported. The underlying NVTabular `Dataset` object is stored in the `data` attribute, and should be used for updating NVTabular `Workflow` statistics:: workflow = nvt.Workflow(...) dataset = KerasSequenceLoader(...) workflow.update_stats(, record_stats=True) Parameters ------------- - paths_or_dataset: str or list(str) Either a string representing a file pattern (see `tf.glob` for pattern rules), a list of filenames to be iterated through, or a Dataset object, in which case `buffer_size`, `engine`, and `reader_kwargs` will be ignored - batch_size: int Number of samples to yield at each iteration - label_names: list(str) Column name of the target variable in the dataframe specified by `paths_or_dataset` - feature_columns: list(tf.feature_column) or None A list of TensorFlow feature columns representing the inputs exposed to the model to be trained. Columns with parent columns will climb the parent tree, and the names of the columns in the unique set of terminal columns will be used as the column names. If left as None, must specify `cat_names` and `cont_names` - cat_names: list(str) or None List of categorical column names. Ignored if `feature_columns` is specified - cont_names: list(str) or None List of continuous column names. Ignored if `feature_columns` is specified - engine: {'csv', 'parquet', None}, default None String specifying the type of read engine to use. If left as `None`, will try to infer the engine type from the file extension. - shuffle: bool, default True Whether to shuffle chunks of batches before iterating through them. - buffer_size: float or int If `0 < buffer_size < 1`, `buffer_size` will refer to the fraction of total GPU memory to occupy with a buffered chunk. If `1 < buffer_size < batch_size`, the number of rows read for a buffered chunk will be equal to `int(buffer_size*batch_size)`. Otherwise, if `buffer_size > batch_size`, `buffer_size` rows will be read in each chunk (except for the last chunk in a dataset, which will, in general, be smaller). Larger chunk sizes will lead to more efficiency and randomness, but require more memory. - device: None Which GPU device to load from. Ignored for now - parts_per_chunk: int Number of dataset partitions with size dictated by `buffer_size` to load and concatenate asynchronously. More partitions leads to better epoch-level randomness but can negatively impact throughput - reader_kwargs: dict extra kwargs to pass when instantiating the underlying `nvtabular.Dataset` sparse_list : list(str) or None list with column names of columns that should be represented as sparse tensors sparse_max : dict dictionary of key: column_name + value: integer representing max sequence length for column sparse_dense : bool bool value to activate transforming sparse tensors to dense """ _use_nnz = True def __init__( self, paths_or_dataset, batch_size, label_names=None, feature_columns=None, cat_names=None, cont_names=None, engine=None, shuffle=True, seed_fn=None, buffer_size=0.1, device=None, parts_per_chunk=1, reader_kwargs=None, global_size=None, global_rank=None, drop_last=False, sparse_names=None, sparse_max=None, sparse_as_dense=False, schema=None, ): device = device or 0 device = "cpu" if not HAS_GPU else device dataset = _validate_dataset( paths_or_dataset, batch_size, buffer_size, engine, device, reader_kwargs ) schema = _get_schema(dataset) if not schema else schema cat_names, cont_names = _validate_schema( feature_columns, cat_names, cont_names, schema=schema ) DataLoader.__init__( self, dataset, batch_size, shuffle, cat_names=cat_names, cont_names=cont_names, label_names=label_names, seed_fn=seed_fn, parts_per_chunk=parts_per_chunk, device=device, global_size=global_size, global_rank=global_rank, drop_last=drop_last, sparse_names=sparse_names, sparse_max=sparse_max, sparse_as_dense=sparse_as_dense, ) self._map_fns = [] def __len__(self): """ recreating since otherwise Keras yells at you """ # TODO: what's a better way to do this inheritance # of the appropriate methods? A Metaclass? DataLoader.stop(self) return DataLoader.__len__(self) def __getitem__(self, idx): """ implemented exclusively for consistency with Keras Does not leverage passed idx in any way """ return DataLoader.__next__(self)
[docs] def map(self, fn): """ Applying a function to each batch. This can for instance be used to add `sample_weight` to the model. """ self._map_fns.append(fn) return self
@contextlib.contextmanager def _get_device_ctx(self, dev): # with tf.device("/device:GPU:{}".format(dev)) as tf_device: # # tf.device changes the cupy cuda device, which breaks us on multigpu # # force cupy to still use the device we expect # cupy.cuda.Device(dev).use() # yield tf_device # commenting out since device statements cause # RuntimeErrors when exiting if two dataloaders # are running at once (e.g. train and validation) if dev != "cpu": yield tf.device("/GPU:" + str(dev)) else: # yield tf.device("/device:CPU:0") def _split_fn(self, tensor, idx, axis=0): return tf.split(tensor, idx, axis=axis) def _tensor_split(self, tensor, idx, axis=0): """ Same function as above but need this method for api match. """ return tf.split(tensor, idx, axis=axis) @property def _LONG_DTYPE(self): return tf.int64 @property def _FLOAT32_DTYPE(self): return tf.float32 def _pack(self, gdf): if isinstance(gdf, np.ndarray): return gdf elif hasattr(gdf, "to_dlpack") and callable(getattr(gdf, "to_dlpack")): return gdf.to_dlpack() elif hasattr(gdf, "to_numpy") and callable(getattr(gdf, "to_numpy")): gdf = gdf.to_numpy() if isinstance(gdf[0], list): gdf = np.stack(gdf) return gdf return gdf.toDlpack() def _unpack(self, gdf): if hasattr(gdf, "shape"): return tf.convert_to_tensor(gdf) return from_dlpack(gdf) def _to_tensor(self, gdf, dtype=None): if gdf.empty: return # checks necessary because of this bug # if len(gdf.shape) == 1 or gdf.shape[1] == 1: dlpack = self._pack(gdf) elif gdf.shape[0] == 1: dlpack = self._pack(gdf.values[0]) else: dlpack = self._pack(gdf.values.T) # catch error caused by tf eager context # not being initialized try: x = self._unpack(dlpack) except AssertionError: tf.random.uniform((1,)) x = self._unpack(dlpack) # if rank is already two it is already in list format if gdf.shape[0] == 1 and not tf.rank(x) == 2: # batch size 1 so got squashed to a vector x = tf.expand_dims(x, 0) elif len(gdf.shape) == 1 or len(x.shape) == 1: # sort of a generic check for any other # len(shape)==1 case, could probably # be more specific x = tf.expand_dims(x, -1) elif gdf.shape[1] > 1: # matrix which means we had to transpose # for the bug above, so untranspose x = tf.transpose(x) return x def _pull_values_offsets(self, values_offset): """ values_offset is either a tuple (values, offsets) or just values. Values is a tensor. This method is used to turn a tensor into its sparse representation """ # pull_values_offsets, return values offsets diff_offsets diff_offsets = None if isinstance(values_offset, tuple): values = tf.reshape(values_offset[0], [-1]) diff_offsets = tf.cast(tf.reshape(values_offset[1], [-1]), dtype=tf.int64) offsets = tf.math.cumsum(diff_offsets) else: values = tf.reshape(values_offset, [-1]) offsets = tf.arange(tf.shape(values)[0], dtype=tf.int64) diff_offsets = offsets[1:] - offsets[:-1] num_rows = len(offsets) return values, offsets, diff_offsets, num_rows def _get_max_seq_len(self, diff_offsets): # get_max_seq_len, return int return int(tf.math.reduce_max(diff_offsets)) def _get_indices(self, offsets, diff_offsets): # Building the indices to reconstruct the sparse tensors row_ids = tf.range(len(offsets), dtype=tf.int64) row_ids_repeated = tf.repeat(row_ids, diff_offsets) row_offset_repeated = tf.repeat(offsets, diff_offsets) col_ids = tf.range(len(row_offset_repeated), dtype=tf.int64) - row_offset_repeated indices = tf.concat( values=[tf.expand_dims(row_ids_repeated, -1), tf.expand_dims(col_ids, -1)], axis=1 ) return indices def _get_sparse_tensor(self, values, indices, num_rows, seq_limit): sparse_tensor = tf.sparse.SparseTensor( indices=indices, values=values, dense_shape=[num_rows, seq_limit] ) return sparse_tensor def _build_sparse_tensor(self, values, offsets, diff_offsets, num_rows, seq_limit): ragged = tf.RaggedTensor.from_row_lengths(values=values, row_lengths=diff_offsets) tensor = tf.RaggedTensor.from_tensor(ragged.to_tensor(shape=[None, seq_limit])).to_sparse() if self.sparse_as_dense: tensor = tf.sparse.to_dense(tensor) return tensor def _handle_tensors(self, cats, conts, labels): to_return = super()._handle_tensors(cats, conts, labels) for map_fn in self._map_fns: to_return = map_fn(*to_return) return to_return
[docs]class KerasSequenceValidater(tf.keras.callbacks.Callback): # TODO: document _supports_tf_logs = True def __init__(self, dataloader): super().__init__() self.dataloader = dataloader
[docs] def on_epoch_end(self, epoch, logs=None): logs = logs if logs is not None else {} for X, y_true in self.dataloader: y_pred = self.model(X) # TODO: how do we want to handle the multi-output case? for metric in self.model.metrics: metric.update_state(y_true, y_pred) set_logs = {} for metric in self.model.metrics: set_logs[f"val_{}"] = metric.result().numpy() logs.update(set_logs) print(set_logs) return logs