tfjs Documentation

A WebGL accelerated JavaScript library for training and deploying ML models.

RAW Doc

Docker & Self-Hosted Deployment Verified OCI

Deploy tfjs on any server, homelab, or cloud container platform.

1-Click Terminal Command
bash (docker run)
docker run -d --name tfjs --restart always tensorflow/tfjs:latest

1. Project Overview & Quickstart (tensorflow/tfjs)

File: docs/OPTIMIZATION_PURE_GPU_PIPELINE.md

Optimization Topic - How to write a high performance GPU pipeline

This document shows how to write a high performance GPU pipeline with TF.js. A typical ML pipeline is composed of some ML steps (e.g. taking some inputs, running a model and getting outputs) and non-ML steps (e.g. image processing, rendering, etc.). General rule of thumb is that if the whole pipeline can run on GPU without downloading any data to CPU, it is usually much faster than a fragmented pipeline that requires data transfer between CPU and GPU. This additional time for GPU to CPU and CPU to GPU sync adds to the pipeline latency.

In the latest TF.js release (version 3.13.0), we provide a new tensor
API that allows data to stay in GPU. Our users are pretty faimilar with
tensor.dataSync() and tensor.data(), however both will download data from GPU to CPU. In the latest release, we have added tensor.dataToGPU(),
which returns a GPUData type.

For webgl backend

The example below shows how to get the tensor texture by calling
dataToGPU. Also see this example code for details.

javascript
// Getting the texture that holds the data on GPU.
// data has below fields:
// {texture, texShape, tensorRef}
// texture: A WebGLTexture.
// texShape: the [height, width] of the texture.
// tensorRef: the tensor associated with the texture.
//
// Here the tensor has a shape of [1, videoHeight, videoWidth, 4],
// which represents an image. In this case, we can specify the texture
// to use the image's shape as its shape.
const data =
      tensor.dataToGPU({customTexShape: [videoHeight, videoWidth]});

// Once we have the texture, we can bind it and pass it to the downstream
// webgl processing steps to use the texture.
gl.bindTexture(gl.TEXTURE_2D, data.texture);

// Some webgl processing steps.

// Once all the processing is done, we can render the result on the canvas.
gl.blitFramebuffer(0, 0, videoWidth, videoHeight, 0, videoHeight, videoWidth, 0, gl.COLOR_BUFFER_BIT, gl.LINEAR);

// Remember to dispose the texture after use, otherwise, there will be
// memory leak.
data.tensorRef.dispose();

The texture from dataToGPU() has a specific format. The data is densely
packed, meaning all four channels (r, g, b, a) in a texel is used to store
data. This is how the texture with tensor shape = [2, 3, 4] looks like:

text
000|001|002|003   010|011|012|013   020|021|022|023

    100|101|102|103   110|111|112|113   120|121|122|123

So if a tensor's shape is [1, height, width, 4] or [height, width, 4], the way we store the data aligns with how image is stored, therefore if the
downstream processing assumes the texture is an image, it doesn't need to
do any additional coordination conversion. Each texel is mapped to a CSS pixel in an image. However, if the tensor has other shapes, downstream
processing steps may need to reformat the data onto another texture. The texShape field has the [height, width] information of the texture, which can be used for transformation.

For webgpu backend

The example below shows how to get the tensor buffer by calling
dataToGPU. Also see this example code for details.

javascript
// Getting the buffer that holds the data on GPU.
// data has below fields:
// {buffer, bufSize, tensorRef}
// buffer: A GPUBuffer.
// bufSize: the size of the buffer.
// tensorRef: the tensor associated with the buffer.
//
// Unlike webgl backend, There is no parameter for dataToGPU API on the webgpu
// backend, and the size of returned buffer is the same as the tensor size.
const data = tensor.dataToGPU();

// Once we have the buffer, we can bind it and pass it to the downstream
// webgpu processing steps to use the buffer.
const uniformBindGroup = device.createBindGroup({
  layout: pipeline.getBindGroupLayout(0),
  entries: [
    {
      binding: 1,
      resource: {
        buffer: data.buffer,
      },
    },
    ...
  ],
});

// Some webgpu processing steps.

// Defines the mapping between resources of all GPUBindGroup objects
passEncoder.setBindGroup(0, uniformBindGroup);

// Some webgpu processing steps.

// Remember to dispose the buffer after use, otherwise, there will be
// memory leak.
data.tensorRef.dispose();

File: e2e/benchmarks/browserstack-benchmark/README.md

Benchmark on multiple devices

:warning: To use this tool, you need to sign up for BrowserStack's Automate service.

The Multi-device benchmark tool can benchmark the performance (time, memory) of model inference on a collection of remote devices. Using this tool you will be able to:

  • Select a collection of BrowserStack devices, based on the following fields:
    • OS
    • OS version
    • Browser
    • Browser version
    • Device
  • Select a backend:
    • WASM
    • WebGL
    • CPU
  • Set the number of rounds for model inference.
  • Select a model to benchmark.

Usage

  1. Export your access key of BrowserStack's Automate service:
shell
export BROWSERSTACK_USERNAME=YOUR_USERNAME
export BROWSERSTACK_ACCESS_KEY=YOUR_ACCESS_KEY
  1. Download and run the tool:
shell
git clone https://github.com/tensorflow/tfjs.git
cd tfjs/e2e/benchmarks/browserstack-benchmark
yarn install

node app.js

Then you can see > Running socket on port: 8001 on your Command-line interface.

  1. Open http://localhost:8001/ and start to benchmark.
    3.1 If you want to benchmark code snippet. Please update benchmarkCodeSnippet with your code snippet before running node app.js and select codeSnippet in model name:

  2. When the benchmark is complete, you can see the benchmark results in the webpage, like:

Command Line Arguments

The following are supported options arguments which trigger options features:

  • --benchmarks

    • Runs benchmarks from a user-specified, pre-configured JSON file.
    shell
    node app.js --benchmarks=relative_file_path.json

    A pre-configuration file consists of a JSON object with the following format:

    text
    {
      "benchmark": {
        "model": ["model_name"], //List of one or more custom or official models to be benchmarked
        "numRuns": positive_integer,
        "backend": ["backend_name"] //List of one or more backends to be benchmarked
      },
      "browsers": {
        "local": {},  // Benchmark on your local device
        "unique_identifier_laptop_or_desktop": {
          "base": "BrowserStack",
          "browser": "browser_name",
          "browser_version": "browser_version",
          "os": "os_name",
          "os_version": "os_version",
          "device": null
        },
        "unique_identifier_mobile_device": {
          "base": "BrowserStack",
          "browser": "iphone_or_android",
          "browser_version": null,
          "os": "os_name",
          "os_version": "os_version",
          "device": "device_name"
        }
      }
    }

    Each model in the model list will be run on each backend in the backend list. Each model-backend combination will run on every browser. If you would like to test specific backends on specific models, the recommended method is to create multiple configuration files.

    For more examples of documentation, refer to the links below:
    List of officially supported TFJS browsers
    Example benchmark pre-configuration

  • --cloud

    • Runs GCP compatible version of benchmarking by blocking the local server.
    shell
    node app.js --cloud
  • --firestore

    • Pushes successful benchmark results to a Firestore database.
    shell
    node app.js --firestore
  • --h, --help

    • Shows help menu and all optional arguments in the shell window.
    shell
    node app.js --h

    or

    shell
    node app.js --help
  • --period

    • Runs a part of models specified in --benchmarks's file in a cycle and the part of models to run is determined by the date of a month. The value could be or 1~31 (representing Sunday to Saturday). This argument takes effect only if --benchmarks is set.
    shell
    node app.js --period=15
  • --date

    • Set the date for selecting models and this works only if --period is set. The value could be 1~31. If it is not declared, the date will the real date at runtime.
    shell
    node app.js --period=15 --date=1
  • --maxBenchmarks

    • Sets maximum for number of benchmarks run in parallel. Expects a positive integer.
    shell
    node app.js --maxBenchmarks=positive_integer
  • --maxTries

    • Sets maximum for number of tries a given benchmark has to succeed. Expects a positive integer.
    shell
    node app.js --maxTries=positive_integer
  • --outfile

    • Writes results to an accessible external file, benchmark_results.js or benchmark_results.json. Expects 'html' or 'json'. If you set it as 'html', benchmark_results.js will be generated and you could review the benchmark results by openning benchmark_result.html file.
    shell
    node app.js --outfile=js
  • --v, --version

    • Shows node version in use.
    shell
    node app.js --v

    or

    shell
    node app.js --version
  • --localBuild

    • Uses local build dependencies, instead of public CDNs. (When using localBuild targets, please make sure you have built the targets (eg. run yarn build-individual-link-package tfjs-backend-webgl) you need.)
    shell
    node app.js --localBuild=core,webgl,wasm,cpu,layers,converter,automl
  • --npmVersion

    • Specify the npm version of TFJS library to benchmark. By default the latest version of TFJS will be benchmarked.
shell
node app.js --npmVersion=4.4.0

Custom model

The custom model is supported, but is constrained by:

If you want to benchmark more complex models with customized input preprocessing logic, you need to add your model with load and predictFunc methods into tfjs/e2e/benchmarks/model_config.js, following this example PR.

About this tool

The tool contains:

  • A test runner - Karma:
  • A node server. app.js runs the test runner and send the benchmark results back to the webpage.
  • A webpage.

Thanks, BrowserStack, for providing supports.


File: e2e/benchmarks/local-benchmark/README.md

Benchmark custom models

The custom model in the local benchmark tool currently only supports tf.GraphModel or tf.LayersModel.

If you want to benchmark more complex TensorFlow.js models with customized input preprocessing logic, you need to implement load and predictFunc methods, following this example PR.

Models in local file system

If you have a TensorFlow.js model in local file system, you can benchmark it by: locally host the local benchmark tool and the model on a http server. In addition, if the online local benchmark tool is blocked by CORS problems when fetching custom models, this solution also works.

Example

You can benchmark the MobileNet model in local file system through the following steps:

  1. Download the tool.
shell
git clone https://github.com/tensorflow/tfjs.git
cd tfjs/e2e/benchmarks/
  1. Download the model.
shell
wget -O model.tar.gz "https://tfhub.dev/google/tfjs-model/imagenet/mobilenet_v2_130_224/classification/3/default/1?tfjs-format=compressed"
mkdir model
tar -xf model.tar.gz -C model/
  1. Run a http server to host the model and the local benchmark tool.
text
npx http-server
  1. Open http://127.0.0.1:8080/local-benchmark/ through the browser.
  2. Select custom in the models field.
  3. Fill http://127.0.0.1:8080/model/model.json into the modelUrl field.
  4. Run benchmark.

Paths to custom models

The benchmark tool suopports three kinds of paths to the custom models.

URL

Examples:

LocalStorage

Store the model in LocalStorage at first. Run the following codes in the browser console:

javascript
const localStorageModel = tf.sequential(
     {layers: [tf.layers.dense({units: 1, inputShape: [3]})]});
const saveResults = await localStorageModel.save('localstorage://my-model-1');

Then use "localstorage://my-model-1" as the custom model URL.

IndexDB

Store the model in IndexDB at first. Run the following codes in the browser console:

javascript
const indexDBModel = tf.sequential(
     {layers: [tf.layers.dense({units: 1, inputShape: [3]})]});
const saveResults = await indexDBModel.save('indexeddb://my-model-1');

Then use "indexeddb://my-model-1" as the custom model URL.

Variable input shapes

If input shapes for you model contain dynamic dimension (i.e. for mobilenet
shape = [-1, 224, 224, 3]), you are requires to set it to a valid shape
[1, 224, 224, 3] before you can perform the benchmark.
In the Inputs section you will see an input box for you to update the shape.
Once the shape is set, you can click the 'Run benchmark' button again to run
the benchmark.

Benchmark test

It's easy to set up a web server to host benchmarks and run against them via e2e/benchmarks/local-benchmark/index.html. You can manually specify the optional url parameters as needed. Here are the list of supported url parameters:

  • Model related parameters:

    architecture: same as architecture (only certain models has it, such as MobileNetV3 and posenet)
    benchmark: same as models
    inputSize: same as inputSizes
    inputType: same as inputTypes
    modelUrl: same as modelUrl, for custom models only
    ${InputeName}Shape: the input shape array, separated by comma, for custom models only. For example, bodypix's graph model has an input named sub_2, then users could add 'sub_2Shape=1,1,1,3' in the URL to populate its shape.

  • Environment related parameters:

    backend: same as backend
    localBuild: local build name list, separated by comma. The name is in short form (in general the name without the tfjs- and backend- prefixes, for example webgl for tfjs-backend-webgl, core for tfjs-core). Example: 'webgl,core'.
    run: same as numRuns
    task: correctness to "Test correctness" or performance to "Run benchmark"
    warmup: same as numWarmups


File: e2e/README.md

Usage

This package contains integration and e2e tests for TensorFlow.js.

To run the tests in local:

Filter tests by tag:

js
export TAGS=#SMOKE,#REGRESSION,#GOLDEN
yarn test

Filter tests by grep:

js
yarn test --grep cpu

Add new tests:

When creating new test, add at least one tag to the test description.

Supported tags:

  • SMOKE:
    Smoke tests should be light weight. Run in every PR and nightly
    builds. Criteria for smoke test is that the test should run fast and only
    test critical CUJ.
  • REGRESSION:
    Regression tests compare results across backends, previous
    builds, with other platform, etc. Run in nightly builds. Need additional
    steps to run in local.
  • GOLDEN:
    Golden tests validate pretrained model outputs against the prebuilt golden
    outputs. Need additional steps to run in local.

To add a new tag: Extend the TAGS list in integration_tests/util


File: scripts/cloud_funcs/README.md

This directory contains the following Google Cloud Functions.

`trigger_nightly`

Programmatically triggers a Cloud Build on master. This function is called by the Cloud Scheduler around 4am "America/New York" time every day (configurable via the Cloud Scheduler UI).
You can also trigger the function manually via the Cloud UI.

Command to re-deploy:

sh
gcloud functions deploy nightly_tfjs \
  --runtime nodejs14 \
  --trigger-topic nightly_tfjs

If a build was triggered by nightly, there is a substitution variable _NIGHTLY=true.
You can forward the substitution as the NIGHTLY environment variable so the scripts can use it, by specifying env: ['NIGHTLY=$_NIGHTLY'] in cloudbuild.yml. E.g. integration_tests/benchmarks/benchmark_cloud.sh uses the NIGHTLY bit to always run on nightly.

`send_email`

Sends an email and a chat message with the nightly build status. Every build sends a message to the cloud-builds topic with its build information. The send_email function is subscribed to that topic and ignores all builds (e.g. builds triggered by pull requests) except for the nightly build and sends an email to an internal mailing list with its build status around 4:40am.

Command to re-deploy:

sh
gcloud functions deploy send_email \
  --runtime nodejs14 \
  --stage-bucket learnjs-174218_cloudbuild \
  --trigger-topic cloud-builds \
  --set-env-vars MAILGUN_API_KEY="[API_KEY_HERE]",HANGOUTS_URL="[URL_HERE]"

`sync_reactnative`

Makes a request to browserStack to sync the current build of the tfjs-react-native integration app to browserstack. The app itself is stored in a GCP bucket. This needs to be done at least once every 30 days and is triggered via cloud scheduler via the sync_reactnative topic.
Currently set to run weekly on Thursdays at 3AM.

Command to re-deploy:

sh
gcloud functions deploy sync_reactnative \
  --runtime nodejs14 \
  --trigger-topic sync_reactnative \
  --set-env-vars HANGOUTS_URL="[URL_HERE]",BOTS_HANGOUTS_URL="[URL_HERE]"

The pipeline

The pipeline looks like this:

  1. At 4am, Cloud Scheduler writes to nightly_tfjs topic
  2. That triggers the nightly_tfjs function, which starts a build programmatically
  3. That build runs and writes its status to cloud-builds topic
  4. That triggers the send_email function, which sends email and chat with the build status.

File: tfjs-automl/README.md

AutoML Edge API

This packages provides a set of APIs to load and run models produced by AutoML
Edge.

Installation

If you are using npm/yarn

sh
npm i @tensorflow/tfjs-automl

If you are using CDN:

html
<script src="https://cdn.jsdelivr.net/npm/@tensorflow/tfjs-automl"></script>

We support the following types of AutoML Edge models:

  1. Image classification
  2. Object detection

Image classification

AutoML Image classification model will output the following set of files:

  • model.json, the model topology
  • dict.txt, a newline-separated list of labels
  • One or more of *.bin files which hold the weights

Make sure you can access those files as static assets from your web app by serving them locally or on Google Cloud Storage.

Demo

The image classification demo lives in
demo/img_classification. To run it:

sh
cd demo/img_classification
yarn
yarn watch

This will start a local HTTP server on port 1234 that serves the demo.

Loading the model

js
import * as automl from '@tensorflow/tfjs-automl';
const modelUrl = 'model.json'; // URL to the model.json file.
const model = await automl.loadImageClassification(modelUrl);

If you do not want (or cannot) load the model over HTTP you can also load the model separately and directly use the constructor.
This is particularly relevant for non-browser platforms.

The following pseudocode demonstrates this approach:

js
import * as automl from '@tensorflow/tfjs-automl';
import * as tf from '@tensorflow/tfjs';
// You can load the graph model using any IO handler
const graphModel = await tf.loadGraphModel(string|io.IOHandler); // a url or ioHandler instance
// You can load the dictionary using any api available to the platform
const dict = loadDictionary("path/to/dict.txt");
const model = new automl.ImageClassificationModel(graphModel, dict);

Making a prediction

The AutoML library takes care of any image preprocessing
(normalize, resize, crop). The input img you provide can be
HTMLImageElement,
HTMLCanvasElement,
HTMLVideoElement,
ImageData or
a 3D Tensor:

html
js
const img = document.getElementById('img');
const options = {centerCrop: true};
const predictions = await model.classify(img, options);

options is optional and has the following properties:

  • centerCrop - Defaults to true. Since the ML model expects a square image,
    we need to resize. If true, the image will be cropped first to the center before
    resizing.

The result predictions is a sorted list of predicted labels and their
probabilities:

js
[
  {label: "daisy", prob: 0.931},
  {label: "dandelion", prob: 0.027},
  {label: "roses", prob: 0.013},
  ...
]

Advanced usage

Advanced users can access the underlying
GraphModel via
model.graphModel. The GraphModel allows users to call lower level methods
such as predict(), execute() and executeAsync() which return tensors.

model.dictionary gives you access to the ordered list of labels.

Object detection

AutoML Object detection model will output the following set of files:

  • model.json, the model topology
  • dict.txt, a newline-separated list of labels
  • One or more of *.bin files which hold the weights

Make sure you can access those files as static assets from your web app by serving them locally or on Google Cloud Storage.

Demo

The object detection demo lives in
demo/object_detection. To run it:

sh
cd demo/object_detection
yarn
yarn watch

This will start a local HTTP server on port 1234 that serves the demo.

Loading the model

js
import * as automl from '@tensorflow/tfjs-automl';
const modelUrl = 'model.json'; // URL to the model.json file.
const model = await automl.loadObjectDetection(modelUrl);

If you do not want (or cannot) load the model over HTTP you can also load the model separately and directly use the constructor.
This is particularly relevant for non-browser platforms.

The following pseudocode demonstrates this approach:

js
import * as automl from '@tensorflow/tfjs-automl';
import * as tf from '@tensorflow/tfjs';
// You can load the graph model using any IO handler
const graphModel = await tf.loadGraphModel(string|io.IOHandler); // a url or ioHandler instance
// You can load the dictionary using any api available to the platform
const dict = readDictionary("path/to/dict.txt");
const model = new automl.ObjectDetectionModel(graphModel, dict);

Making a prediction

The AutoML library takes care of any image preprocessing
(normalize, resize, crop). The input img you provide can be
HTMLImageElement,
HTMLCanvasElement,
HTMLVideoElement,
ImageData or
a 3D Tensor:

html
js
const img = document.getElementById('img');
const options = {score: 0.5, iou: 0.5, topk: 20};
const predictions = await model.detect(img, options);

options is optional and has the following properties:

  • score - Probability score between 0 and 1. Defaults to 0.5. Boxes with score lower than this threshold will be ignored.
  • topk - Only the topk most likely objects are returned. The actual number of objects might be less than this number.
  • iou - Intersection over union threshold. IoU is a metric between 0 and 1 used to measure the overlap of two boxes. The predicted boxes will not overlap more than the specified threshold.

The result predictions is a sorted list of predicted objects:

js
[
  {
    box: {
      left: 105.1,
      top: 22.2,
      width: 70.6,
      height: 55.7
    },
    label: "Tomato",
    score: 0.972
  },
  ...
]

Advanced usage

Advanced users can access the underlying
GraphModel via
model.graphModel. The GraphModel allows users to call lower level methods
such as predict(), execute() and executeAsync() which return tensors.

model.dictionary gives you access to the ordered list of labels.


File: tfjs-backend-cpu/README.md

Usage

This package implements a JavaScript based CPU backend to TensorFlow.js.

Importing the backend

Note: this backend is included by default in @tensorflow/tfjs.

Via NPM

js
// Import @tensorflow/tfjs-core
import * as tf from '@tensorflow/tfjs-core';
// Adds the CPU backend to the global backend registry.
import '@tensorflow/tfjs-backend-cpu';

Via a script tag

html
<script src="https://cdn.jsdelivr.net/npm/@tensorflow/tfjs-core"></script>


<script src="https://cdn.jsdelivr.net/npm/@tensorflow/tfjs-backend-cpu"></script>

You can also get ES2017 code using the following links

html
<script src="https://cdn.jsdelivr.net/npm/@tensorflow/[email protected]/dist/tf-core.es2017.js"></script>


<script src="https://cdn.jsdelivr.net/npm/@tensorflow/[email protected]/dist/tf-backend-cpu.es2017.js"></script>

File: tfjs-backend-nodegl/demo/README.md

MobileNet tfjs-backend-nodegl Demo

This is a very early demo to show how tfjs-backend-nodegl can be used for headless WebGL acceleration.

To run this demo, perform the following:

  1. Move into tfjs-backend-nodegl (parent directory of this demo folder):
sh
$ cd tfjs-backend-nodegl
  1. Build package and compile TypeScript:
sh
$ yarn && yarn tsc
  1. Move into the demo directory:
sh
$ cd demo
  1. Prep and build demo:
sh
$ yarn
  1. Run demo:
sh
$ node run_mobilenet_inference.js dog.jpg

Expected output:

sh
$ node run_mobilenet_inference.js dog.jpg
Platform node has already been set. Overwriting the platform with [object Object].
  - gl.VERSION: OpenGL ES 3.0 (ANGLE 2.1.0.9512a0ef062a)
  - gl.RENDERER: ANGLE (Intel Inc., Intel(R) Iris(TM) Plus Graphics 640, OpenGL 4.1 core)
  - Loading model...
  - Mobilenet load: 6450.763924002647ms
  - Coldstarting model...
  - Mobilenet cold start: 297.92842200398445ms
  - Running inference (100x) ...
  - Mobilenet inference: (100x) : 35.75772546708584ms

File: tfjs-backend-nodegl/README.md

Headless WebGL backend for TensorFlow.js via Node.js

** This project is under heavy development **

This new backend will provide a light-weight headless WebGL runtime for TensorFlow.js running under Node.js. This new backend is powered by the node-gles module which uses ANGLE to provide an integration layer to system GL runtime. This package aims to provide a think acceleration engine for IoT, desktop, and Node.js applications where CUDA (size/OS compatibility) is not an option.

2. Official Technical Reference & Guides (tensorflow/docs)

File: README.md

TensorFlow Documentation

These are the source files for the guide and tutorials on
tensorflow.org.

To contribute to the TensorFlow documentation, please read
CONTRIBUTING.md, the
TensorFlow docs contributor guide,
and the style guide.

To file a docs issue, use the issue tracker in the
tensorflow/tensorflow repo.

And join the TensorFlow documentation contributors on the
TensorFlow Forum.

Community translations

Community translations
are located in the
tensorflow/docs-l10n repo. These docs
are contributed, reviewed, and maintained by the community as best-effort. To
participate as a translator or reviewer, see the site/<lang>/README.md, join
the language mailing list, and submit a pull request.

License

Apache License 2.0


File: site/en/r1/guide/extend/architecture.md

TensorFlow Architecture

We designed TensorFlow for large-scale distributed training and inference, but
it is also flexible enough to support experimentation with new machine
learning models and system-level optimizations.

This document describes the system architecture that makes this
combination of scale and flexibility possible. It assumes that you have basic familiarity
with TensorFlow programming concepts such as the computation graph, operations,
and sessions. See this document for an introduction to
these topics. Some familiarity with distributed TensorFlow
will also be helpful.

This document is for developers who want to extend TensorFlow in some way not
supported by current APIs, hardware engineers who want to optimize for
TensorFlow, implementers of machine learning systems working on scaling and
distribution, or anyone who wants to look under Tensorflow's hood. By the end of this document
you should understand the TensorFlow architecture well enough to read
and modify the core TensorFlow code.

Overview

The TensorFlow runtime is a cross-platform library. Figure 1 illustrates its
general architecture. A C API separates user level code in different languages
from the core runtime.

Figure 1

This document focuses on the following layers:

  • Client:
    • Defines the computation as a dataflow graph.
    • Initiates graph execution using a session.
  • Distributed Master
    • Prunes a specific subgraph from the graph, as defined by the arguments
      to Session.run().
    • Partitions the subgraph into multiple pieces that run in different
      processes and devices.
    • Distributes the graph pieces to worker services.
    • Initiates graph piece execution by worker services.
  • Worker Services (one for each task)
    • Schedule the execution of graph operations using kernel implementations
      appropriate to the available hardware (CPUs, GPUs, etc).
    • Send and receive operation results to and from other worker services.
  • Kernel Implementations
    • Perform the computation for individual graph operations.

Figure 2 illustrates the interaction of these components. "/job:worker/task:0" and
"/job:ps/task:0" are both tasks with worker services. "PS" stands for "parameter
server": a task responsible for storing and updating the model's parameters.
Other tasks send updates to these parameters as they work on optimizing the
parameters. This particular division of labor between tasks is not required, but
is common for distributed training.

Figure 2

Note that the Distributed Master and Worker Service only exist in
distributed TensorFlow. The single-process version of TensorFlow includes a
special Session implementation that does everything the distributed master does
but only communicates with devices in the local process.

The following sections describe the core TensorFlow layers in greater detail and
step through the processing of an example graph.

Client

Users write the client TensorFlow program that builds the computation graph.
This program can either directly compose individual operations or use a
convenience library like the Estimators API to compose neural network layers and
other higher-level abstractions. TensorFlow supports multiple client
languages, and we have prioritized Python and C++, because our internal users
are most familiar with these languages. As features become more established,
we typically port them to C++, so that users can access an optimized
implementation from all client languages. Most of the training libraries are
still Python-only, but C++ does have support for efficient inference.

The client creates a session, which sends the graph definition to the
distributed master as a tf.GraphDef
protocol buffer. When the client evaluates a node or nodes in the
graph, the evaluation triggers a call to the distributed master to initiate
computation.

In Figure 3, the client has built a graph that applies weights (w) to a
feature vector (x), adds a bias term (b) and saves the result in a variable
(s).

Figure 3

Code

  • tf.Session

Distributed master

The distributed master:

  • prunes the graph to obtain the subgraph required to evaluate the nodes
    requested by the client,
  • partitions the graph to obtain graph pieces for
    each participating device, and
  • caches these pieces so that they may be re-used in subsequent steps.

Since the master sees the overall computation for
a step, it applies standard optimizations such as common subexpression
elimination and constant folding. It then coordinates execution of the
optimized subgraphs across a set of tasks.

Figure 4

Figure 5 shows a possible partition of our example graph. The distributed
master has grouped the model parameters in order to place them together on the
parameter server.

Figure 5

Where graph edges are cut by the partition, the distributed master inserts
send and receive nodes to pass information between the distributed tasks
(Figure 6).

Figure 6

The distributed master then ships the graph pieces to the distributed tasks.

Figure 7

Code

Worker Service

The worker service in each task:

  • handles requests from the master,
  • schedules the execution of the kernels for the operations that comprise a
    local subgraph, and
  • mediates direct communication between tasks.

We optimize the worker service for running large graphs with low overhead. Our
current implementation can execute tens of thousands of subgraphs per second,
which enables a large number of replicas to make rapid, fine-grained training
steps. The worker service dispatches kernels to local devices and runs kernels
in parallel when possible, for example by using multiple CPU cores or GPU
streams.

We specialize Send and Recv operations for each pair of source and destination
device types:

  • Transfers between local CPU and GPU devices use the
    cudaMemcpyAsync() API to overlap computation and data transfer.
  • Transfers between two local GPUs use peer-to-peer DMA, to avoid an expensive
    copy via the host CPU.

For transfers between tasks, TensorFlow uses multiple protocols, including:

  • gRPC over TCP.
  • RDMA over Converged Ethernet.

We also have preliminary support for NVIDIA's NCCL library for multi-GPU
communication, see:
tf.contrib.nccl.

Figure 8

Code

Kernel Implementations

The runtime contains over 200 standard operations including mathematical, array
manipulation, control flow, and state management operations. Each of these
operations can have kernel implementations optimized for a variety of devices.
Many of the operation kernels are implemented using Eigen::Tensor, which uses
C++ templates to generate efficient parallel code for multicore CPUs and GPUs;
however, we liberally use libraries like cuDNN where a more efficient kernel
implementation is possible. We have also implemented
quantization, which enables
faster inference in environments such as mobile devices and high-throughput
datacenter applications, and use the
gemmlowp low-precision matrix library to
accelerate quantized computation.

If it is difficult or inefficient to represent a subcomputation as a composition
of operations, users can register additional kernels that provide an efficient
implementation written in C++. For example, we recommend registering your own
fused kernels for some performance critical operations, such as the ReLU and
Sigmoid activation functions and their corresponding gradients. The XLA Compiler has an
experimental implementation of automatic kernel fusion.

Code


File: site/en/r1/guide/extend/bindings.md

TensorFlow in other languages

Background

This document is intended as a guide for those interested in the creation or
development of TensorFlow functionality in other programming languages. It
describes the features of TensorFlow and recommended steps for making the same
available in other programming languages.

Python was the first client language supported by TensorFlow and currently
supports the most features. More and more of that functionality is being moved
into the core of TensorFlow (implemented in C++) and exposed via a C API.
Client languages should use the language's foreign function interface
(FFI)
to call into
this C API to provide TensorFlow functionality.

Overview

Providing TensorFlow functionality in a programming language can be broken down
into broad categories:

  • Run a predefined graph: Given a GraphDef (or
    MetaGraphDef) protocol message, be able to create a session, run queries,
    and get tensor results. This is sufficient for a mobile app or server that
    wants to run inference on a pre-trained model.
  • Graph construction: At least one function per defined
    TensorFlow op that adds an operation to the graph. Ideally these functions
    would be automatically generated so they stay in sync as the op definitions
    are modified.
  • Gradients (AKA automatic differentiation): Given a graph and a list of
    input and output operations, add operations to the graph that compute the
    partial derivatives (gradients) of the inputs with respect to the outputs.
    Allows for customization of the gradient function for a particular operation
    in the graph.
  • Functions: Define a subgraph that may be called in multiple places in the
    main GraphDef. Defines a FunctionDef in the FunctionDefLibrary
    included in a GraphDef.
  • Control Flow: Construct "If" and "While" with user-specified subgraphs.
    Ideally these work with gradients (see above).
  • Neural Network library: A number of components that together support the
    creation of neural network models and training them (possibly in a
    distributed setting). While it would be convenient to have this available in
    other languages, there are currently no plans to support this in languages
    other than Python. These libraries are typically wrappers over the features
    described above.

At a minimum, a language binding should support running a predefined graph, but
most should also support graph construction. The TensorFlow Python API provides
all these features.

Current Status

New language support should be built on top of the C API. However, as you can
see in the table below, not all functionality is available in C yet. Providing
more functionality in the C API is an ongoing project.

Feature Python C
Run a predefined Graph tf.import_graph_def, tf.Session TF_GraphImportGraphDef, TF_NewSession
Graph construction with generated op functions Yes Yes (The C API supports client languages that do this)
Gradients tf.gradients
Functions tf.python.framework.function.Defun
Control Flow tf.cond, tf.while_loop
Neural Network library tf.train, tf.nn, tf.contrib.layers, tf.contrib.slim

Run a predefined graph

A language binding is expected to define the following classes:

  • Graph: A graph representing a TensorFlow computation. Consists of
    operations (represented in the client language by Operations) and
    corresponds to a TF_Graph in the C API. Mainly used as an argument when
    creating new Operation objects and when starting a Session. Also
    supports iterating through the operations in the graph
    (TF_GraphNextOperation), looking up operations by name
    (TF_GraphOperationByName), and converting to and from a GraphDef
    protocol message (TF_GraphToGraphDef and TF_GraphImportGraphDef in the C
    API).
  • Operation: Represents a computation node in the graph. Corresponds to a
    TF_Operation in the C API.
  • Output: Represents one of the outputs of an operation in the graph. Has a
    DataType (and eventually a shape). May be passed as an input argument to a
    function for adding operations to a graph, or to a Session's Run()
    method to fetch that output as a tensor. Corresponds to a TF_Output in the
    C API.
  • Session: Represents a client to a particular instance of the TensorFlow
    runtime. Its main job is to be constructed with a Graph and some options
    and then field calls to Run() the graph. Corresponds to a TF_Session in
    the C API.
  • Tensor: Represents an N-dimensional (rectangular) array with elements all
    the same DataType. Gets data into and out of a Session's Run() call.
    Corresponds to a TF_Tensor in the C API.
  • DataType: An enumerant with all the possible tensor types supported by
    TensorFlow. Corresponds to TF_DataType in the C API and often referred to
    as dtype in the Python API.

Graph construction

TensorFlow has many ops, and the list is not static, so we recommend generating
the functions for adding ops to a graph instead of writing them by individually
by hand (though writing a few by hand is a good way to figure out what the
generator should generate). The information needed to generate a function is
contained in an OpDef protocol message.

There are a few ways to get a list of the OpDefs for the registered ops:

  • TF_GetAllOpList in the C API retrieves all registered OpDef protocol
    messages. This can be used to write the generator in the client language.
    This requires that the client language have protocol buffer support in order
    to interpret the OpDef messages.
  • The C++ function OpRegistry::Global()->GetRegisteredOps() returns the same
    list of all registered OpDefs (defined in
    tensorflow/core/framework/op.h). This can be used to write the generator
    in C++ (particularly useful for languages that do not have protocol buffer
    support).
  • The ASCII-serialized version of that list is periodically checked in to
    tensorflow/core/ops/ops.pbtxt by an automated process.

The OpDef specifies the following:

  • Name of the op in CamelCase. For generated functions follow the conventions
    of the language. For example, if the language uses snake_case, use that
    instead of CamelCase for the op's function name.
  • A list of inputs and outputs. The types for these may be polymorphic by
    referencing attributes, as described in the inputs and outputs section of
    Adding an op.
  • A list of attributes, along with their default values (if any). Note that
    some of these will be inferred (if they are determined by an input), some
    will be optional (if they have a default), and some will be required (no
    default).
  • Documentation for the op in general and the inputs, outputs, and
    non-inferred attributes.
  • Some other fields that are used by the runtime and can be ignored by the
    code generators.

An OpDef can be converted into the text of a function that adds that op to the
graph using the TF_OperationDescription C API (wrapped in the language's FFI):

  • Start with TF_NewOperation() to create the TF_OperationDescription*.
  • Call TF_AddInput() or TF_AddInputList() once per input (depending on
    whether the input has a list type).
  • Call TF_SetAttr*() functions to set non-inferred attributes. May skip
    attributes with defaults if you don't want to override the default value.
  • Set optional fields if necessary:
    • TF_SetDevice(): force the operation onto a specific device.
    • TF_AddControlInput(): add requirements that another operation finish
      before this operation starts running
    • TF_SetAttrString("_kernel") to set the kernel label (rarely used)
    • TF_ColocateWith() to colocate one op with another
  • Call TF_FinishOperation() when done. This adds the operation to the graph,
    after which it can't be modified.

The existing examples run the code generator as part of the build process (using
a Bazel genrule). Alternatively, the code generator can be run by an automated
cron process, possibly checking in the result. This creates a risk of divergence
between the generated code and the OpDefs checked into the repository, but is
useful for languages where code is expected to be generated ahead of time like
go get for Go and cargo ops for Rust. At the other end of the spectrum, for
some languages the code could be generated dynamically from
tensorflow/core/ops/ops.pbtxt.

Handling Constants

Calling code will be much more concise if users can provide constants to input
arguments. The generated code should convert those constants to operations that
are added to the graph and used as input to the op being instantiated.

Optional parameters

If the language allows for optional parameters to a function (like keyword
arguments with defaults in Python), use them for optional attributes, operation
names, devices, control inputs etc. In some languages, these optional parameters
can be set using dynamic scopes (like "with" blocks in Python). Without these
features, the library may resort to the "builder pattern", as is done in the C++
version of the TensorFlow API.

Name scopes

It is a good idea to have support for naming graph operations using some sort of
scoping hierarchy, especially considering the fact that TensorBoard relies on it
to display large graphs in a reasonable way. The existing Python and C++ APIs
take different approaches: In Python, the "directory" part of the name
(everything up to the last "/") comes from with blocks. In effect, there is a
thread-local stack with the scopes defining the name hierarchy. The last
component of the name is either supplied explicitly by the user (using the
optional name keyword argument) or defaults to the name of the type of the op
being added. In C++ the "directory" part of the name is stored in an explicit
Scope object. The NewSubScope() method appends to that part of the name and
returns a new Scope. The last component of the name is set using the
WithOpName() method, and like Python defaults to the name of the type of op
being added. Scope objects are explicitly passed around to specify the name of
the context.

Wrappers

It may make sense to keep the generated functions private for some ops so that
wrapper functions that do a little bit of additional work can be used instead.
This also gives an escape hatch for supporting features outside the scope of
generated code.

One use of a wrapper is for supporting SparseTensor input and output. A
SparseTensor is a tuple of 3 dense tensors: indices, values, and shape. values
is a vector size [n], shape is a vector size [rank], and indices is a matrix
size [n, rank]. There are some sparse ops that use this triple to represent a
single sparse tensor.

Another reason to use wrappers is for ops that hold state. There are a few such
ops (e.g. a variable) that have several companion ops for operating on that
state. The Python API has classes for these ops where the constructor creates
the op, and methods on that class add operations to the graph that operate on
the state.

Other Considerations

  • It is good to have a list of keywords used to rename op functions and
    arguments that collide with language keywords (or other symbols that will
    cause trouble, like the names of library functions or variables referenced
    in the generated code).
  • The function for adding a Const operation to a graph typically is a
    wrapper since the generated function will typically have redundant
    DataType inputs.

Gradients, functions and control flow

At this time, support for gradients, functions and control flow operations ("if"
and "while") is not available in languages other than Python. This will be
updated when the C API provides necessary support.


File: site/en/r1/guide/extend/cc.md

C++ API

Note: The instructions in this doc require
building from source. You probably want to build from
the master branch of TensorFlow.

TensorFlow's C++ API provides mechanisms for constructing and executing a data
flow graph. The API is designed to be simple and concise: graph operations are
clearly expressed using a "functional" construction style, including easy
specification of names, device placement, etc., and the resulting graph can be
efficiently run and the desired outputs fetched in a few lines of code. This
guide explains the basic concepts and data structures needed to get started with
TensorFlow graph construction and execution in C++.

The C++ API is only designed to work with TensorFlow bazel build.
If you need a stand-alone option, use the C API.
See these instructions
for details on how to include TensorFlow as a subproject (instead of building
your project from inside TensorFlow, as in this example).

The Basics

Let's start with a simple example that illustrates graph construction and
execution using the C++ API.

c++
// tensorflow/cc/example/example.cc

#include "tensorflow/cc/client/client_session.h"
#include "tensorflow/cc/ops/standard_ops.h"
#include "tensorflow/core/framework/tensor.h"

int main() {
  using namespace tensorflow;
  using namespace tensorflow::ops;
  Scope root = Scope::NewRootScope();
  // Matrix A = [3 2; -1 0]
  auto A = Const(root, { {3.f, 2.f}, {-1.f, 0.f} });
  // Vector b = [3 5]
  auto b = Const(root, { {3.f, 5.f} });
  // v = Ab^T
  auto v = MatMul(root.WithOpName("v"), A, b, MatMul::TransposeB(true));
  std::vector<Tensor> outputs;
  ClientSession session(root);
  // Run and fetch v
  TF_CHECK_OK(session.Run({v}, &outputs));
  // Expect outputs[0] == [19; -3]
  LOG(INFO) << outputs[0].matrix<float>();
  return 0;
}

Place this example code in the file tensorflow/cc/example/example.cc inside a
clone of the
TensorFlow
GitHub repository. Also place a
BUILD file in the same directory with the following contents:

python
load("//tensorflow:tensorflow.bzl", "tf_cc_binary")

tf_cc_binary(
    name = "example",
    srcs = ["example.cc"],
    deps = [
        "//tensorflow/cc:cc_ops",
        "//tensorflow/cc:client_session",
        "//tensorflow/core:tensorflow",
    ],
)

Use tf_cc_binary rather than Bazel's native cc_binary to link in necessary
symbols from libtensorflow_framework.so. You should be able to build and run
the example using the following command (be sure to run ./configure in your
build sandbox first):

shell
bazel run -c opt //tensorflow/cc/example:example

This example shows some of the important features of the C++ API such as the
following:

  • Constructing tensor constants from C++ nested initializer lists
  • Constructing and naming of TensorFlow operations
  • Specifying optional attributes to operation constructors
  • Executing and fetching the tensor values from the TensorFlow session.

We will delve into the details of each below.

Graph Construction

Scope

tensorflow::Scope is the main data structure that holds the current state
of graph construction. A Scope acts as a handle to the graph being
constructed, as well as storing TensorFlow operation properties. The Scope
object is the first argument to operation constructors, and operations that use
a given Scope as their first argument inherit that Scope's properties, such
as a common name prefix. Multiple Scopes can refer to the same graph, as
explained further below.

Create a new Scope object by calling Scope::NewRootScope. This creates
some resources such as a graph to which operations are added. It also creates a
tensorflow::Status object which will be used to indicate errors encountered
when constructing operations. The Scope class has value semantics, thus, a
Scope object can be freely copied and passed around.

The Scope object returned by Scope::NewRootScope is referred
to as the root scope. "Child" scopes can be constructed from the root scope by
calling various member functions of the Scope class, thus forming a hierarchy
of scopes. A child scope inherits all of the properties of the parent scope and
typically has one property added or changed. For instance, NewSubScope(name)
appends name to the prefix of names for operations created using the returned
Scope object.

Here are some of the properties controlled by a Scope object:

  • Operation names
  • Set of control dependencies for an operation
  • Device placement for an operation
  • Kernel attribute for an operation

Please refer to tensorflow::Scope for the complete list of member functions
that let you create child scopes with new properties.

Operation Constructors

You can create graph operations with operation constructors, one C++ class per
TensorFlow operation. Unlike the Python API which uses snake-case to name the
operation constructors, the C++ API uses camel-case to conform to C++ coding
style. For instance, the MatMul operation has a C++ class with the same name.

Using this class-per-operation method, it is possible, though not recommended,
to construct an operation as follows:

c++
// Not recommended
MatMul m(scope, a, b);

Instead, we recommend the following "functional" style for constructing
operations:

c++
// Recommended
auto m = MatMul(scope, a, b);

The first parameter for all operation constructors is always a Scope object.
Tensor inputs and mandatory attributes form the rest of the arguments.

For optional arguments, constructors have an optional parameter that allows
optional attributes. For operations with optional arguments, the constructor's
last optional parameter is a struct type called [operation]:Attrs that
contains data members for each optional attribute. You can construct such
Attrs in multiple ways:

  • You can specify a single optional attribute by constructing an Attrs object
    using the static functions provided in the C++ class for the operation. For
    example:
c++
auto m = MatMul(scope, a, b, MatMul::TransposeA(true));
  • You can specify multiple optional attributes by chaining together functions
    available in the Attrs struct. For example:
c++
auto m = MatMul(scope, a, b, MatMul::TransposeA(true).TransposeB(true));

// Or, alternatively
auto m = MatMul(scope, a, b, MatMul::Attrs().TransposeA(true).TransposeB(true));

The arguments and return values of operations are handled in different ways
depending on their type:

  • For operations that return single tensors, the object returned by
    the operation object can be passed directly to other operation
    constructors. For example:
c++
auto m = MatMul(scope, x, W);
auto sum = Add(scope, m, bias);
  • For operations producing multiple outputs, the object returned by the
    operation constructor has a member for each of the outputs. The names of those
    members are identical to the names present in the OpDef for the
    operation. For example:
c++
auto u = Unique(scope, a);
// u.y has the unique values and u.idx has the unique indices
auto m = Add(scope, u.y, b);
  • Operations producing a list-typed output return an object that can
    be indexed using the [] operator. That object can also be directly passed to
    other constructors that expect list-typed inputs. For example:
c++
auto s = Split(scope, 0, a, 2);
// Access elements of the returned list.
auto b = Add(scope, s[0], s[1]);
// Pass the list as a whole to other constructors.
auto c = Concat(scope, s, 0);

Constants

You may pass many different types of C++ values directly to tensor
constants. You may explicitly create a tensor constant by calling the
tensorflow::ops::Const function from various kinds of C++ values. For
example:

  • Scalars
c++
auto f = Const(scope, 42.0f);
auto s = Const(scope, "hello world!");
  • Nested initializer lists
c++
// 2x2 matrix
auto c1 = Const(scope, { {1, 2}, {2, 4} });
// 1x3x1 tensor
auto c2 = Const(scope, { { {1}, {2}, {3} } });
// 1x2x0 tensor
auto c3 = ops::Const(scope, { { {}, {} } });
  • Shapes explicitly specified
c++
// 2x2 matrix with all elements = 10
auto c1 = Const(scope, 10, /* shape */ {2, 2});
// 1x3x2x1 tensor
auto c2 = Const(scope, {1, 2, 3, 4, 5, 6}, /* shape */ {1, 3, 2, 1});

You may directly pass constants to other operation constructors, either by
explicitly constructing one using the Const function, or implicitly as any of
the above types of C++ values. For example:

c++
// [1 1] * [41; 1]
auto x = MatMul(scope, { {1, 1} }, { {41}, {1} });
// [1 2 3 4] + 10
auto y = Add(scope, {1, 2, 3, 4}, 10);

Graph Execution

When executing a graph, you will need a session. The C++ API provides a
tensorflow::ClientSession class that will execute ops created by the
operation constructors. TensorFlow will automatically determine which parts of
the graph need to be executed, and what values need feeding. For example:

c++
Scope root = Scope::NewRootScope();
auto c = Const(root, { {1, 1} });
auto m = MatMul(root, c, { {41}, {1} });

ClientSession session(root);
std::vector<Tensor> outputs;
session.Run({m}, &outputs);
// outputs[0] == {42}

Similarly, the object returned by the operation constructor can be used as the
argument to specify a value being fed when executing the graph. Furthermore, the
value to feed can be specified with the different kinds of C++ values used to
specify tensor constants. For example:

c++
Scope root = Scope::NewRootScope();
auto a = Placeholder(root, DT_INT32);
// [3 3; 3 3]
auto b = Const(root, 3, {2, 2});
auto c = Add(root, a, b);
ClientSession session(root);
std::vector<Tensor> outputs;

// Feed a <- [1 2; 3 4]
session.Run({ {a, { {1, 2}, {3, 4} } } }, {c}, &outputs);
// outputs[0] == [4 5; 6 7]

Please see the tensorflow::Tensor documentation for more information on how
to use the execution output.


File: site/en/r1/guide/extend/filesystem.md

Adding a Custom Filesystem Plugin

Background

The TensorFlow framework is often used in multi-process and
multi-machine environments, such as Google data centers, Google Cloud
Machine Learning, Amazon Web Services (AWS), and on-site distributed clusters.
In order to both share and save certain types of state produced by TensorFlow,
the framework assumes the existence of a reliable, shared filesystem. This
shared filesystem has numerous uses, for example:

  • Checkpoints of state are often saved to a distributed filesystem for
    reliability and fault-tolerance.
  • Training processes communicate with TensorBoard by writing event files
    to a directory, which TensorBoard watches. A shared filesystem allows this
    communication to work even when TensorBoard runs in a different process or
    machine.

There are many different implementations of shared or distributed filesystems in
the real world, so TensorFlow provides an ability for users to implement a
custom FileSystem plugin that can be registered with the TensorFlow runtime.
When the TensorFlow runtime attempts to write to a file through the FileSystem
interface, it uses a portion of the pathname to dynamically select the
implementation that should be used for filesystem operations. Thus, adding
support for your custom filesystem requires implementing a FileSystem
interface, building a shared object containing that implementation, and loading
that object at runtime in whichever process needs to write to that filesystem.

Note that TensorFlow already includes many filesystem implementations, such as:

  • A standard POSIX filesystem

    Note: NFS filesystems often mount as a POSIX interface, and so standard
    TensorFlow can work on top of NFS-mounted remote filesystems.

  • HDFS - the Hadoop File System

  • GCS - Google Cloud Storage filesystem

  • S3 - Amazon Simple Storage Service filesystem

  • A "memory-mapped-file" filesystem

The rest of this guide describes how to implement a custom filesystem.

Implementing a custom filesystem plugin

To implement a custom filesystem plugin, you must do the following:

  • Implement subclasses of RandomAccessFile, WriteableFile,
    AppendableFile, and ReadOnlyMemoryRegion.
  • Implement the FileSystem interface as a subclass.
  • Register the FileSystem implementation with an appropriate prefix pattern.
  • Load the filesystem plugin in a process that wants to write to that
    filesystem.

The FileSystem interface

The FileSystem interface is an abstract C++ interface defined in
file_system.h.
An implementation of the FileSystem interface should implement all relevant
the methods defined by the interface. Implementing the interface requires
defining operations such as creating RandomAccessFile, WritableFile, and
implementing standard filesystem operations such as FileExists, IsDirectory,
GetMatchingPaths, DeleteFile, and so on. An implementation of these
interfaces will often involve translating the function's input arguments to
delegate to an already-existing library function implementing the equivalent
functionality in your custom filesystem.

For example, the PosixFileSystem implementation implements DeleteFile using
the POSIX unlink() function; CreateDir simply calls mkdir(); GetFileSize
involves calling stat() on the file and then returns the filesize as reported
by the return of the stat object. Similarly, for the HDFSFileSystem
implementation, these calls simply delegate to the libHDFS implementation of
similar functionality, such as hdfsDelete for
DeleteFile.

We suggest looking through these code examples to get an idea of how different
filesystem implementations call their existing libraries. Examples include:

The File interfaces

Beyond operations that allow you to query and manipulate files and directories
in a filesystem, the FileSystem interface requires you to implement factories
that return implementations of abstract objects such as the
RandomAccessFile,
the WritableFile, so that TensorFlow code and read and write to files in that
FileSystem implementation.

To implement a RandomAccessFile, you must implement a single interface called
Read(), in which the implementation must provide a way to read from an offset
within a named file.

For example, below is the implementation of RandomAccessFile for the POSIX
filesystem, which uses the pread() random-access POSIX function to implement
read. Notice that the particular implementation must know how to retry or
propagate errors from the underlying filesystem.

c++
class PosixRandomAccessFile : public RandomAccessFile {
     public:
      PosixRandomAccessFile(const string& fname, int fd)
          : filename_(fname), fd_(fd) {}
      ~PosixRandomAccessFile() override { close(fd_); }

      Status Read(uint64 offset, size_t n, StringPiece* result,
                  char* scratch) const override {
        Status s;
        char* dst = scratch;
        while (n > 0 && s.ok()) {
          ssize_t r = pread(fd_, dst, n, static_cast<off_t>(offset));
          if (r > 0) {
            dst += r;
            n -= r;
            offset += r;
          } else if (r == 0) {
            s = Status(error::OUT_OF_RANGE, "Read less bytes than requested");
          } else if (errno == EINTR || errno == EAGAIN) {
            // Retry
          } else {
            s = IOError(filename_, errno);
          }
        }
        *result = StringPiece(scratch, dst - scratch);
        return s;
      }

     private:
      string filename_;
      int fd_;
    };

To implement the WritableFile sequential-writing abstraction, one must implement
a few interfaces, such as Append(), Flush(), Sync(), and Close().

For example, below is the implementation of WritableFile for the POSIX
filesystem, which takes a FILE object in its constructor and uses standard
posix functions on that object to implement the interface.

c++
class PosixWritableFile : public WritableFile {
     public:
      PosixWritableFile(const string& fname, FILE* f)
          : filename_(fname), file_(f) {}

      ~PosixWritableFile() override {
        if (file_ != NULL) {
          fclose(file_);
        }
      }

      Status Append(const StringPiece& data) override {
        size_t r = fwrite(data.data(), 1, data.size(), file_);
        if (r != data.size()) {
          return IOError(filename_, errno);
        }
        return Status::OK();
      }

      Status Close() override {
        Status result;
        if (fclose(file_) != 0) {
          result = IOError(filename_, errno);
        }
        file_ = NULL;
        return result;
      }

      Status Flush() override {
        if (fflush(file_) != 0) {
          return IOError(filename_, errno);
        }
        return Status::OK();
      }

      Status Sync() override {
        Status s;
        if (fflush(file_) != 0) {
          s = IOError(filename_, errno);
        }
        return s;
      }

     private:
      string filename_;
      FILE* file_;
    };

For more details, please see the documentations of those interfaces, and look at
example implementations for inspiration.

Registering and loading the filesystem

Once you have implemented the FileSystem implementation for your custom
filesystem, you need to register it under a "scheme" so that paths prefixed with
that scheme are directed to your implementation. To do this, you call
REGISTER_FILE_SYSTEM::

text
REGISTER_FILE_SYSTEM("foobar", FooBarFileSystem);

When TensorFlow tries to operate on a file whose path starts with foobar://,
it will use the FooBarFileSystem implementation.

c++
string filename = "foobar://path/to/file.txt";
    std::unique_ptr<WritableFile> file;

    // Calls FooBarFileSystem::NewWritableFile to return
    // a WritableFile class, which happens to be the FooBarFileSystem's
    // WritableFile implementation.
    TF_RETURN_IF_ERROR(env->NewWritableFile(filename, &file));

Next, you must build a shared object containing this implementation. An example
of doing so using bazel's cc_binary rule can be found
here,
but you may use any build system to do so. See the section on building the op library for similar
instructions.

The result of building this target is a .so shared object file.

Lastly, you must dynamically load this implementation in the process. In Python,
you can call the tf.load_file_system_library(file_system_library) function,
passing the path to the shared object. Calling this in your client program loads
the shared object in the process, thus registering your implementation as
available for any file operations going through the FileSystem interface. You
can see
test_file_system.py
for an example.

What goes through this interface?

Almost all core C++ file operations within TensorFlow use the FileSystem
interface, such as the CheckpointWriter, the EventsWriter, and many other
utilities. This means implementing a FileSystem implementation allows most of
your TensorFlow programs to write to your shared filesystem.

In Python, the gfile and file_io classes bind underneath to the `FileSystem
implementation via SWIG, which means that once you have loaded this filesystem
library, you can do:

text
with gfile.Open("foobar://path/to/file.txt") as w:

  w.write("hi")

When you do this, a file containing "hi" will appear in the "/path/to/file.txt"
of your shared filesystem.


File: site/en/r1/guide/extend/formats.md

Reading custom file and record formats

PREREQUISITES:

We divide the task of supporting a file format into two pieces:

  • File formats: We use a reader tf.data.Dataset to read raw records (which
    are typically represented by scalar string tensors, but can have more
    structure) from a file.
  • Record formats: We use decoder or parsing ops to turn a string record
    into tensors usable by TensorFlow.

For example, to re-implement tf.contrib.data.make_csv_dataset function, we
could use tf.data.TextLineDataset to extract the records, and then
use tf.data.Dataset.map and tf.decode_csv to parses the CSV records from
each line of text in the dataset.

[TOC]

Writing a `Dataset` for a file format

A tf.data.Dataset represents a sequence of elements, which can be the
individual records in a file. There are several examples of "reader" datasets
that are already built into TensorFlow:

Each of these implementations comprises three related classes:

  • A tensorflow::DatasetOpKernel subclass (e.g. TextLineDatasetOp), which
    tells TensorFlow how to construct a dataset object from the inputs to and
    attrs of an op, in its MakeDataset() method.

  • A tensorflow::GraphDatasetBase subclass (e.g. TextLineDatasetOp::Dataset),
    which represents the immutable definition of the dataset itself, and tells
    TensorFlow how to construct an iterator object over that dataset, in its
    MakeIteratorInternal() method.

  • A tensorflow::DatasetIterator<Dataset> subclass (e.g.
    TextLineDatasetOp::Dataset::Iterator), which represents the mutable state
    of an iterator over a particular dataset, and tells TensorFlow how to get the
    next element from the iterator, in its GetNextInternal() method.

The most important method is the GetNextInternal() method, since it defines
how to actually read records from the file and represent them as one or more
Tensor objects.

To create a new reader dataset called (for example) MyReaderDataset, you will
need to:

  1. In C++, define subclasses of tensorflow::DatasetOpKernel,
    tensorflow::GraphDatasetBase, and tensorflow::DatasetIterator<Dataset>
    that implement the reading logic.
  2. In C++, register a new reader op and kernel with the name
    "MyReaderDataset".
  3. In Python, define a subclass of tf.data.Dataset called MyReaderDataset.

You can put all the C++ code in a single file, such as
my_reader_dataset_op.cc. It will help if you are
familiar with the adding an op how-to. The following skeleton
can be used as a starting point for your implementation:

text
/* Detailed source-code truncated for AI context efficiency. */

The last step is to build the C++ code and add a Python wrapper. The easiest way
to do this is by compiling a dynamic
library
(e.g. called "my_reader_dataset_op.so"), and adding a Python class
that subclasses tf.data.Dataset to wrap it. An example Python program is
given here:

python
import tensorflow as tf

# Assumes the file is in the current working directory.
my_reader_dataset_module = tf.load_op_library("./my_reader_dataset_op.so")

class MyReaderDataset(tf.data.Dataset):

  def __init__(self):
    super(MyReaderDataset, self).__init__()
    # Create any input attrs or tensors as members of this class.

  def _as_variant_tensor(self):
    # Actually construct the graph node for the dataset op.
    #
    # This method will be invoked when you create an iterator on this dataset
    # or a dataset derived from it.
    return my_reader_dataset_module.my_reader_dataset()

  # The following properties define the structure of each element: a scalar
  # `tf.string` tensor. Change these properties to match the `output_dtypes()`
  # and `output_shapes()` methods of `MyReaderDataset::Dataset` if you modify
  # the structure of each element.
  @property
  def output_types(self):
    return tf.string

  @property
  def output_shapes(self):
    return tf.TensorShape([])

  @property
  def output_classes(self):
    return tf.Tensor

if __name__ == "__main__":
  # Create a MyReaderDataset and print its elements.
  with tf.Session() as sess:
    iterator = MyReaderDataset().make_one_shot_iterator()
    next_element = iterator.get_next()
    try:
      while True:
        print(sess.run(next_element))  # Prints "MyReader!" ten times.
    except tf.errors.OutOfRangeError:
      pass

You can see some examples of Dataset wrapper classes in
tensorflow/python/data/ops/dataset_ops.py.

Writing an Op for a record format

Generally this is an ordinary op that takes a scalar string record as input, and
so follow the instructions to add an Op.
You may optionally take a scalar string key as input, and include that in error
messages reporting improperly formatted data. That way users can more easily
track down where the bad data came from.

Examples of Ops useful for decoding records:

  • tf.parse_single_example (and tf.parse_example)
  • tf.decode_csv
  • tf.decode_raw

Note that it can be useful to use multiple Ops to decode a particular record
format. For example, you may have an image saved as a string in
a tf.train.Example protocol buffer.
Depending on the format of that image, you might take the corresponding output
from a tf.parse_single_example op and call tf.image.decode_jpeg,
tf.image.decode_png, or tf.decode_raw. It is common to take the output
of tf.decode_raw and use tf.slice and tf.reshape to extract pieces.


File: site/en/r1/guide/extend/model_files.md

A Tool Developer's Guide to TensorFlow Model Files

Most users shouldn't need to care about the internal details of how TensorFlow
stores data on disk, but you might if you're a tool developer. For example, you
may want to analyze models, or convert back and forth between TensorFlow and
other formats. This guide tries to explain some of the details of how you can
work with the main files that hold model data, to make it easier to develop
those kind of tools.

[TOC]

Protocol Buffers

All of TensorFlow's file formats are based on
Protocol Buffers, so to
start it's worth getting familiar with how they work. The summary is that you
define data structures in text files, and the protobuf tools generate classes in
C, Python, and other languages that can load, save, and access the data in a
friendly way. We often refer to Protocol Buffers as protobufs, and I'll use
that convention in this guide.

GraphDef

The foundation of computation in TensorFlow is the Graph object. This holds a
network of nodes, each representing one operation, connected to each other as
inputs and outputs. After you've created a Graph object, you can save it out
by calling as_graph_def(), which returns a GraphDef object.

The GraphDef class is an object created by the ProtoBuf library from the
definition in
tensorflow/core/framework/graph.proto. The protobuf tools parse
this text file, and generate the code to load, store, and manipulate graph
definitions. If you see a standalone TensorFlow file representing a model, it's
likely to contain a serialized version of one of these GraphDef objects
saved out by the protobuf code.

This generated code is used to save and load the GraphDef files from disk. The code that actually loads the model looks like this:

python
graph_def = graph_pb2.GraphDef()

This line creates an empty GraphDef object, the class that's been created
from the textual definition in graph.proto. This is the object we're going to
populate with the data from our file.

python
with open(FLAGS.graph, "rb") as f:

Here we get a file handle for the path we've passed in to the script

python
if FLAGS.input_binary:
    graph_def.ParseFromString(f.read())
  else:
    text_format.Merge(f.read(), graph_def)

Text or Binary?

There are actually two different formats that a ProtoBuf can be saved in.
TextFormat is a human-readable form, which makes it nice for debugging and
editing, but can get large when there's numerical data like weights stored in
it. You can see a small example of that in
graph_run_run2.pbtxt.

Binary format files are a lot smaller than their text equivalents, even though
they're not as readable for us. In this script, we ask the user to supply a
flag indicating whether the input file is binary or text, so we know the right
function to call. You can find an example of a large binary file inside the
inception_v3 archive,
as inception_v3_2016_08_28_frozen.pb.

The API itself can be a bit confusing - the binary call is actually
ParseFromString(), whereas you use a utility function from the text_format
module to load textual files.

Nodes

Once you've loaded a file into the graph_def variable, you can now access the
data inside it. For most practical purposes, the important section is the list
of nodes stored in the node member. Here's the code that loops through those:

python
for node in graph_def.node

Each node is a NodeDef object, defined in
tensorflow/core/framework/node_def.proto. These
are the fundamental building blocks of TensorFlow graphs, with each one defining
a single operation along with its input connections. Here are the members of a
NodeDef, and what they mean.

`name`

Every node should have a unique identifier that's not used by any other nodes
in the graph. If you don't specify one as you're building a graph using the
Python API, one reflecting the name of operation, such as "MatMul",
concatenated with a monotonically increasing number, such as "5", will be
picked for you. The name is used when defining the connections between nodes,
and when setting inputs and outputs for the whole graph when it's run.

`op`

This defines what operation to run, for example "Add", "MatMul", or
"Conv2D". When a graph is run, this op name is looked up in a registry to
find an implementation. The registry is populated by calls to the
REGISTER_OP() macro, like those in
tensorflow/core/ops/nn_ops.cc.

`input`

A list of strings, each one of which is the name of another node, optionally
followed by a colon and an output port number. For example, a node with two
inputs might have a list like ["some_node_name", "another_node_name"], which
is equivalent to ["some_node_name:0", "another_node_name:0"], and defines the
node's first input as the first output from the node with the name
"some_node_name", and a second input from the first output of
"another_node_name"

`device`

In most cases you can ignore this, since it defines where to run a node in a
distributed environment, or when you want to force the operation onto CPU or
GPU.

`attr`

This is a key/value store holding all the attributes of a node. These are the
permanent properties of nodes, things that don't change at runtime such as the
size of filters for convolutions, or the values of constant ops. Because there
can be so many different types of attribute values, from strings, to ints, to
arrays of tensor values, there's a separate protobuf file defining the data
structure that holds them, in
tensorflow/core/framework/attr_value.proto.

Each attribute has a unique name string, and the expected attributes are listed
when the operation is defined. If an attribute isn't present in a node, but it
has a default listed in the operation definition, that default is used when the
graph is created.

You can access all of these members by calling node.name, node.op, etc. in
Python. The list of nodes stored in the GraphDef is a full definition of the
model architecture.

Freezing

One confusing part about this is that the weights usually aren't stored inside
the file format during training. Instead, they're held in separate checkpoint
files, and there are Variable ops in the graph that load the latest values
when they're initialized. It's often not very convenient to have separate files
when you're deploying to production, so there's the
freeze_graph.py script that takes a graph definition and a set
of checkpoints and freezes them together into a single file.

What this does is load the GraphDef, pull in the values for all the variables
from the latest checkpoint file, and then replace each Variable op with a
Const that has the numerical data for the weights stored in its attributes.
It then strips away all the extraneous nodes that aren't used for forward
inference, and saves out the resulting GraphDef into an output file.

Weight Formats

If you're dealing with TensorFlow models that represent neural networks, one of
the most common problems is extracting and interpreting the weight values. A
common way to store them, for example in graphs created by the freeze_graph
script, is as Const ops containing the weights as Tensors. These are
defined in
tensorflow/core/framework/tensor.proto, and contain information
about the size and type of the data, as well as the values themselves. In
Python, you get a TensorProto object from a NodeDef representing a Const
op by calling something like some_node_def.attr['value'].tensor.

This will give you an object representing the weights data. The data itself
will be stored in one of the lists with the suffix _val as indicated by the
type of the object, for example float_val for 32-bit float data types.

The ordering of convolution weight values is often tricky to deal with when
converting between different frameworks. In TensorFlow, the filter weights for
the Conv2D operation are stored on the second input, and are expected to be
in the order [filter_height, filter_width, input_depth, output_depth], where
filter_count increasing by one means moving to an adjacent value in memory.

Hopefully this rundown gives you a better idea of what's going on inside
TensorFlow model files, and will help you if you ever need to manipulate them.


File: site/en/r1/guide/extend/op.md

Adding a New Op

Note: To guarantee that your C++ custom ops are ABI compatible with TensorFlow's
official pip packages, please follow the guide at
Custom op repository. It has an
end-to-end code example, as well as Docker images for building and distributing
your custom ops.

If you'd like to create an op that isn't covered by the existing TensorFlow
library, we recommend that you first try writing the op in Python as
a composition of existing Python ops or functions. If that isn't possible, you
can create a custom C++ op. There are several reasons why you might want to
create a custom C++ op:

  • It's not easy or possible to express your operation as a composition of
    existing ops.
  • It's not efficient to express your operation as a composition of existing
    primitives.
  • You want to hand-fuse a composition of primitives that a future compiler
    would find difficult fusing.

For example, imagine you want to implement something like "median pooling",
similar to the "MaxPool" operator, but computing medians over sliding windows
instead of maximum values. Doing this using a composition of operations may be
possible (e.g., using ExtractImagePatches and TopK), but may not be as
performance- or memory-efficient as a native operation where you can do
something more clever in a single, fused operation. As always, it is typically
first worth trying to express what you want using operator composition, only
choosing to add a new operation if that proves to be difficult or inefficient.

To incorporate your custom op you'll need to:

  1. Register the new op in a C++ file. Op registration defines an interface
    (specification) for the op's functionality, which is independent of the
    op's implementation. For example, op registration defines the op's name and
    the op's inputs and outputs. It also defines the shape function
    that is used for tensor shape inference.
  2. Implement the op in C++. The implementation of an op is known
    as a kernel, and it is the concrete implementation of the specification you
    registered in Step 1. There can be multiple kernels for different input /
    output types or architectures (for example, CPUs, GPUs).
  3. Create a Python wrapper (optional). This wrapper is the public API that's
    used to create the op in Python. A default wrapper is generated from the
    op registration, which can be used directly or added to.
  4. Write a function to compute gradients for the op (optional).
  5. Test the op. We usually do this in Python for convenience, but you can also
    test the op in C++. If you define gradients, you can verify them with the
    Python tf.test.compute_gradient_error.
    See
    relu_op_test.py as
    an example that tests the forward functions of Relu-like operators and
    their gradients.

PREREQUISITES:

[TOC]

Define the op's interface

You define the interface of an op by registering it with the TensorFlow system.
In the registration, you specify the name of your op, its inputs (types and
names) and outputs (types and names), as well as docstrings and
any attrs the op might require.

To see how this works, suppose you'd like to create an op that takes a tensor of
int32s and outputs a copy of the tensor, with all but the first element set to
zero. To do this, create a file named zero_out.cc. Then add a call to the
REGISTER_OP macro that defines the interface for your op:

c++
#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/shape_inference.h"

using namespace tensorflow;

REGISTER_OP("ZeroOut")
    .Input("to_zero: int32")
    .Output("zeroed: int32")
    .SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) {
      c->set_output(0, c->input(0));
      return Status::OK();
    });

This ZeroOut op takes one tensor to_zero of 32-bit integers as input, and
outputs a tensor zeroed of 32-bit integers. The op also uses a shape function
to ensure that the output tensor is the same shape as the input tensor. For
example, if the input is a tensor of shape [10, 20], then this shape function
specifies that the output shape is also [10, 20].

A note on naming: The op name must be in CamelCase and it must be unique
among all other ops that are registered in the binary.

Implement the kernel for the op

After you define the interface, provide one or more implementations of the op.
To create one of these kernels, create a class that extends OpKernel and
overrides the Compute method. The Compute method provides one context
argument of type OpKernelContext*, from which you can access useful things
like the input and output tensors.

Add your kernel to the file you created above. The kernel might look something
like this:

c++
#include "tensorflow/core/framework/op_kernel.h"

using namespace tensorflow;

class ZeroOutOp : public OpKernel {
 public:
  explicit ZeroOutOp(OpKernelConstruction* context) : OpKernel(context) {}

  void Compute(OpKernelContext* context) override {
    // Grab the input tensor
    const Tensor& input_tensor = context->input(0);
    auto input = input_tensor.flat();

    // Create an output tensor
    Tensor* output_tensor = NULL;
    OP_REQUIRES_OK(context, context->allocate_output(0, input_tensor.shape(),
                                                     &output_tensor));
    auto output_flat = output_tensor->flat();

    // Set all but the first element of the output tensor to 0.
    const int N = input.size();
    for (int i = 1; i < N; i++) {
      output_flat(i) = 0;
    }

    // Preserve the first input value if possible.
    if (N > 0) output_flat(0) = input(0);
  }
};

After implementing your kernel, you register it with the TensorFlow system. In
the registration, you specify different constraints under which this kernel
will run. For example, you might have one kernel made for CPUs, and a separate
one for GPUs.

To do this for the ZeroOut op, add the following to zero_out.cc:

c++
REGISTER_KERNEL_BUILDER(Name("ZeroOut").Device(DEVICE_CPU), ZeroOutOp);

Important: Instances of your OpKernel may be accessed concurrently.
Your Compute method must be thread-safe. Guard any access to class
members with a mutex. Or better yet, don't share state via class members!
Consider using a ResourceMgr
to keep track of op state.

Multi-threaded CPU kernels

To write a multi-threaded CPU kernel, the Shard function in
work_sharder.h
can be used. This function shards a computation function across the
threads configured to be used for intra-op threading (see
intra_op_parallelism_threads in
config.proto).

GPU kernels

A GPU kernel is implemented in two parts: the OpKernel and the CUDA kernel and
its launch code.

Sometimes the OpKernel implementation is common between a CPU and GPU kernel,
such as around inspecting inputs and allocating outputs. In that case, a
suggested implementation is to:

  1. Define the OpKernel templated on the Device and the primitive type of the
    tensor.
  2. To do the actual computation of the output, the Compute function calls a
    templated functor struct.
  3. The specialization of that functor for the CPUDevice is defined in the same
    file, but the specialization for the GPUDevice is defined in a .cu.cc file,
    since it will be compiled with the CUDA compiler.

Here is an example implementation.

c++
// kernel_example.h
#ifndef KERNEL_EXAMPLE_H_
#define KERNEL_EXAMPLE_H_

template <typename Device, typename T>
struct ExampleFunctor {
  void operator()(const Device& d, int size, const T* in, T* out);
};

#if GOOGLE_CUDA
// Partially specialize functor for GpuDevice.
template <typename Eigen::GpuDevice, typename T>
struct ExampleFunctor {
  void operator()(const Eigen::GpuDevice& d, int size, const T* in, T* out);
};
#endif

#endif KERNEL_EXAMPLE_H_
text
/* Detailed source-code truncated for AI context efficiency. */
c++
// kernel_example.cu.cc
#ifdef GOOGLE_CUDA
#define EIGEN_USE_GPU
#include "example.h"
#include "tensorflow/core/util/gpu_kernel_helper.h"

using namespace tensorflow;

using GPUDevice = Eigen::GpuDevice;

// Define the CUDA kernel.
template <typename T>
__global__ void ExampleCudaKernel(const int size, const T* in, T* out) {
  for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < size;
       i += blockDim.x * gridDim.x) {
    out[i] = 2 * ldg(in + i);
  }
}

// Define the GPU implementation that launches the CUDA kernel.
template <typename T>
void ExampleFunctor<GPUDevice, T>::operator()(
    const GPUDevice& d, int size, const T* in, T* out) {
  // Launch the cuda kernel.
  //
  // See core/util/gpu_kernel_helper.h for example of computing
  // block count and thread_per_block count.
  int block_count = 1024;
  int thread_per_block = 20;
  ExampleCudaKernel<T>
      <<>>(size, in, out);
}

// Explicitly instantiate functors for the types of OpKernels registered.
template struct ExampleFunctor<GPUDevice, float>;
template struct ExampleFunctor<GPUDevice, int32>;

#endif  // GOOGLE_CUDA

Build the op library

Compile the op using your system compiler (TensorFlow binary installation)

You should be able to compile zero_out.cc with a C++ compiler such as g++
or clang available on your system. The binary PIP package installs the header
files and the library that you need to compile your op in locations that are
system specific. However, the TensorFlow python library provides the
get_include function to get the header directory, and the get_lib directory
has a shared object to link against.
Here are the outputs of these functions on an Ubuntu machine.

bash
$ python
>>> import tensorflow as tf
>>> tf.sysconfig.get_include()
'/usr/local/lib/python3.6/site-packages/tensorflow/include'
>>> tf.sysconfig.get_lib()
'/usr/local/lib/python3.6/site-packages/tensorflow'

Assuming you have g++ installed, here is the sequence of commands you can use
to compile your op into a dynamic library.

bash
TF_CFLAGS=( $(python -c 'import tensorflow as tf; print(" ".join(tf.sysconfig.get_compile_flags()))') )
TF_LFLAGS=( $(python -c 'import tensorflow as tf; print(" ".join(tf.sysconfig.get_link_flags()))') )
g++ -std=c++11 -shared zero_out.cc -o zero_out.so -fPIC ${TF_CFLAGS[@]} ${TF_LFLAGS[@]} -O2

On macOS, the additional flag "-undefined dynamic_lookup" is required when
building the .so file.

Note on gcc version >=5: gcc uses the new C++
ABI since version 5.
TensorFlow 2.8 and earlier were built with gcc4 that uses the older ABI. If
you are using these versions of TensorFlow and are trying to compile your op
library with gcc>=5, add -D_GLIBCXX_USE_CXX11_ABI=0 to the command line to
make the library compatible with the older ABI. TensorFlow 2.9+ packages are
compatible with the newer ABI by default.

Compile the op using bazel (TensorFlow source installation)

If you have TensorFlow sources installed, you can make use of TensorFlow's build
system to compile your op. Place a BUILD file with following Bazel build rule in
the tensorflow/core/user_ops directory.

python
load("//tensorflow:tensorflow.bzl", "tf_custom_op_library")

tf_custom_op_library(
    name = "zero_out.so",
    srcs = ["zero_out.cc"],
)

Run the following command to build zero_out.so.

bash
$ bazel build --config opt //tensorflow/core/user_ops:zero_out.so

As explained above, if you are compiling with gcc>=5 add --cxxopt="-D_GLIBCXX_USE_CXX11_ABI=0"
to the bazel command line.

Note: Although you can create a shared library (a .so file) with the
standard cc_library rule, we strongly recommend that you use the
tf_custom_op_library macro. It adds some required dependencies, and
performs checks to ensure that the shared library is compatible with
TensorFlow's plugin loading mechanism.

Use the op in Python

TensorFlow Python API provides the
tf.load_op_library function to
load the dynamic library and register the op with the TensorFlow
framework. load_op_library returns a Python module that contains the Python
wrappers for the op and the kernel. Thus, once you have built the op, you can
do the following to run it from Python:

python
import tensorflow as tf
zero_out_module = tf.load_op_library('./zero_out.so')
with tf.Session(''):
  zero_out_module.zero_out([[1, 2], [3, 4]]).eval()

# Prints
array([[1, 0], [0, 0]], dtype=int32)

Keep in mind, the generated function will be given a snake_case name (to comply
with PEP8). So, if your op is
named ZeroOut in the C++ files, the python function will be called zero_out.

To make the op available as a regular function import-able from a Python
module, it maybe useful to have the load_op_library call in a Python source
file as follows:

python
import tensorflow as tf

zero_out_module = tf.load_op_library('./zero_out.so')
zero_out = zero_out_module.zero_out

Verify that the op works

A good way to verify that you've successfully implemented your op is to write a
test for it. Create the file
zero_out_op_test.py with the contents:

python
import tensorflow as tf

class ZeroOutTest(tf.test.TestCase):
  def testZeroOut(self):
    zero_out_module = tf.load_op_library('./zero_out.so')
    with self.test_session():
      result = zero_out_module.zero_out([5, 4, 3, 2, 1])
      self.assertAllEqual(result.eval(), [5, 0, 0, 0, 0])

if __name__ == "__main__":
  tf.test.main()

Then run your test (assuming you have tensorflow installed):

sh
$ python zero_out_op_test.py

Building advanced features into your op

Now that you know how to build a basic (and somewhat restricted) op and
implementation, we'll look at some of the more complicated things you will
typically need to build into your op. This includes:

Conditional checks and validation

The example above assumed that the op applied to a tensor of any shape. What
if it only applied to vectors? That means adding a check to the above OpKernel
implementation.

c++
void Compute(OpKernelContext* context) override {
    // Grab the input tensor
    const Tensor& input_tensor = context->input(0);

    OP_REQUIRES(context, TensorShapeUtils::IsVector(input_tensor.shape()),
                errors::InvalidArgument("ZeroOut expects a 1-D vector."));
    // ...
  }

This asserts that the input is a vector, and returns having set the
InvalidArgument status if it isn't. The
OP_REQUIRES macro takes three arguments:

Alternatively, if you want to test whether a Status object returned from some
function is an error, and if so return it, use
OP_REQUIRES_OK. Both of these macros return from the
function on error.

Op registration

Attrs

Ops can have attrs, whose values are set when the op is added to a graph. These
are used to configure the op, and their values can be accessed both within the
kernel implementation and in the types of inputs and outputs in the op
registration. Prefer using an input instead of an attr when possible, since
inputs are more flexible. This is because attrs are constants and must be
defined at graph construction time. In contrast, inputs are Tensors whose
values can be dynamic; that is, inputs can change every step, be set using a
feed, etc. Attrs are used for things that can't be done with inputs: any
configuration that affects the signature (number or type of inputs or outputs)
or that can't change from step-to-step.

You define an attr when you register the op, by specifying its name and type
using the Attr method, which expects a spec of the form:

text
<name>: <attr-type-expr>

where <name> begins with a letter and can be composed of alphanumeric
characters and underscores, and <attr-type-expr> is a type expression of the
form described below.

For example, if you'd like the ZeroOut op to preserve a user-specified index,
instead of only the 0th element, you can register the op like so:

c++
REGISTER_OP("ZeroOut")
    .Attr("preserve_index: int")
    .Input("to_zero: int32")
    .Output("zeroed: int32");

(Note that the set of attribute types is different from the
tf.DType used for inputs and outputs.)

Your kernel can then access this attr in its constructor via the context
parameter:

c++
class ZeroOutOp : public OpKernel {
 public:
  explicit ZeroOutOp(OpKernelConstruction* context) : OpKernel(context) {
    // Get the index of the value to preserve
    OP_REQUIRES_OK(context,
                   context->GetAttr("preserve_index", &preserve_index_));
    // Check that preserve_index is positive
    OP_REQUIRES(context, preserve_index_ >= 0,
                errors::InvalidArgument("Need preserve_index >= 0, got ",
                                        preserve_index_));
  }
  void Compute(OpKernelContext* context) override {
    // ...
  }
 private:
  int preserve_index_;
};

which can then be used in the Compute method:

c++
void Compute(OpKernelContext* context) override {
    // ...

    // We're using saved attr to validate potentially dynamic input
    // So we check that preserve_index is in range
    OP_REQUIRES(context, preserve_index_ < input.dimension(0),
                errors::InvalidArgument("preserve_index out of range"));

    // Set all the elements of the output tensor to 0
    const int N = input.size();
    for (int i = 0; i < N; i++) {
      output_flat(i) = 0;
    }

    // Preserve the requested input value
    output_flat(preserve_index_) = input(preserve_index_);
  }

Attr types

The following types are supported in an attr:

  • string: Any sequence of bytes (not required to be UTF8).
  • int: A signed integer.
  • float: A floating point number.
  • bool: True or false.
  • type: One of the (non-ref) values of DataType.
  • shape: A TensorShapeProto.
  • tensor: A TensorProto.
  • list(<type>): A list of <type>, where <type> is one of the above types.
    Note that list(list(<type>)) is invalid.

See also: op_def_builder.cc:FinalizeAttr for a definitive list.

Default values & constraints

Attrs may have default values, and some types of attrs can have constraints. To
define an attr with constraints, you can use the following <attr-type-expr>s:

  • {'<string1>', '<string2>'}: The value must be a string that has either the
    value <string1> or <string2>. The name of the type, string, is implied
    when you use this syntax. This emulates an enum:

    c++
    REGISTER_OP("EnumExample")
        .Attr("e: {'apple', 'orange'}");
  • {<type1>, <type2>}: The value is of type type, and must be one of
    <type1> or <type2>, where <type1> and <type2> are supported
    tf.DType. You don't specify
    that the type of the attr is type. This is implied when you have a list of
    types in {...}. For example, in this case the attr t is a type that must
    be an int32, a float, or a bool:

    c++
    REGISTER_OP("RestrictedTypeExample")
        .Attr("t: {int32, float, bool}");
  • There are shortcuts for common type constraints:

    • numbertype: Type type restricted to the numeric (non-string and
      non-bool) types.
    • realnumbertype: Like numbertype without complex types.
    • quantizedtype: Like numbertype but just the quantized number types.

    The specific lists of types allowed by these are defined by the functions
    (like NumberTypes()) in
    tensorflow/core/framework/types.h.
    In this example the attr t must be one of the numeric types:

    c++
    REGISTER_OP("NumberType")
        .Attr("t: numbertype");

    For this op:

    python
    tf.number_type(t=tf.int32)  # Valid
    tf.number_type(t=tf.bool)   # Invalid

    Lists can be combined with other lists and single types. The following
    op allows attr t to be any of the numeric types, or the bool type:

    c++
    REGISTER_OP("NumberOrBooleanType")
        .Attr("t: {numbertype, bool}");

    For this op:

    python
    tf.number_or_boolean_type(t=tf.int32)  # Valid
    tf.number_or_boolean_type(t=tf.bool)   # Valid
    tf.number_or_boolean_type(t=tf.string) # Invalid
  • int >= <n>: The value must be an int whose value is greater than or equal to
    <n>, where <n> is a natural number.

    For example, the following op registration specifies that the attr a must
    have a value that is at least 2:

    c++
    REGISTER_OP("MinIntExample")
        .Attr("a: int >= 2");
  • list(<type>) >= <n>: A list of type <type> whose length is greater than
    or equal to <n>.

    For example, the following op registration specifies that the attr a is a
    list of types (either int32 or float), and that there must be at least 3
    of them:

    c++
    REGISTER_OP("TypeListExample")
        .Attr("a: list({int32, float}) >= 3");

To set a default value for an attr (making it optional in the generated code),
add = <default> to the end, as in:

c++
REGISTER_OP("AttrDefaultExample")
    .Attr("i: int = 0");

The supported syntax of the default value is what would be used in the proto
representation of the resulting GraphDef definition.

Here are examples for how to specify a default for all types:

c++
REGISTER_OP("AttrDefaultExampleForAllTypes")
   .Attr("s: string = 'foo'")
   .Attr("i: int = 0")
   .Attr("f: float = 1.0")
   .Attr("b: bool = true")
   .Attr("ty: type = DT_INT32")
   .Attr("sh: shape = { dim { size: 1 } dim { size: 2 } }")
   .Attr("te: tensor = { dtype: DT_INT32 int_val: 5 }")
   .Attr("l_empty: list(int) = []")
   .Attr("l_int: list(int) = [2, 3, 5, 7]");

Note in particular that the values of type type
use tf.DType.

Polymorphism

Type Polymorphism

For ops that can take different types as input or produce different output
types, you can specify an attr in
an input or output type in the op registration. Typically
you would then register an OpKernel for each supported type.

For instance, if you'd like the ZeroOut op to work on floats
in addition to int32s, your op registration might look like:

c++
REGISTER_OP("ZeroOut")
    .Attr("T: {float, int32}")
    .Input("to_zero: T")
    .Output("zeroed: T");

Your op registration now specifies that the input's type must be float, or
int32, and that its output will be the same type, since both have type T.

A note on naming: Inputs, outputs, and attrs generally should be
given snake_case names. The one exception is attrs that are used as the type
of an input or in the type of an input. Those attrs can be inferred when the
op is added to the graph and so don't appear in the op's function. For
example, this last definition of ZeroOut will generate a Python function that
looks like:

python
def zero_out(to_zero, name=None):
  """...
  Args:
    to_zero: A `Tensor`. Must be one of the following types:
        `float32`, `int32`.
    name: A name for the operation (optional).

  Returns:
    A `Tensor`. Has the same type as `to_zero`.
  """

If to_zero is passed an int32 tensor, then T is automatically set to
int32 (well, actually DT_INT32). Those inferred attrs are given
Capitalized or CamelCase names.

Compare this with an op that has a type attr that determines the output
type:

c++
REGISTER_OP("StringToNumber")
    .Input("string_tensor: string")
    .Output("output: out_type")
    .Attr("out_type: {float, int32} = DT_FLOAT");
    .Doc(R"doc(
Converts each string in the input Tensor to the specified numeric type.
)doc");

In this case, the user has to specify the output type, as in the generated
Python:

python
def string_to_number(string_tensor, out_type=None, name=None):
  """Converts each string in the input Tensor to the specified numeric type.

  Args:
    string_tensor: A `Tensor` of type `string`.
    out_type: An optional `tf.DType` from: `tf.float32, tf.int32`.
      Defaults to `tf.float32`.
    name: A name for the operation (optional).

  Returns:
    A `Tensor` of type `out_type`.
  """
c++
#include "tensorflow/core/framework/op_kernel.h"

class ZeroOutInt32Op : public OpKernel {
  // as before
};

class ZeroOutFloatOp : public OpKernel {
 public:
  explicit ZeroOutFloatOp(OpKernelConstruction* context)
      : OpKernel(context) {}

  void Compute(OpKernelContext* context) override {
    // Grab the input tensor
    const Tensor& input_tensor = context->input(0);
    auto input = input_tensor.flat<float>();

    // Create an output tensor
    Tensor* output = NULL;
    OP_REQUIRES_OK(context,
                   context->allocate_output(0, input_tensor.shape(), &output));
    auto output_flat = output->template flat<float>();

    // Set all the elements of the output tensor to 0
    const int N = input.size();
    for (int i = 0; i < N; i++) {
      output_flat(i) = 0;
    }

    // Preserve the first input value
    if (N > 0) output_flat(0) = input(0);
  }
};

// Note that TypeConstraint("T") means that attr "T" (defined
// in the op registration above) must be "int32" to use this template
// instantiation.
REGISTER_KERNEL_BUILDER(
    Name("ZeroOut")
    .Device(DEVICE_CPU)
    .TypeConstraint("T"),
    ZeroOutInt32Op);
REGISTER_KERNEL_BUILDER(
    Name("ZeroOut")
    .Device(DEVICE_CPU)
    .TypeConstraint<float>("T"),
    ZeroOutFloatOp);

To preserve backwards compatibility, you should
specify a default value when adding an attr to
an existing op:

c++
REGISTER_OP("ZeroOut")
  .Attr("T: {float, int32} = DT_INT32")
  .Input("to_zero: T")
  .Output("zeroed: T")

Let's say you wanted to add more types, say double:

c++
REGISTER_OP("ZeroOut")
    .Attr("T: {float, double, int32}")
    .Input("to_zero: T")
    .Output("zeroed: T");

Instead of writing another OpKernel with redundant code as above, often you
will be able to use a C++ template instead. You will still have one kernel
registration (REGISTER_KERNEL_BUILDER call) per overload.

c++
template <typename T>
class ZeroOutOp : public OpKernel {
 public:
  explicit ZeroOutOp(OpKernelConstruction* context) : OpKernel(context) {}
  
  void Compute(OpKernelContext* context) override {
    // Grab the input tensor
    const Tensor& input_tensor = context->input(0);
    auto input = input_tensor.flat<T>();
    
    // Create an output tensor
    Tensor* output = NULL;
    OP_REQUIRES_OK(context,
                   context->allocate_output(0, input_tensor.shape(), &output));
    auto output_flat = output->template flat<T>();
    
    // Set all the elements of the output tensor to 0
    const int N = input.size();
    for (int i = 0; i < N; i++) {
      output_flat(i) = 0;
    }
    
    // Preserve the first input value
    if (N > 0) output_flat(0) = input(0);
  }
};

// Note that TypeConstraint("T") means that attr "T" (defined
// in the op registration above) must be "int32" to use this template
// instantiation.
REGISTER_KERNEL_BUILDER(
    Name("ZeroOut")
    .Device(DEVICE_CPU)
    .TypeConstraint("T"),
    ZeroOutOp);
REGISTER_KERNEL_BUILDER(
    Name("ZeroOut")
    .Device(DEVICE_CPU)
    .TypeConstraint<float>("T"),
    ZeroOutOp<float>);
REGISTER_KERNEL_BUILDER(
    Name("ZeroOut")
    .Device(DEVICE_CPU)
    .TypeConstraint<double>("T"),
    ZeroOutOp<double>);

If you have more than a couple overloads, you can put the registration in a
macro.

c++
#include "tensorflow/core/framework/op_kernel.h"

#define REGISTER_KERNEL(type)                                       \
  REGISTER_KERNEL_BUILDER(                                          \
      Name("ZeroOut").Device(DEVICE_CPU).TypeConstraint<type>("T"), \
      ZeroOutOp<type>)

REGISTER_KERNEL(int32);
REGISTER_KERNEL(float);
REGISTER_KERNEL(double);

#undef REGISTER_KERNEL

Depending on the list of types you are registering the kernel for, you may be
able to use a macro provided by
tensorflow/core/framework/register_types.h:

c++
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/framework/register_types.h"

REGISTER_OP("ZeroOut")
    .Attr("T: realnumbertype")
    .Input("to_zero: T")
    .Output("zeroed: T");

template <typename T>
class ZeroOutOp : public OpKernel { ... };

#define REGISTER_KERNEL(type)                                       \
  REGISTER_KERNEL_BUILDER(                                          \
      Name("ZeroOut").Device(DEVICE_CPU).TypeConstraint<type>("T"), \
      ZeroOutOp<type>)

TF_CALL_REAL_NUMBER_TYPES(REGISTER_KERNEL);

#undef REGISTER_KERNEL

List Inputs and Outputs

In addition to being able to accept or produce different types, ops can consume
or produce a variable number of tensors.

In the next example, the attr T holds a list of types, and is used as the
type of both the input in and the output out. The input and output are
lists of tensors of that type (and the number and types of tensors in the output
are the same as the input, since both have type T).

c++
REGISTER_OP("PolymorphicListExample")
    .Attr("T: list(type)")
    .Input("in: T")
    .Output("out: T");

You can also place restrictions on what types can be specified in the list. In
this next case, the input is a list of float and double tensors. The op
accepts, for example, input types (float, double, float) and in that case the
output type would also be (float, double, float).

c++
REGISTER_OP("ListTypeRestrictionExample")
    .Attr("T: list({float, double})")
    .Input("in: T")
    .Output("out: T");

If you want all the tensors in a list to be of the same type, you might do
something like:

c++
REGISTER_OP("IntListInputExample")
    .Attr("N: int")
    .Input("in: N * int32")
    .Output("out: int32");

This accepts a list of int32 tensors, and uses an int attr N to
specify the length of the list.

This can be made type polymorphic as well. In the next
example, the input is a list of tensors (with length "N") of the same (but
unspecified) type ("T"), and the output is a single tensor of matching type:

c++
REGISTER_OP("SameListInputExample")
    .Attr("N: int")
    .Attr("T: type")
    .Input("in: N * T")
    .Output("out: T");

By default, tensor lists have a minimum length of 1. You can change that default
using
a ">=" constraint on the corresponding attr.
In this next example, the input is a list of at least 2 int32 tensors:

c++
REGISTER_OP("MinLengthIntListExample")
    .Attr("N: int >= 2")
    .Input("in: N * int32")
    .Output("out: int32");

The same syntax works with "list(type)" attrs:

c++
REGISTER_OP("MinimumLengthPolymorphicListExample")
    .Attr("T: list(type) >= 3")
    .Input("in: T")
    .Output("out: T");

Inputs and Outputs

To summarize the above, an op registration can have multiple inputs and outputs:

c++
REGISTER_OP("MultipleInsAndOuts")
    .Input("y: int32")
    .Input("z: float")
    .Output("a: string")
    .Output("b: int32");

Each input or output spec is of the form:

text
<name>:

where <name> begins with a letter and can be composed of alphanumeric
characters and underscores. `` is one of the following type
expressions:

  • <type>, where <type> is a supported input type (e.g. float, int32,
    string). This specifies a single tensor of the given type.

    See
    tf.DType.

    c++
    REGISTER_OP("BuiltInTypesExample")
        .Input("integers: int32")
        .Input("complex_numbers: complex64");
  • <attr-type>, where <attr-type> is the name of an Attr with type
    type or list(type) (with a possible type restriction). This syntax allows
    for polymorphic ops.

    c++
    REGISTER_OP("PolymorphicSingleInput")
        .Attr("T: type")
        .Input("in: T");
    
    REGISTER_OP("RestrictedPolymorphicSingleInput")
        .Attr("T: {int32, int64}")
        .Input("in: T");

    Referencing an attr of type list(type) allows you to accept a sequence of
    tensors.

    c++
    REGISTER_OP("ArbitraryTensorSequenceExample")
        .Attr("T: list(type)")
        .Input("in: T")
        .Output("out: T");
    
    REGISTER_OP("RestrictedTensorSequenceExample")
        .Attr("T: list({int32, int64})")
        .Input("in: T")
        .Output("out: T");

    Note that the number and types of tensors in the output out is the same as
    in the input in, since both are of type T.

  • For a sequence of tensors with the same type: <number> * <type>, where
    <number> is the name of an Attr with type int. The <type> can
    either be a tf.DType,
    or the name of an attr with type type. As an example of the first, this
    op accepts a list of int32 tensors:

    c++
    REGISTER_OP("Int32SequenceExample")
        .Attr("NumTensors: int")
        .Input("in: NumTensors * int32")

    Whereas this op accepts a list of tensors of any type, as long as they are all
    the same:

    c++
    REGISTER_OP("SameTypeSequenceExample")
        .Attr("NumTensors: int")
        .Attr("T: type")
        .Input("in: NumTensors * T")
  • For a reference to a tensor: Ref(<type>), where <type> is one of the
    previous types.

A note on naming: Any attr used in the type of an input will be inferred. By
convention those inferred attrs use capital names (like T or N).
Otherwise inputs, outputs, and attrs have names like function parameters
(e.g. num_outputs). For more details, see the
earlier note on naming.

For more details, see
tensorflow/core/framework/op_def_builder.h.

Backwards compatibility

Let's assume you have written a nice, custom op and shared it with others, so
you have happy customers using your operation. However, you'd like to make
changes to the op in some way.

In general, changes to existing, checked-in specifications must be
backwards-compatible: changing the specification of an op must not break prior
serialized GraphDef protocol buffers constructed from older specifications.
The details of GraphDef compatibility are
described here.

There are several ways to preserve backwards-compatibility.

  1. Any new attrs added to an operation must have default values defined, and
    with that default value the op must have the original behavior. To change an
    operation from not polymorphic to polymorphic, you must give a default
    value to the new type attr to preserve the original signature by default. For
    example, if your operation was:

    text
    REGISTER_OP("MyGeneralUnaryOp")
        .Input("in: float")
        .Output("out: float");

    you can make it polymorphic in a backwards-compatible way using:

    text
    REGISTER_OP("MyGeneralUnaryOp")
        .Input("in: T")
        .Output("out: T")
        .Attr("T: numerictype = DT_FLOAT");
  2. You can safely make a constraint on an attr less restrictive. For example,
    you can change from {int32, int64} to {int32, int64, float} or type.
    Or you may change from {"apple", "orange"} to {"apple", "banana", "orange"} or string.

  3. You can change single inputs / outputs into list inputs / outputs, as long as
    the default for the list type matches the old signature.

  4. You can add a new list input / output, if it defaults to empty.

  5. Namespace any new ops you create, by prefixing the op names with something
    unique to your project. This avoids having your op colliding with any ops
    that might be included in future versions of TensorFlow.

  6. Plan ahead! Try to anticipate future uses for the op. Some signature changes
    can't be done in a compatible way (for example, making a list of the same
    type into a list of varying types).

The full list of safe and unsafe changes can be found in
tensorflow/core/framework/op_compatibility_test.cc.
If you cannot make your change to an operation backwards compatible, then create
a new operation with a new name with the new semantics.

Also note that while these changes can maintain GraphDef compatibility, the
generated Python code may change in a way that isn't compatible with old
callers. The Python API may be kept compatible by careful changes in a
hand-written Python wrapper, by keeping the old signature except possibly adding
new optional arguments to the end. Generally incompatible changes may only be
made when TensorFlow changes major versions, and must conform to the
GraphDef version semantics.

GPU Support

You can implement different OpKernels and register one for CPU and another for
GPU, just like you can register kernels for different types.
There are several examples of kernels with GPU support in
tensorflow/core/kernels/.
Notice some kernels have a CPU version in a .cc file, a GPU version in a file
ending in _gpu.cu.cc, and some code shared in common in a .h file.

For example, the tf.pad has
everything but the GPU kernel in tensorflow/core/kernels/pad_op.cc.
The GPU kernel is in
tensorflow/core/kernels/pad_op_gpu.cu.cc,
and the shared code is a templated class defined in
tensorflow/core/kernels/pad_op.h.
We organize the code this way for two reasons: it allows you to share common
code among the CPU and GPU implementations, and it puts the GPU implementation
into a separate file so that it can be compiled only by the GPU compiler.

One thing to note, even when the GPU kernel version of pad is used, it still
needs its "paddings" input in CPU memory. To mark that inputs or outputs are
kept on the CPU, add a HostMemory() call to the kernel registration, e.g.:

c++
#define REGISTER_GPU_KERNEL(T)                         \
  REGISTER_KERNEL_BUILDER(Name("Pad")                  \
                              .Device(DEVICE_GPU)      \
                              .TypeConstraint<T>("T")  \
                              .HostMemory("paddings"), \
                          PadOp<GPUDevice, T>)

Compiling the kernel for the GPU device

Look at
cuda_op_kernel.cu.cc
for an example that uses a CUDA kernel to implement an op. The
tf_custom_op_library accepts a gpu_srcs argument in which the list of source
files containing the CUDA kernels (*.cu.cc files) can be specified. For use
with a binary installation of TensorFlow, the CUDA kernels have to be compiled
with NVIDIA's nvcc compiler. Here is the sequence of commands you can use to
compile the
cuda_op_kernel.cu.cc
and
cuda_op_kernel.cc
into a single dynamically loadable library:

bash
nvcc -std=c++11 -c -o cuda_op_kernel.cu.o cuda_op_kernel.cu.cc \
  ${TF_CFLAGS[@]} -D GOOGLE_CUDA=1 -x cu -Xcompiler -fPIC

g++ -std=c++11 -shared -o cuda_op_kernel.so cuda_op_kernel.cc \
  cuda_op_kernel.cu.o ${TF_CFLAGS[@]} -fPIC -lcudart ${TF_LFLAGS[@]}

cuda_op_kernel.so produced above can be loaded as usual in Python, using the
tf.load_op_library function.

Note that if your CUDA libraries are not installed in /usr/local/lib64,
you'll need to specify the path explicitly in the second (g++) command above.
For example, add -L /usr/local/cuda-8.0/lib64/ if your CUDA is installed in
/usr/local/cuda-8.0.

Note in some linux settings, additional options to nvcc compiling step are needed. Add -D_MWAITXINTRIN_H_INCLUDED to the nvcc command line to avoid errors from mwaitxintrin.h.

Implement the gradient in Python

Given a graph of ops, TensorFlow uses automatic differentiation
(backpropagation) to add new ops representing gradients with respect to the
existing ops.
To make automatic differentiation work for new ops, you must register a gradient
function which computes gradients with respect to the ops' inputs given
gradients with respect to the ops' outputs.

Mathematically, if an op computes \(y = f(x)\) the registered gradient op
converts gradients \(\partial L/ \partial y\) of loss \(L\) with respect to
\(y\) into gradients \(\partial L/ \partial x\) with respect to \(x\) via
the chain rule:

$$\frac{\partial L}{\partial x}
= \frac{\partial L}{\partial y} \frac{\partial y}{\partial x}
= \frac{\partial L}{\partial y} \frac{\partial f}{\partial x}.$$

In the case of ZeroOut, only one entry in the input affects the output, so the
gradient with respect to the input is a sparse "one hot" tensor. This is
expressed as follows:

python
from tensorflow.python.framework import ops
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import sparse_ops

@ops.RegisterGradient("ZeroOut")
def _zero_out_grad(op, grad):
  """The gradients for `zero_out`.

  Args:
    op: The `zero_out` `Operation` that we are differentiating, which we can use
      to find the inputs and outputs of the original op.
    grad: Gradient with respect to the output of the `zero_out` op.

  Returns:
    Gradients with respect to the input of `zero_out`.
  """
  to_zero = op.inputs[0]
  shape = array_ops.shape(to_zero)
  index = array_ops.zeros_like(shape)
  first_grad = array_ops.reshape(grad, [-1])[0]
  to_zero_grad = sparse_ops.sparse_to_dense([index], shape, first_grad, 0)
  return [to_zero_grad]  # List of one Tensor, since we have one input

Details about registering gradient functions with
tf.RegisterGradient:

  • For an op with one output, the gradient function will take an
    tf.Operation op and a
    tf.Tensor grad and build new ops
    out of the tensors
    op.inputs[i],
    op.outputs[i], and grad. Information
    about any attrs can be found via
    tf.Operation.get_attr.

  • If the op has multiple outputs, the gradient function will take op and
    grads, where grads is a list of gradients with respect to each output.
    The result of the gradient function must be a list of Tensor objects
    representing the gradients with respect to each input.

  • If there is no well-defined gradient for some input, such as for integer
    inputs used as indices, the corresponding returned gradient should be
    None. For example, for an op taking a floating point tensor x and an
    integer index i, the gradient function would return [x_grad, None].

  • If there is no meaningful gradient for the op at all, you often will not have
    to register any gradient, and as long as the op's gradient is never needed,
    you will be fine. In some cases, an op has no well-defined gradient but can
    be involved in the computation of the gradient. Here you can use
    ops.NotDifferentiable to automatically propagate zeros backwards.

Note that at the time the gradient function is called, only the data flow graph
of ops is available, not the tensor data itself. Thus, all computation must be
performed using other tensorflow ops, to be run at graph execution time.

Shape functions in C++

The TensorFlow API has a feature called "shape inference" that provides
information about the shapes of tensors without having to execute the
graph. Shape inference is supported by "shape functions" that are registered for
each op type in the C++ REGISTER_OP declaration, and perform two roles:
asserting that the shapes of the inputs are compatible during graph
construction, and specifying the shapes for the outputs.

Shape functions are defined as operations on the
shape_inference::InferenceContext class. For example, in the shape function
for ZeroOut:

c++
.SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) {
      c->set_output(0, c->input(0));
      return Status::OK();
    });

c->set_output(0, c->input(0)); declares that the first output's shape should
be set to the first input's shape. If the output is selected by its index as in the above example, the second parameter of set_output should be a ShapeHandle object. You can create an empty ShapeHandle object by its default constructor. The ShapeHandle object for an input with index idx can be obtained by c->input(idx).

There are a number of common shape functions
that apply to many ops, such as shape_inference::UnchangedShape which can be
found in common_shape_fns.h and used as follows:

c++
REGISTER_OP("ZeroOut")
    .Input("to_zero: int32")
    .Output("zeroed: int32")
    .SetShapeFn(::tensorflow::shape_inference::UnchangedShape);

A shape function can also constrain the shape of an input. For the version of
ZeroOut with a vector shape constraint, the shape function
would be as follows:

c++
.SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) {
      ::tensorflow::shape_inference::ShapeHandle input;
      TF_RETURN_IF_ERROR(c->WithRank(c->input(0), 1, &input));
      c->set_output(0, input);
      return Status::OK();
    });

The WithRank call validates that the input shape c->input(0) has
a shape with exactly one dimension (or if the input shape is unknown,
the output shape will be a vector with one unknown dimension).

If your op is polymorphic with multiple inputs, you can use
members of InferenceContext to determine the number of shapes to check, and
Merge to validate that the shapes are all compatible (alternatively, access
attributes that indicate the lengths, with InferenceContext::GetAttr, which
provides access to the attributes of the op).

c++
.SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) {
      ::tensorflow::shape_inference::ShapeHandle input;
      ::tensorflow::shape_inference::ShapeHandle output;
      for (size_t i = 0; i < c->num_inputs(); ++i) {
        TF_RETURN_IF_ERROR(c->WithRank(c->input(i), 2, &input));
        TF_RETURN_IF_ERROR(c->Merge(output, input, &output));
      }
      c->set_output(0, output);
      return Status::OK();
    });

Since shape inference is an optional feature, and the shapes of tensors may vary
dynamically, shape functions must be robust to incomplete shape information for
any of the inputs. The Merge method in InferenceContext
allows the caller to assert that two shapes are the same, even if either
or both of them do not have complete information. Shape functions are defined
for all of the core TensorFlow ops and provide many different usage examples.

The InferenceContext class has a number of functions that can be used to
define shape function manipulations. For example, you can validate that a
particular dimension has a very specific value using InferenceContext::Dim and
InferenceContext::WithValue; you can specify that an output dimension is the
sum / product of two input dimensions using InferenceContext::Add and
InferenceContext::Multiply. See the InferenceContext class for
all of the various shape manipulations you can specify. The following example sets
shape of the first output to (n, 3), where first input has shape (n, ...)

c++
.SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) {
    c->set_output(0, c->Matrix(c->Dim(c->input(0), 0), 3));
    return Status::OK();
});

If you have a complicated shape function, you should consider adding a test for
validating that various input shape combinations produce the expected output
shape combinations. You can see examples of how to write these tests in some
our
core ops tests.
(The syntax of INFER_OK and INFER_ERROR are a little cryptic, but try to be
compact in representing input and output shape specifications in tests. For
now, see the surrounding comments in those tests to get a sense of the shape
string specification).

Build a pip package for your custom op

To build a pip package for your op, see the
tensorflow/custom-op example. This
guide shows how to build custom ops from the TensorFlow pip package instead
of building TensorFlow from source.


File: site/en/r1/guide/performance/benchmarks.md

Benchmarks

Overview

A selection of image classification models were tested across multiple platforms
to create a point of reference for the TensorFlow community. The
Methodology section details how the tests were executed and has
links to the scripts used.

Results for image classification models

InceptionV3 (arXiv:1512.00567), ResNet-50
(arXiv:1512.03385), ResNet-152
(arXiv:1512.03385), VGG16
(arXiv:1409.1556), and
AlexNet
were tested using the ImageNet data set. Tests were
run on Google Compute Engine, Amazon Elastic Compute Cloud (Amazon EC2), and an
NVIDIA® DGX-1™. Most of the tests were run with both synthetic and real data.
Testing with synthetic data was done by using a tf.Variable set to the same
shape as the data expected by each model for ImageNet. We believe it is
important to include real data measurements when benchmarking a platform. This
load tests both the underlying hardware and the framework at preparing data for
actual training. We start with synthetic data to remove disk I/O as a variable
and to set a baseline. Real data is then used to verify that the TensorFlow
input pipeline and the underlying disk I/O are saturating the compute units.

Training with NVIDIA® DGX-1™ (NVIDIA® Tesla® P100)

Details and additional results are in the Details for NVIDIA® DGX-1™ (NVIDIA®
Tesla® P100)
section.

Training with NVIDIA® Tesla® K80

Details and additional results are in the Details for Google Compute Engine
(NVIDIA® Tesla® K80)
and
Details for Amazon EC2 (NVIDIA® Tesla®
K80)
sections.

Distributed training with NVIDIA® Tesla® K80

Details and additional results are in the Details for Amazon EC2 Distributed
(NVIDIA® Tesla® K80)

section.

Compare synthetic with real data training

NVIDIA® Tesla® P100

NVIDIA® Tesla® K80

Details for NVIDIA® DGX-1™ (NVIDIA® Tesla® P100)

Environment

  • Instance type: NVIDIA® DGX-1™
  • GPU: 8x NVIDIA® Tesla® P100
  • OS: Ubuntu 16.04 LTS with tests run via Docker
  • CUDA / cuDNN: 8.0 / 5.1
  • TensorFlow GitHub hash: b1e174e
  • Benchmark GitHub hash: 9165a70
  • Build Command: bazel build -c opt --copt=-march="haswell" --config=cuda //tensorflow/tools/pip_package:build_pip_package
  • Disk: Local SSD
  • DataSet: ImageNet
  • Test Date: May 2017

Batch size and optimizer used for each model are listed in the table below. In
addition to the batch sizes listed in the table, InceptionV3, ResNet-50,
ResNet-152, and VGG16 were tested with a batch size of 32. Those results are in
the other results section.

Options InceptionV3 ResNet-50 ResNet-152 AlexNet VGG16
Batch size per GPU 64 64 64 512 64
Optimizer sgd sgd sgd sgd sgd

Configuration used for each model.

Model variable_update local_parameter_device
InceptionV3 parameter_server cpu
ResNet50 parameter_server cpu
ResNet152 parameter_server cpu
AlexNet replicated (with NCCL) n/a
VGG16 replicated (with NCCL) n/a

Results

Training synthetic data

GPUs InceptionV3 ResNet-50 ResNet-152 AlexNet VGG16
1 142 219 91.8 2987 154
2 284 422 181 5658 295
4 569 852 356 10509 584
8 1131 1734 716 17822 1081

Training real data

GPUs InceptionV3 ResNet-50 ResNet-152 AlexNet VGG16
1 142 218 91.4 2890 154
2 278 425 179 4448 284
4 551 853 359 7105 534
8 1079 1630 708 N/A 898

Training AlexNet with real data on 8 GPUs was excluded from the graph and table
above due to it maxing out the input pipeline.

Other Results

The results below are all with a batch size of 32.

Training synthetic data

GPUs InceptionV3 ResNet-50 ResNet-152 VGG16
1 128 195 82.7 144
2 259 368 160 281
4 520 768 317 549
8 995 1485 632 820

Training real data

GPUs InceptionV3 ResNet-50 ResNet-152 VGG16
1 130 193 82.4 144
2 257 369 159 253
4 507 760 317 457
8 966 1410 609 690

Details for Google Compute Engine (NVIDIA® Tesla® K80)

Environment

  • Instance type: n1-standard-32-k80x8
  • GPU: 8x NVIDIA® Tesla® K80
  • OS: Ubuntu 16.04 LTS
  • CUDA / cuDNN: 8.0 / 5.1
  • TensorFlow GitHub hash: b1e174e
  • Benchmark GitHub hash: 9165a70
  • Build Command: bazel build -c opt --copt=-march="haswell" --config=cuda //tensorflow/tools/pip_package:build_pip_package
  • Disk: 1.7 TB Shared SSD persistent disk (800 MB/s)
  • DataSet: ImageNet
  • Test Date: May 2017

Batch size and optimizer used for each model are listed in the table below. In
addition to the batch sizes listed in the table, InceptionV3 and ResNet-50 were
tested with a batch size of 32. Those results are in the other results
section.

Options InceptionV3 ResNet-50 ResNet-152 AlexNet VGG16
Batch size per GPU 64 64 32 512 32
Optimizer sgd sgd sgd sgd sgd

The configuration used for each model was variable_update equal to
parameter_server and local_parameter_device equal to cpu.

Results

Training synthetic data

GPUs InceptionV3 ResNet-50 ResNet-152 AlexNet VGG16
1 30.5 51.9 20.0 656 35.4
2 57.8 99.0 38.2 1209 64.8
4 116 195 75.8 2328 120
8 227 387 148 4640 234

Training real data

GPUs InceptionV3 ResNet-50 ResNet-152 AlexNet VGG16
1 30.6 51.2 20.0 639 34.2
2 58.4 98.8 38.3 1136 62.9
4 115 194 75.4 2067 118
8 225 381 148 4056 230

Other Results

Training synthetic data

GPUs InceptionV3 (batch size 32) ResNet-50 (batch size 32)
1 29.3 49.5
2 55.0 95.4
4 109 183
8 216 362

Training real data

GPUs InceptionV3 (batch size 32) ResNet-50 (batch size 32)
1 29.5 49.3
2 55.4 95.3
4 110 186
8 216 359

Details for Amazon EC2 (NVIDIA® Tesla® K80)

Environment

  • Instance type: p2.8xlarge
  • GPU: 8x NVIDIA® Tesla® K80
  • OS: Ubuntu 16.04 LTS
  • CUDA / cuDNN: 8.0 / 5.1
  • TensorFlow GitHub hash: b1e174e
  • Benchmark GitHub hash: 9165a70
  • Build Command: bazel build -c opt --copt=-march="haswell" --config=cuda //tensorflow/tools/pip_package:build_pip_package
  • Disk: 1TB Amazon EFS (burst 100 MiB/sec for 12 hours, continuous 50
    MiB/sec)
  • DataSet: ImageNet
  • Test Date: May 2017

Batch size and optimizer used for each model are listed in the table below. In
addition to the batch sizes listed in the table, InceptionV3 and ResNet-50 were
tested with a batch size of 32. Those results are in the other results
section.

Options InceptionV3 ResNet-50 ResNet-152 AlexNet VGG16
Batch size per GPU 64 64 32 512 32
Optimizer sgd sgd sgd sgd sgd

Configuration used for each model.

Model variable_update local_parameter_device
InceptionV3 parameter_server cpu
ResNet-50 replicated (without NCCL) gpu
ResNet-152 replicated (without NCCL) gpu
AlexNet parameter_server gpu
VGG16 parameter_server gpu

Results

Training synthetic data

GPUs InceptionV3 ResNet-50 ResNet-152 AlexNet VGG16
1 30.8 51.5 19.7 684 36.3
2 58.7 98.0 37.6 1244 69.4
4 117 195 74.9 2479 141
8 230 384 149 4853 260

Training real data

GPUs InceptionV3 ResNet-50 ResNet-152 AlexNet VGG16
1 30.5 51.3 19.7 674 36.3
2 59.0 94.9 38.2 1227 67.5
4 118 188 75.2 2201 136
8 228 373 149 N/A 242

Training AlexNet with real data on 8 GPUs was excluded from the graph and table
above due to our EFS setup not providing enough throughput.

Other Results

Training synthetic data

GPUs InceptionV3 (batch size 32) ResNet-50 (batch size 32)
1 29.9 49.0
2 57.5 94.1
4 114 184
8 216 355

Training real data

GPUs InceptionV3 (batch size 32) ResNet-50 (batch size 32)
1 30.0 49.1
2 57.5 95.1
4 113 185
8 212 353

Details for Amazon EC2 Distributed (NVIDIA® Tesla® K80)

Environment

  • Instance type: p2.8xlarge
  • GPU: 8x NVIDIA® Tesla® K80
  • OS: Ubuntu 16.04 LTS
  • CUDA / cuDNN: 8.0 / 5.1
  • TensorFlow GitHub hash: b1e174e
  • Benchmark GitHub hash: 9165a70
  • Build Command: bazel build -c opt --copt=-march="haswell" --config=cuda //tensorflow/tools/pip_package:build_pip_package
  • Disk: 1.0 TB EFS (burst 100 MB/sec for 12 hours, continuous 50 MB/sec)
  • DataSet: ImageNet
  • Test Date: May 2017

The batch size and optimizer used for the tests are listed in the table. In
addition to the batch sizes listed in the table, InceptionV3 and ResNet-50 were
tested with a batch size of 32. Those results are in the other results
section.

Options InceptionV3 ResNet-50 ResNet-152
Batch size per GPU 64 64 32
Optimizer sgd sgd sgd

Configuration used for each model.

Model variable_update local_parameter_device cross_replica_sync
InceptionV3 distributed_replicated n/a True
ResNet-50 distributed_replicated n/a True
ResNet-152 distributed_replicated n/a True

To simplify server setup, EC2 instances (p2.8xlarge) running worker servers also
ran parameter servers. Equal numbers of parameter servers and worker servers were
used with the following exceptions:

  • InceptionV3: 8 instances / 6 parameter servers
  • ResNet-50: (batch size 32) 8 instances / 4 parameter servers
  • ResNet-152: 8 instances / 4 parameter servers

Results

Training synthetic data

GPUs InceptionV3 ResNet-50 ResNet-152
1 29.7 52.4 19.4
8 229 378 146
16 459 751 291
32 902 1388 565
64 1783 2744 981

Other Results

Training synthetic data

GPUs InceptionV3 (batch size 32) ResNet-50 (batch size 32)
1 29.2 48.4
8 219 333
16 427 667
32 820 1180
64 1608 2315

Methodology

This
script
was run on the various platforms to generate the above results.

In order to create results that are as repeatable as possible, each test was run
5 times and then the times were averaged together. GPUs are run in their default
state on the given platform. For NVIDIA® Tesla® K80 this means leaving on GPU
Boost
.
For each test, 10 warmup steps are done and then the next 100 steps are
averaged.


File: site/en/r1/guide/performance/datasets.md

Data Input Pipeline Performance

GPUs and TPUs can radically reduce the time required to execute a single
training step. Achieving peak performance requires an efficient input pipeline
that delivers data for the next step before the current step has finished. The
tf.data API helps to build flexible and efficient input pipelines. This
document explains the tf.data API's features and best practices for building
high performance TensorFlow input pipelines across a variety of models and
accelerators.

This guide does the following:

  • Illustrates that TensorFlow input pipelines are essentially an
    ETL process.
  • Describes common performance optimizations in the context of the tf.data
    API.
  • Discusses the performance implications of the order in which you apply
    transformations.
  • Summarizes the best practices for designing performant TensorFlow input
    pipelines.

Input Pipeline Structure

A typical TensorFlow training input pipeline can be framed as an ETL process:

  1. Extract: Read data from persistent storage -- either local (e.g. HDD or
    SSD) or remote (e.g. GCS or
    HDFS).
  2. Transform: Use CPU cores to parse and perform preprocessing operations
    on the data such as image decompression, data augmentation transformations
    (such as random crop, flips, and color distortions), shuffling, and batching.
  3. Load: Load the transformed data onto the accelerator device(s) (for
    example, GPU(s) or TPU(s)) that execute the machine learning model.

This pattern effectively utilizes the CPU, while reserving the accelerator for
the heavy lifting of training your model. In addition, viewing input pipelines
as an ETL process provides structure that facilitates the application of
performance optimizations.

When using the tf.estimator.Estimator API, the first two phases (Extract and
Transform) are captured in the input_fn passed to
tf.estimator.Estimator.train. In code, this might look like the following
(naive, sequential) implementation:

text
def parse_fn(example):
  "Parse TFExample records and perform simple data augmentation."
  example_fmt = {
    "image": tf.FixedLengthFeature((), tf.string, ""),
    "label": tf.FixedLengthFeature((), tf.int64, -1)
  }
  parsed = tf.parse_single_example(example, example_fmt)
  image = tf.image.decode_image(parsed["image"])
  image = _augment_helper(image)  # augments image using slice, reshape, resize_bilinear
  return image, parsed["label"]

def input_fn():
  files = tf.data.Dataset.list_files("/path/to/dataset/train-*.tfrecord")
  dataset = files.interleave(tf.data.TFRecordDataset)
  dataset = dataset.shuffle(buffer_size=FLAGS.shuffle_buffer_size)
  dataset = dataset.map(map_func=parse_fn)
  dataset = dataset.batch(batch_size=FLAGS.batch_size)
  return dataset

The next section builds on this input pipeline, adding performance
optimizations.

Optimizing Performance

As new computing devices (such as GPUs and TPUs) make it possible to train
neural networks at an increasingly fast rate, the CPU processing is prone to
becoming the bottleneck. The tf.data API provides users with building blocks
to design input pipelines that effectively utilize the CPU, optimizing each step
of the ETL process.

Pipelining

To perform a training step, you must first extract and transform the training
data and then feed it to a model running on an accelerator. However, in a naive
synchronous implementation, while the CPU is preparing the data, the accelerator
is sitting idle. Conversely, while the accelerator is training the model, the
CPU is sitting idle. The training step time is thus the sum of both CPU
pre-processing time and the accelerator training time.

Pipelining overlaps the preprocessing and model execution of a training
step. While the accelerator is performing training step N, the CPU is
preparing the data for step N+1. Doing so reduces the step time to the maximum
(as opposed to the sum) of the training and the time it takes to extract and
transform the data.

Without pipelining, the CPU and the GPU/TPU sit idle much of the time:

With pipelining, idle time diminishes significantly:

The tf.data API provides a software pipelining mechanism through the
tf.data.Dataset.prefetch transformation, which can be used to decouple the
time data is produced from the time it is consumed. In particular, the
transformation uses a background thread and an internal buffer to prefetch
elements from the input dataset ahead of the time they are requested. Thus, to
achieve the pipelining effect illustrated above, you can add prefetch(1) as
the final transformation to your dataset pipeline (or prefetch(n) if a single
training step consumes n elements).

To apply this change to our running example, change:

text
dataset = dataset.batch(batch_size=FLAGS.batch_size)
return dataset

to:

text
dataset = dataset.batch(batch_size=FLAGS.batch_size)
dataset = dataset.prefetch(buffer_size=FLAGS.prefetch_buffer_size)
return dataset

Note that the prefetch transformation will yield benefits any time there is an
opportunity to overlap the work of a "producer" with the work of a "consumer."
The preceding recommendation is simply the most common application.

Parallelize Data Transformation

When preparing a batch, input elements may need to be pre-processed. To this
end, the tf.data API offers the tf.data.Dataset.map transformation, which
applies a user-defined function (for example, parse_fn from the running
example) to each element of the input dataset. Because input elements are
independent of one another, the pre-processing can be parallelized across
multiple CPU cores. To make this possible, the map transformation provides the
num_parallel_calls argument to specify the level of parallelism. For example,
the following diagram illustrates the effect of setting num_parallel_calls=2
to the map transformation:

Choosing the best value for the num_parallel_calls argument depends on your
hardware, characteristics of your training data (such as its size and shape),
the cost of your map function, and what other processing is happening on the
CPU at the same time; a simple heuristic is to use the number of available CPU
cores. For instance, if the machine executing the example above had 4 cores, it
would have been more efficient to set num_parallel_calls=4. On the other hand,
setting num_parallel_calls to a value much greater than the number of
available CPUs can lead to inefficient scheduling, resulting in a slowdown.

To apply this change to our running example, change:

text
dataset = dataset.map(map_func=parse_fn)

to:

text
dataset = dataset.map(map_func=parse_fn, num_parallel_calls=FLAGS.num_parallel_calls)

Furthermore, if your batch size is in the hundreds or thousands, your pipeline
will likely additionally benefit from parallelizing the batch creation. To this
end, the tf.data API provides the tf.data.experimental.map_and_batch
transformation, which effectively "fuses" the map and batch transformations.

To apply this change to our running example, change:

text
dataset = dataset.map(map_func=parse_fn, num_parallel_calls=FLAGS.num_parallel_calls)
dataset = dataset.batch(batch_size=FLAGS.batch_size)

to:

text
dataset = dataset.apply(tf.contrib.data.map_and_batch(
    map_func=parse_fn, batch_size=FLAGS.batch_size))

Parallelize Data Extraction

In a real-world setting, the input data may be stored remotely (for example,
GCS or HDFS), either because the input data would not fit locally or because the
training is distributed and it would not make sense to replicate the input data
on every machine. A dataset pipeline that works well when reading data locally
might become bottlenecked on I/O when reading data remotely because of the
following differences between local and remote storage:

  • Time-to-first-byte: Reading the first byte of a file from remote storage
    can take orders of magnitude longer than from local storage.
  • Read throughput: While remote storage typically offers large aggregate
    bandwidth, reading a single file might only be able to utilize a small
    fraction of this bandwidth.

In addition, once the raw bytes are read into memory, it may also be necessary
to deserialize or decrypt the data
(e.g. protobuf), which adds
additional overhead. This overhead is present irrespective of whether the data
is stored locally or remotely, but can be worse in the remote case if data is
not prefetched effectively.

To mitigate the impact of the various data extraction overheads, the tf.data
API offers the tf.data.experimental.parallel_interleave transformation. Use this
transformation to parallelize the execution of and interleave the contents of
other datasets (such as data file readers). The
number of datasets to overlap can be specified by the cycle_length argument.

The following diagram illustrates the effect of supplying cycle_length=2 to
the parallel_interleave transformation:

To apply this change to our running example, change:

text
dataset = files.interleave(tf.data.TFRecordDataset)

to:

text
dataset = files.apply(tf.contrib.data.parallel_interleave(
    tf.data.TFRecordDataset, cycle_length=FLAGS.num_parallel_readers))

The throughput of remote storage systems can vary over time due to load or
network events. To account for this variance, the parallel_interleave
transformation can optionally use prefetching. (See
tf.data.experimental.parallel_interleave for details).

By default, the parallel_interleave transformation provides a deterministic
ordering of elements to aid reproducibility. As an alternative to prefetching
(which may be ineffective in some cases), the parallel_interleave
transformation also provides an option that can boost performance at the expense
of ordering guarantees. In particular, if the sloppy argument is set to true,
the transformation may depart from its otherwise deterministic ordering, by
temporarily skipping over files whose elements are not available when the next
element is requested.

Performance Considerations

The tf.data API is designed around composable transformations to provide its
users with flexibility. Although many of these transformations are commutative,
the ordering of certain transformations has performance implications.

Map and Batch

Invoking the user-defined function passed into the map transformation has
overhead related to scheduling and executing the user-defined function.
Normally, this overhead is small compared to the amount of computation performed
by the function. However, if map does little work, this overhead can dominate
the total cost. In such cases, we recommend vectorizing the user-defined
function (that is, have it operate over a batch of inputs at once) and apply the
batch transformation before the map transformation.

Map and Cache

The tf.data.Dataset.cache transformation can cache a dataset, either in
memory or on local storage. If the user-defined function passed into the map
transformation is expensive, apply the cache transformation after the map
transformation as long as the resulting dataset can still fit into memory or
local storage. If the user-defined function increases the space required to
store the dataset beyond the cache capacity, consider pre-processing your data
before your training job to reduce resource usage.

Map and Interleave / Prefetch / Shuffle

A number of transformations, including interleave, prefetch, and shuffle,
maintain an internal buffer of elements. If the user-defined function passed
into the map transformation changes the size of the elements, then the
ordering of the map transformation and the transformations that buffer elements
affects the memory usage. In general, we recommend choosing the order that
results in lower memory footprint, unless different ordering is desirable for
performance (for example, to enable fusing of the map and batch transformations).

Repeat and Shuffle

The tf.data.Dataset.repeat transformation repeats the input data a finite (or
infinite) number of times; each repetition of the data is typically referred to
as an epoch. The tf.data.Dataset.shuffle transformation randomizes the
order of the dataset's examples.

If the repeat transformation is applied before the shuffle transformation,
then the epoch boundaries are blurred. That is, certain elements can be repeated
before other elements appear even once. On the other hand, if the shuffle
transformation is applied before the repeat transformation, then performance
might slow down at the beginning of each epoch related to initialization of the
internal state of the shuffle transformation. In other words, the former
(repeat before shuffle) provides better performance, while the latter
(shuffle before repeat) provides stronger ordering guarantees.

When possible, we recommend using the fused
tf.data.experimental.shuffle_and_repeat transformation, which combines the best of
both worlds (good performance and strong ordering guarantees). Otherwise, we
recommend shuffling before repeating.

Summary of Best Practices

Here is a summary of the best practices for designing input pipelines:

  • Use the prefetch transformation to overlap the work of a producer and
    consumer. In particular, we recommend adding prefetch(n) (where n is the
    number of elements / batches consumed by a training step) to the end of your
    input pipeline to overlap the transformations performed on the CPU with the
    training done on the accelerator.
  • Parallelize the map transformation by setting the num_parallel_calls
    argument. We recommend using the number of available CPU cores for its value.
  • If you are combining pre-processed elements into a batch using the batch
    transformation, we recommend using the fused map_and_batch transformation;
    especially if you are using large batch sizes.
  • If you are working with data stored remotely and / or requiring
    deserialization, we recommend using the parallel_interleave
    transformation to overlap the reading (and deserialization) of data from
    different files.
  • Vectorize cheap user-defined functions passed in to the map transformation
    to amortize the overhead associated with scheduling and executing the
    function.
  • If your data can fit into memory, use the cache transformation to cache it
    in memory during the first epoch, so that subsequent epochs can avoid the
    overhead associated with reading, parsing, and transforming it.
  • If your pre-processing increases the size of your data, we recommend
    applying the interleave, prefetch, and shuffle first (if possible) to
    reduce memory usage.
  • We recommend applying the shuffle transformation before the repeat
    transformation, ideally using the fused shuffle_and_repeat transformation.