R interface to TensorFlow Dataset API

Overview

The TensorFlow Dataset API provides various facilities for creating scalable input pipelines for TensorFlow models, including:

  • Reading data from a variety of formats including CSV files and TFRecords files (the standard binary format for TensorFlow training data).

  • Transforming datasets in a variety of ways including mapping arbitrary functions against them.

  • Shuffling, batching, and repeating datasets over a number of epochs.

  • Streaming interface to data for reading arbitrarily large datasets.

  • Reading and transforming data are TensorFlow graph operations, so are executed in C++ and in parallel with model training.

The R interface to TensorFlow datasets provides access to the Dataset API, including high-level convenience functions for easy integration with the keras and tfestimators R packages.

Installation

To use tfdatasets you need to install both the R package as well as TensorFlow itself.

First, install the tfdatasets R package from CRAN as follows:

Then, use the install_tensorflow() function to install TensorFlow:

Creating a Dataset

To create a dataset, use one of the dataset creation functions. Dataset can be created from delimted text files, TFRecords files, as well as from in-memory data.

Text Files

For example, to create a dataset from a text file, first create a specification for how records will be decoded from the file, then call text_line_dataset() with the file to be read and the specification:

TensorFlow Dataset
Petal.Length : <tf.float32> 1.4 1.4 1.3 1.5 1.4 1.7 1.4 1.5 1.4 1.5 1.5 1.6 1.4 1.1 1.2 1...
Sepal.Length : <tf.float32> 5.1 4.9 4.7 4.6 5 5.4 4.6 5 4.4 4.9 5.4 4.8 4.8 4.3 5.8 5.7 5...
Petal.Width  : <tf.float32> 0.2 0.2 0.2 0.2 0.2 0.4 0.3 0.2 0.2 0.1 0.2 0.2 0.1 0.1 0.2 0...
Sepal.Width  : <tf.float32> 3.5 3 3.2 3.1 3.6 3.9 3.4 3.4 2.9 3.1 3.7 3.4 3 3 4 4.4 3.9 3...
Species      : <tf.int32>   1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 ..

In the example above, the csv_record_spec() function is passed an example file which is used to automatically detect column names and types (done by reading up to the first 1,000 lines of the file). You can also provide explicit column names and/or data types using the names and types parameters (note that in this case we don’t pass an example file):

Note that we’ve also specified skip = 1 to indicate that the first row of the CSV that contains column names should be skipped.

Supported column types are integer, double, and character. You can also provide types in a more compact form using single-letter abbreviations (e.g. types = "dddi"). For example:

mtcars_spec <- csv_record_spec("mtcars.csv", types = "dididddiiii")

Parallel Decoding

Decoding lines of text into a record can be computationally expensive. You can parallelize these computations using the parallel_records parameter. For example:

dataset <- text_line_dataset("iris.csv", record_spec = iris_spec, parallel_records = 4)

You can also parallelize the reading of data from storage by requesting that a buffer of records be prefected. You do this with the dataset_prefetch() function. For example:

dataset <- text_line_dataset("iris.csv", record_spec = iris_spec, parallel_records = 4) %>% 
  dataset_prefetch(1000)

If you have multiple input files, you can also parallelize reading of these files both across multiple machines (sharding) and/or on multiple threads per-machine (parallel reads with interleaving). See the section on Reading Multiple Files below for additional details.

TFRecords Files

You can read datasets from TFRecords files using the tfrecord_dataset() function.

In many cases you’ll want to map the records in the dataset into a set of named columns. You can do this using the dataset_map() function along with the tf$parse_single_example() function. for example:

SQLite Databases

You can read datasets from SQLite databases using the sqlite_dataset() function. Note that support for SQLite databases required TensorFlow v1.7 as well as the development version of the tfdatasets package, which you can install via devtools::install_github("rstudio/tfdatasts").

To use sqlite_dataset() you provide the filename of the database, a SQL query to execute, and sql_record_spec() that describes the names and TensorFlow types of columns within the query. For example:

TensorFlow Dataset
disp : <tf.float64> 21 21 22.8 21.4 18.7 18.1 14.3 24.4 22.8 19.2 17.8 16.4 17.3 15.2 10.4..
drat : <tf.int32>   6 6 4 6 8 6 8 4 4 6 6 8 8 8 8 8 8 4 4 4 4 8 8 8 8 4 4 4 8 6 8 4
vs   : <tf.float64> 160 160 108 258 360 225 360 146.7 140.8 167.6 167.6 275.8 275.8 275.8 ..
gear : <tf.int32>   110 110 93 110 175 105 245 62 95 123 123 180 180 180 205 215 230 66 52..
mpg  : <tf.float64> 3.9 3.9 3.85 3.08 3.15 2.76 3.21 3.69 3.92 3.92 3.92 3.07 3.07 3.07 2...
qsec : <tf.float64> 2.62 2.88 2.32 3.21 3.44 3.46 3.57 3.19 3.15 3.44 3.44 4.07 3.73 3.78 ..
hp   : <tf.float64> 16.5 17 18.6 19.4 17 20.2 15.8 20 22.9 18.3 18.9 17.4 17.6 18 18 17.8 ..
am   : <tf.int32>   0 0 1 1 0 1 0 1 1 1 1 0 0 0 0 0 0 1 1 1 1 0 0 0 0 1 0 1 0 0 0 1
wt   : <tf.int32>   1 1 1 0 0 0 0 0 0 0 0 0 0 0 0 0 0 1 1 1 0 0 0 0 0 1 1 1 1 1 1 1
carb : <tf.int32>   4 4 4 3 3 3 3 4 4 4 4 3 3 3 3 3 3 4 4 4 3 3 3 3 3 4 5 5 5 5 5 4
cyl  : <tf.int32>   4 4 1 1 2 1 4 2 2 4 4 3 3 3 4 4 4 1 2 1 1 2 2 4 2 1 2 2 4 6 8 2

Note that for floating point data you must use tf$float64 (reading tf$float32 is not supported for SQLite databases).

Transformations

Mapping

You can map arbitrary transformation functions onto dataset records using the dataset_map() function. For example, to transform the “Species” column into a one-hot encoded vector you would do this:

Note that while dataset_map() is defined using an R function, there are some special constraints on this function which allow it to execute not within R but rather within the TensorFlow graph.

For a dataset created with the csv_dataset() function, the passed record will be named list of tensors (one for each column of the dataset). The return value should be another set of tensors which were created from TensorFlow functions (e.g. tf$one_hot as illustrated above). This function will be converted to a TensorFlow graph operation that performs the transformation within native code.

Parallel Mapping

If these transformations are computationally expensive they can be executed on multiple threads using the num_parallel_calls parameter. For example:

You can control the maximum number of processed elements that will be buffered when processing in parallel using the dataset_prefetch() transformation. For example:

Filtering

You can filter the elements of a dataset using the dataset_filter() function, which takes a predicate function that returns a boolean tensor for records that should be included. For example:

Note that the functions used inside the predicate must be tensor operations (e.g. tf$not_equal, tf$less, etc.). R generic methods for relational operators (e.g. <, >, <=, etc.) and logical operators (e.g. !, &, |, etc.) are provided so you can use shorthand syntax for most common comparisons (as illustrated above).

Features and Response

A common transformation is taking a column oriented dataset (e.g. one created by csv_dataset() or tfrecord_dataset()) and transforming it into a two-element list with features (“x”) and response (“y”). You can use the dataset_prepare() function to do this type of transformation. For example:

mtcars_dataset <- text_line_dataset("mtcars.csv", record_spec = mtcars_spec) %>% 
  dataset_prepare(x = c(mpg, disp), y = cyl)

iris_dataset <- text_line_dataset("iris.csv", record_spec = iris_spec) %>% 
  dataset_prepare(x = -Species, y = Species)

The dataset_prepare() function also accepts standard R formula syntax for defining features and response:

mtcars_dataset <- text_line_dataset("mtcars.csv", record_spec = mtcars_spec) %>% 
  dataset_prepare(cyl ~ mpg + disp)

Shuffling and Batching

There are several functions which control how batches are drawn from the dataset. For example, the following specifies that data will be drawn in batches of 128 from a shuffled window of 1000 records, and that the dataset will be repeated for 10 epochs:

Complete Example

Here’s a complete example of using the various dataset transformation functions together. We’ll read the mtcars dataset from a CSV, filter it on some threshold values, map it into x and y components for modeling, and specify desired shuffling and batch iteration behavior:

dataset <- text_line_dataset("mtcars.csv", record_spec = mtcars_spec) %>%
  dataset_prefetch(1000) %>% 
  dataset_filter(function(record) {
    record$mpg >= 20 & record$cyl >= 6L
  }) %>% 
  dataset_prepare(cyl ~ mpg + disp) %>% 
  dataset_shuffle(1000) %>% 
  dataset_batch(128) %>% 
  dataset_repeat(10)

Reading Datasets

You read batches of data from a dataset by using tensors that yield the next batch. You can obtain this tensor from a dataset via the next_batch() function. For example:

dataset <- text_line_dataset("mtcars.csv", record_spec = mtcars_spec) %>% 
  dataset_prepare(cyl ~ mpg + disp) %>% 
  dataset_shuffle(20) %>% 
  dataset_batch(5)
batch <- next_batch(dataset)
batch
$x
Tensor("IteratorGetNext_13:0", shape=(?, 2), dtype=float32)

$y
Tensor("IteratorGetNext_13:1", shape=(?,), dtype=int32)

As you can see batch isn’t the data itself but rather a tensor that will yield the next batch of data when it is evaluated:

$x
     [,1] [,2]
[1,] 21.0  160
[2,] 21.0  160
[3,] 22.8  108
[4,] 21.4  258
[5,] 18.7  360

$y
[1] 6 6 4 6 8

Dataset Iteration

If you are iterating over an entire dataset by evaluating the next_batch() tensor you will need to determine at what point to stop iteration. There are a couple of possible approaches to controlling/detecting when iteration should end.

One approach is to create a dataset that yields batches infinitely (traversing the dataset multiple times with different batches randomly drawn). In this case you’d use another mechanism like a global step counter or detecting a learning plateau:

library(tfdatasets)
dataset <- text_line_dataset("mtcars.csv", record_spec = mtcars_spec) %>% 
  dataset_prepare(x = c(mpg, disp), y = cyl) %>% 
  dataset_shuffle(5000) %>% 
  dataset_batch(128) %>% 
  dataset_repeat() # repeat infinitely

batch <- next_batch(dataset)

steps <- 200
for (i in 1:steps) {
  # use batch$x and batch$y tensors
}

The call to dataset_repeat() with no count parameter requests that the dataset be traversed infinitely.

Another approach is to detect when all batches have been yielded from the dataset. When the tensor reaches the end of iteration a runtime error will occur. You can catch and ignore the error when it occurs by wrapping your iteration code in the with_dataset() function:

Using with tfestimators

Models created with tfestimators use an input function to consume data for training, evaluation, and prediction. For example, here is an example of using an input function to feed data from an in-memory R data frame to an estimators model:

model %>% train(
  input_fn(mtcars, features = c(mpg, disp), response = cyl,
           batch_size = 128, epochs = 3)
)

If you are using tfdatasets with the tfestimators package, you can create an estimators input function directly from a dataset as follows:

dataset <- text_line_dataset("mtcars.csv", record_spec = mtcars_spec) %>% 
  dataset_batch(128) %>% 
  dataset_repeat(3)

model %>% train(
  input_fn(dataset, features = c(mpg, disp), response = cyl)
)

Note that we don’t use the dataset_prepare() or next_batch() functions in this example. Rather, these functions are used under the hood to provide the input_fn interface expected by tfestimators models.

As with dataset_prepare(), you can alternatively use an R formula to specify features and response:

model %>% train(
  input_fn(dataset, cyl ~ mpg + disp)
)

Using with Keras

Keras models are often trained by passing in-memory arrays directly to the fit function. For example:

However, this requires loading data into an R data frame or matrix before calling fit. You can use the train_on_batch() function to stream data one batch at a time, however the reading and processing of the input data is still being done serially and outside of native code.

Alternatively, Keras enables you to pass a dataset directly as the generator argument to fit_generator() and validate_generator(). Here’s a complete example that uses datasets to read from TFRecord files containing MNIST digits:

Note that all data preprocessing (e.g. one-hot encoding of the response variable) is done within the dataset_map() operation.

Reading Multiple Files

If you have multiple input files you can process them in parallel both across machines (sharding) and/or on multiple threads per-machine (parallel reads with interleaving). The read_files() function provides a high-level interface to parallel file reading.

The read_files() function takes a set of files and a read function along with various options to orchestrate parallel reading. For example, the following function reads all CSV files in a directory using the text_line_dataset() function:

dataset <- read_files("data/*.csv", text_line_dataset, record_spec = mtcars_spec,
                      parallel_files = 4, parallel_interleave = 16) %>% 
  dataset_prefetch(5000) %>% 
  dataset_shuffle(1000) %>% 
  dataset_batch(128) %>% 
  dataset_repeat(3)

The parallel_files argument requests that 4 files be processed in parallel and the parallel_interleave argument requests that blocks of 16 consecutive records from each file be interleaved in the resulting dataset.

Note that because we are processing files in parallel we do not pass the parallel_records argument to text_line_dataset(), since we are already parallelizing at the file level.

Multiple Machines

If you are training on multiple machines and the training supervisor passes a shard index to your training script, you can also parallelizing reading by sharding the file list. For example: