478 lines
16 KiB
Python
478 lines
16 KiB
Python
from dataclasses import dataclass
|
|
from typing import Callable, List, Optional, Tuple, Dict
|
|
import functools
|
|
|
|
import torch
|
|
import tensorflow as tf
|
|
|
|
from tml.common.batch import DataclassBatch
|
|
from tml.projects.home.recap.data.config import RecapDataConfig, TaskData
|
|
from tml.projects.home.recap.data import preprocessors
|
|
from tml.projects.home.recap.config import JobMode
|
|
from tml.projects.home.recap.data.tfe_parsing import get_seg_dense_parse_fn
|
|
from tml.projects.home.recap.data.util import (
|
|
keyed_jagged_tensor_from_tensors_dict,
|
|
sparse_or_dense_tf_to_torch,
|
|
)
|
|
from absl import logging
|
|
import torch.distributed as dist
|
|
|
|
|
|
@dataclass
|
|
class RecapBatch(DataclassBatch):
|
|
"""Holds features and labels from the Recap dataset."""
|
|
|
|
continuous_features: torch.Tensor
|
|
binary_features: torch.Tensor
|
|
discrete_features: torch.Tensor
|
|
sparse_features: "KeyedJaggedTensor" # type: ignore[name-defined] # noqa: F821
|
|
labels: torch.Tensor
|
|
user_embedding: torch.Tensor = None
|
|
user_eng_embedding: torch.Tensor = None
|
|
author_embedding: torch.Tensor = None
|
|
weights: torch.Tensor = None
|
|
|
|
def __post_init__(self):
|
|
if self.weights is None:
|
|
self.weights = torch.ones_like(self.labels)
|
|
for feature_name, feature_value in self.as_dict().items():
|
|
if ("embedding" in feature_name) and (feature_value is None):
|
|
setattr(self, feature_name, torch.empty([0, 0]))
|
|
|
|
|
|
def to_batch(x, sparse_feature_names: Optional[List[str]] = None) -> RecapBatch:
|
|
"""Converts a torch data loader output into `RecapBatch`."""
|
|
|
|
x = tf.nest.map_structure(functools.partial(sparse_or_dense_tf_to_torch, pin_memory=False), x)
|
|
try:
|
|
features_in, labels = x
|
|
except ValueError:
|
|
# For Mode.INFERENCE, we do not expect to recieve labels as part of the input tuple
|
|
features_in, labels = x, None
|
|
|
|
sparse_features = keyed_jagged_tensor_from_tensors_dict({})
|
|
if sparse_feature_names:
|
|
sparse_features = keyed_jagged_tensor_from_tensors_dict(
|
|
{embedding_name: features_in[embedding_name] for embedding_name in sparse_feature_names}
|
|
)
|
|
|
|
user_embedding, user_eng_embedding, author_embedding = None, None, None
|
|
if "user_embedding" in features_in:
|
|
if sparse_feature_names and "meta__user_id" in sparse_feature_names:
|
|
raise ValueError("Only one source of embedding for user is supported")
|
|
else:
|
|
user_embedding = features_in["user_embedding"]
|
|
|
|
if "user_eng_embedding" in features_in:
|
|
if sparse_feature_names and "meta__user_eng_id" in sparse_feature_names:
|
|
raise ValueError("Only one source of embedding for user is supported")
|
|
else:
|
|
user_eng_embedding = features_in["user_eng_embedding"]
|
|
|
|
if "author_embedding" in features_in:
|
|
if sparse_feature_names and "meta__author_id" in sparse_feature_names:
|
|
raise ValueError("Only one source of embedding for user is supported")
|
|
else:
|
|
author_embedding = features_in["author_embedding"]
|
|
|
|
return RecapBatch(
|
|
continuous_features=features_in["continuous"],
|
|
binary_features=features_in["binary"],
|
|
discrete_features=features_in["discrete"],
|
|
sparse_features=sparse_features,
|
|
user_embedding=user_embedding,
|
|
user_eng_embedding=user_eng_embedding,
|
|
author_embedding=author_embedding,
|
|
labels=labels,
|
|
weights=features_in.get("weights", None), # Defaults to torch.ones_like(labels)
|
|
)
|
|
|
|
|
|
def _chain(param, f1, f2):
|
|
"""
|
|
Reduce multiple functions into one chained function
|
|
_chain(x, f1, f2) -> f2(f1(x))
|
|
"""
|
|
output = param
|
|
fns = [f1, f2]
|
|
for f in fns:
|
|
output = f(output)
|
|
return output
|
|
|
|
|
|
def _add_weights(inputs, tasks: Dict[str, TaskData]):
|
|
"""Adds weights based on label sampling for positive and negatives.
|
|
|
|
This is useful for numeric calibration etc. This mutates inputs.
|
|
|
|
Args:
|
|
inputs: A dictionary of strings to tensor-like structures.
|
|
tasks: A dict of string (label) to `TaskData` specifying inputs.
|
|
|
|
Returns:
|
|
A tuple of features and labels; weights are added to features.
|
|
"""
|
|
|
|
weights = []
|
|
for key, task in tasks.items():
|
|
label = inputs[key]
|
|
float_label = tf.cast(label, tf.float32)
|
|
|
|
weights.append(
|
|
float_label / task.pos_downsampling_rate + (1.0 - float_label) / task.neg_downsampling_rate
|
|
)
|
|
|
|
# Ensure we are batch-major (assumes we batch before this call).
|
|
inputs["weights"] = tf.squeeze(tf.transpose(tf.convert_to_tensor(weights)), axis=0)
|
|
return inputs
|
|
|
|
|
|
def get_datetimes(explicit_datetime_inputs):
|
|
"""Compute list datetime strings for train/validation data."""
|
|
datetime_format = "%Y/%m/%d/%H"
|
|
end = datetime.strptime(explicit_datetime_inputs.end_datetime, datetime_format)
|
|
dates = sorted(
|
|
[
|
|
(end - timedelta(hours=i + 1)).strftime(datetime_format)
|
|
for i in range(int(explicit_datetime_inputs.hours))
|
|
]
|
|
)
|
|
return dates
|
|
|
|
|
|
def get_explicit_datetime_inputs_files(explicit_datetime_inputs):
|
|
"""
|
|
Compile list of files for training/validation.
|
|
|
|
Used with DataConfigs that use the `explicit_datetime_inputs` format to specify data.
|
|
For each hour of data, if the directory is missing or empty, we increment a counter to keep
|
|
track of the number of missing data hours.
|
|
Returns only files with a `.gz` extension.
|
|
|
|
Args:
|
|
explicit_datetime_inputs: An `ExplicitDatetimeInputs` object within a `datasets.DataConfig` object
|
|
|
|
Returns:
|
|
data_files: Sorted list of files to read corresponding to data at the desired datetimes
|
|
num_hours_missing: Number of hours that we are missing data
|
|
|
|
"""
|
|
datetimes = get_datetimes(explicit_datetime_inputs)
|
|
folders = [os.path.join(explicit_datetime_inputs.data_root, datetime) for datetime in datetimes]
|
|
data_files = []
|
|
num_hours_missing = 0
|
|
for folder in folders:
|
|
try:
|
|
files = tf.io.gfile.listdir(folder)
|
|
if not files:
|
|
logging.warning(f"{folder} contained no data files")
|
|
num_hours_missing += 1
|
|
data_files.extend(
|
|
[
|
|
os.path.join(folder, filename)
|
|
for filename in files
|
|
if filename.rsplit(".", 1)[-1].lower() == "gz"
|
|
]
|
|
)
|
|
except tf.errors.NotFoundError as e:
|
|
num_hours_missing += 1
|
|
logging.warning(f"Cannot find directory {folder}. Missing one hour of data. Error: \n {e}")
|
|
return sorted(data_files), num_hours_missing
|
|
|
|
|
|
def _map_output_for_inference(
|
|
inputs, tasks: Dict[str, TaskData], preprocessor: tf.keras.Model = None, add_weights: bool = False
|
|
):
|
|
if preprocessor:
|
|
raise ValueError("No preprocessor should be used at inference time.")
|
|
if add_weights:
|
|
raise NotImplementedError()
|
|
|
|
# Add zero weights.
|
|
inputs["weights"] = tf.zeros_like(tf.expand_dims(inputs["continuous"][:, 0], -1))
|
|
for label in tasks:
|
|
del inputs[label]
|
|
return inputs
|
|
|
|
|
|
def _map_output_for_train_eval(
|
|
inputs, tasks: Dict[str, TaskData], preprocessor: tf.keras.Model = None, add_weights: bool = False
|
|
):
|
|
if add_weights:
|
|
inputs = _add_weights_based_on_sampling_rates(inputs, tasks)
|
|
|
|
# Warning this has to happen first as it changes the input
|
|
if preprocessor:
|
|
inputs = preprocessor(inputs)
|
|
|
|
label_values = tf.squeeze(tf.stack([inputs[label] for label in tasks], axis=1), axis=[-1])
|
|
|
|
for label in tasks:
|
|
del inputs[label]
|
|
|
|
return inputs, label_values
|
|
|
|
|
|
def _add_weights_based_on_sampling_rates(inputs, tasks: Dict[str, TaskData]):
|
|
"""Adds weights based on label sampling for positive and negatives.
|
|
|
|
This is useful for numeric calibration etc. This mutates inputs.
|
|
|
|
Args:
|
|
inputs: A dictionary of strings to tensor-like structures.
|
|
tasks: A dict of string (label) to `TaskData` specifying inputs.
|
|
|
|
Returns:
|
|
A tuple of features and labels; weights are added to features.
|
|
"""
|
|
weights = []
|
|
for key, task in tasks.items():
|
|
label = inputs[key]
|
|
float_label = tf.cast(label, tf.float32)
|
|
|
|
weights.append(
|
|
float_label / task.pos_downsampling_rate + (1.0 - float_label) / task.neg_downsampling_rate
|
|
)
|
|
|
|
# Ensure we are batch-major (assumes we batch before this call).
|
|
inputs["weights"] = tf.squeeze(tf.transpose(tf.convert_to_tensor(weights)), axis=0)
|
|
return inputs
|
|
|
|
|
|
class RecapDataset(torch.utils.data.IterableDataset):
|
|
def __init__(
|
|
self,
|
|
data_config: RecapDataConfig,
|
|
dataset_service: Optional[str] = None,
|
|
mode: JobMode = JobMode.TRAIN,
|
|
compression: Optional[str] = "AUTO",
|
|
repeat: bool = False,
|
|
vocab_mapper: tf.keras.Model = None,
|
|
):
|
|
logging.info("***** Labels *****")
|
|
logging.info(list(data_config.tasks.keys()))
|
|
|
|
self._data_config = data_config
|
|
self._parse_fn = get_seg_dense_parse_fn(data_config)
|
|
self._mode = mode
|
|
self._repeat = repeat
|
|
self._num_concurrent_iterators = 1
|
|
self._vocab_mapper = vocab_mapper
|
|
self.dataset_service = dataset_service
|
|
|
|
preprocessor = None
|
|
self._batch_size_multiplier = 1
|
|
if data_config.preprocess:
|
|
preprocessor = preprocessors.build_preprocess(data_config.preprocess, mode=mode)
|
|
if data_config.preprocess.downsample_negatives:
|
|
self._batch_size_multiplier = data_config.preprocess.downsample_negatives.batch_multiplier
|
|
|
|
self._preprocessor = preprocessor
|
|
|
|
if mode == JobMode.INFERENCE:
|
|
if preprocessor is not None:
|
|
raise ValueError("Expect no preprocessor at inference time.")
|
|
should_add_weights = False
|
|
output_map_fn = _map_output_for_inference # (features,)
|
|
else:
|
|
# Only add weights if there is a reason to! If all weights will
|
|
# be equal to 1.0, save bandwidth between DDS and Chief by simply
|
|
# relying on the fact that weights default to 1.0 in `RecapBatch`
|
|
# WARNING: Weights may still be added as a side effect of a preprocessor
|
|
# such as `DownsampleNegatives`.
|
|
should_add_weights = any(
|
|
[
|
|
task_cfg.pos_downsampling_rate != 1.0 or task_cfg.neg_downsampling_rate != 1.0
|
|
for task_cfg in data_config.tasks.values()
|
|
]
|
|
)
|
|
output_map_fn = _map_output_for_train_eval # (features, labels)
|
|
|
|
self._output_map_fn = functools.partial(
|
|
output_map_fn,
|
|
tasks=data_config.tasks,
|
|
preprocessor=preprocessor,
|
|
add_weights=should_add_weights,
|
|
)
|
|
|
|
sparse_feature_names = list(vocab_mapper.vocabs.keys()) if vocab_mapper else None
|
|
|
|
self._tf_dataset = self._create_tf_dataset()
|
|
|
|
self._init_tensor_spec()
|
|
|
|
def _init_tensor_spec(self):
|
|
def _tensor_spec_to_torch_shape(spec):
|
|
if spec.shape is None:
|
|
return None
|
|
shape = [x if x is not None else -1 for x in spec.shape]
|
|
return torch.Size(shape)
|
|
|
|
self.torch_element_spec = tf.nest.map_structure(
|
|
_tensor_spec_to_torch_shape, self._tf_dataset.element_spec
|
|
)
|
|
|
|
def _create_tf_dataset(self):
|
|
if hasattr(self, "_tf_dataset"):
|
|
raise ValueError("Do not call `_create_tf_dataset` more than once.")
|
|
|
|
world_size = dist.get_world_size() if dist.is_initialized() else 1
|
|
per_replica_bsz = (
|
|
self._batch_size_multiplier * self._data_config.global_batch_size // world_size
|
|
)
|
|
|
|
dataset: tf.data.Dataset = self._create_base_tf_dataset(
|
|
batch_size=per_replica_bsz,
|
|
)
|
|
|
|
if self._repeat:
|
|
logging.info("Repeating dataset")
|
|
dataset = dataset.repeat()
|
|
|
|
if self.dataset_service:
|
|
if self._num_concurrent_iterators > 1:
|
|
if not self.machines_config:
|
|
raise ValueError(
|
|
"Must supply a machine_config for autotuning in order to use >1 concurrent iterators"
|
|
)
|
|
dataset = dataset_lib.with_auto_tune_budget(
|
|
dataset,
|
|
machine_config=self.machines_config.chief,
|
|
num_concurrent_iterators=self.num_concurrent_iterators,
|
|
on_chief=False,
|
|
)
|
|
|
|
self.dataset_id, self.job_name = register_dataset(
|
|
dataset=dataset, dataset_service=self.dataset_service, compression=self.compression
|
|
)
|
|
dataset = distribute_from_dataset_id(
|
|
dataset_id=self.dataset_id, # type: ignore[arg-type]
|
|
job_name=self.job_name,
|
|
dataset_service=self.dataset_service,
|
|
compression=self.compression,
|
|
)
|
|
|
|
elif self._num_concurrent_iterators > 1:
|
|
if not self.machines_config:
|
|
raise ValueError(
|
|
"Must supply a machine_config for autotuning in order to use >1 concurrent iterators"
|
|
)
|
|
dataset = dataset_lib.with_auto_tune_budget(
|
|
dataset,
|
|
machine_config=self.machines_config.chief,
|
|
num_concurrent_iterators=self._num_concurrent_iterators,
|
|
on_chief=True,
|
|
)
|
|
|
|
# Vocabulary mapping happens on the training node, not in dds because of size.
|
|
if self._vocab_mapper:
|
|
dataset = dataset.map(self._vocab_mapper)
|
|
|
|
return dataset.prefetch(world_size * 2)
|
|
|
|
def _create_base_tf_dataset(self, batch_size: int):
|
|
if self._data_config.inputs:
|
|
glob = self._data_config.inputs
|
|
filenames = sorted(tf.io.gfile.glob(glob))
|
|
elif self._data_config.explicit_datetime_inputs:
|
|
num_missing_hours_tol = self._data_config.explicit_datetime_inputs.num_missing_hours_tol
|
|
filenames, num_hours_missing = get_explicit_datetime_inputs_files(
|
|
self._data_config.explicit_datetime_inputs,
|
|
increment="hourly",
|
|
)
|
|
if num_hours_missing > num_missing_hours_tol:
|
|
raise ValueError(
|
|
f"We are missing {num_hours_missing} hours of data"
|
|
f"more than tolerance {num_missing_hours_tol}."
|
|
)
|
|
elif self._data_config.explicit_date_inputs:
|
|
num_missing_days_tol = self._data_config.explicit_date_inputs.num_missing_days_tol
|
|
filenames, num_days_missing = get_explicit_datetime_inputs_files(
|
|
self._data_config.explicit_date_inputs,
|
|
increment="daily",
|
|
)
|
|
if num_days_missing > num_missing_days_tol:
|
|
raise ValueError(
|
|
f"We are missing {num_days_missing} days of data"
|
|
f"more than tolerance {num_missing_days_tol}."
|
|
)
|
|
else:
|
|
raise ValueError(
|
|
"Must specifiy either `inputs`, `explicit_datetime_inputs`, or `explicit_date_inputs` in data_config"
|
|
)
|
|
|
|
num_files = len(filenames)
|
|
logging.info(f"Found {num_files} data files")
|
|
if num_files < 1:
|
|
raise ValueError("No data files found")
|
|
|
|
if self._data_config.num_files_to_keep is not None:
|
|
filenames = filenames[: self._data_config.num_files_to_keep]
|
|
logging.info(f"Retaining only {len(filenames)} files.")
|
|
|
|
filenames_ds = (
|
|
tf.data.Dataset.from_tensor_slices(filenames).shuffle(len(filenames))
|
|
# Because of drop_remainder, if our dataset does not fill
|
|
# up a batch, it will emit nothing without this repeat.
|
|
.repeat(-1)
|
|
)
|
|
|
|
if self._data_config.file_batch_size:
|
|
filenames_ds = filenames_ds.batch(self._data_config.file_batch_size)
|
|
|
|
def per_shard_dataset(filename):
|
|
ds = tf.data.TFRecordDataset([filename], compression_type="GZIP")
|
|
return ds.prefetch(4)
|
|
|
|
ds = filenames_ds.interleave(
|
|
per_shard_dataset,
|
|
block_length=4,
|
|
deterministic=False,
|
|
num_parallel_calls=self._data_config.interleave_num_parallel_calls
|
|
or tf.data.experimental.AUTOTUNE,
|
|
)
|
|
|
|
# Combine functions into one map call to reduce overhead.
|
|
map_fn = functools.partial(
|
|
_chain,
|
|
f1=self._parse_fn,
|
|
f2=self._output_map_fn,
|
|
)
|
|
|
|
# Shuffle -> Batch -> Parse is the correct ordering
|
|
# Shuffling needs to be performed before batching otherwise there is not much point
|
|
# Batching happens before parsing because tf.Example parsing is actually vectorized
|
|
# and works much faster overall on batches of data.
|
|
ds = (
|
|
# DANGER DANGER: there is a default shuffle size here.
|
|
ds.shuffle(self._data_config.examples_shuffle_buffer_size)
|
|
.batch(batch_size=batch_size, drop_remainder=True)
|
|
.map(
|
|
map_fn,
|
|
num_parallel_calls=self._data_config.map_num_parallel_calls
|
|
or tf.data.experimental.AUTOTUNE,
|
|
)
|
|
)
|
|
|
|
if self._data_config.cache:
|
|
ds = ds.cache()
|
|
|
|
if self._data_config.ignore_data_errors:
|
|
ds = ds.apply(tf.data.experimental.ignore_errors())
|
|
|
|
options = tf.data.Options()
|
|
options.experimental_deterministic = False
|
|
ds = ds.with_options(options)
|
|
|
|
return ds
|
|
|
|
def _gen(self):
|
|
for x in self._tf_dataset:
|
|
yield to_batch(x)
|
|
|
|
def to_dataloader(self) -> Dict[str, torch.Tensor]:
|
|
return torch.utils.data.DataLoader(self, batch_size=None)
|
|
|
|
def __iter__(self):
|
|
return iter(self._gen())
|