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 withtensor.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 callingdataToGPU. Also see this example code for details.
// 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:
000|001|002|003 010|011|012|013 020|021|022|023
100|101|102|103 110|111|112|113 120|121|122|123So 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 callingdataToGPU. Also see this example code for details.
// 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
- Export your access key of BrowserStack's Automate service:
export BROWSERSTACK_USERNAME=YOUR_USERNAME
export BROWSERSTACK_ACCESS_KEY=YOUR_ACCESS_KEY- Download and run the tool:
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.
Open http://localhost:8001/ and start to benchmark.
3.1 If you want to benchmark code snippet. Please updatebenchmarkCodeSnippetwith your code snippet before runningnode app.jsand selectcodeSnippetinmodel name: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.
shellnode app.js --benchmarks=relative_file_path.jsonA 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.
shellnode app.js --cloud--firestore
- Pushes successful benchmark results to a Firestore database.
shellnode app.js --firestore--h, --help
- Shows help menu and all optional arguments in the shell window.
shellnode app.js --hor
shellnode 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--benchmarksis set.
shellnode app.js --period=15- Runs a part of models specified in
--date
- Set the date for selecting models and this works only if
--periodis set. The value could be 1~31. If it is not declared, the date will the real date at runtime.
shellnode app.js --period=15 --date=1- Set the date for selecting models and this works only if
--maxBenchmarks
- Sets maximum for number of benchmarks run in parallel. Expects a positive integer.
shellnode app.js --maxBenchmarks=positive_integer--maxTries
- Sets maximum for number of tries a given benchmark has to succeed. Expects a positive integer.
shellnode 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.
shellnode app.js --outfile=js--v, --version
- Shows node version in use.
shellnode app.js --vor
shellnode 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.)
shellnode app.js --localBuild=core,webgl,wasm,cpu,layers,converter,automl- Uses local build dependencies, instead of public CDNs. (When using localBuild targets, please make sure you have built the targets (eg. run
--npmVersion
- Specify the npm version of TFJS library to benchmark. By default the latest version of TFJS will be benchmarked.
node app.js --npmVersion=4.4.0Custom model
The custom model is supported, but is constrained by:
- A URL path to the model is required, while the model in local file system is not supported. The following URLs are examples:
- Currently only
tf.GraphModelandtf.LayersModelare supported.
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:
- benchmark_models.js warps the all benchmark logics into a Jasmine spec.
- browser_list.json lists the supported BrowserStack combinations. If you want to add more combinations or refactor this list, you can follow this conversation.
- 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:
- Download the tool.
git clone https://github.com/tensorflow/tfjs.git
cd tfjs/e2e/benchmarks/- Download the model.
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/- Run a http server to host the model and the local benchmark tool.
npx http-server- Open http://127.0.0.1:8080/local-benchmark/ through the browser.
- Select
customin themodelsfield. - Fill
http://127.0.0.1:8080/model/model.jsoninto themodelUrlfield. - Run benchmark.
Paths to custom models
The benchmark tool suopports three kinds of paths to the custom models.
URL
Examples:
- TF Hub: https://tfhub.dev/google/tfjs-model/imagenet/resnet_v2_50/feature_vector/1/default/1
- Storage: https://storage.googleapis.com/tfjs-models/savedmodel/mobilenet_v2_1.0_224/model.json
LocalStorage
Store the model in LocalStorage at first. Run the following codes in the browser console:
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:
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:
export TAGS=#SMOKE,#REGRESSION,#GOLDEN
yarn testFilter tests by grep:
yarn test --grep cpuAdd 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:
gcloud functions deploy nightly_tfjs \
--runtime nodejs14 \
--trigger-topic nightly_tfjsIf 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:
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:
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:
- At 4am, Cloud Scheduler writes to
nightly_tfjstopic - That triggers the
nightly_tfjsfunction, which starts a build programmatically - That build runs and writes its status to
cloud-buildstopic - That triggers the
send_emailfunction, 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
npm i @tensorflow/tfjs-automlIf you are using CDN:
<script src="https://cdn.jsdelivr.net/npm/@tensorflow/tfjs-automl"></script>We support the following types of AutoML Edge models:
Image classification
AutoML Image classification model will output the following set of files:
model.json, the model topologydict.txt, a newline-separated list of labels- One or more of
*.binfiles 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:
cd demo/img_classification
yarn
yarn watchThis will start a local HTTP server on port 1234 that serves the demo.
Loading the model
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:
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 beHTMLImageElement,HTMLCanvasElement,HTMLVideoElement,ImageData or
a 3D Tensor:
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:
[
{label: "daisy", prob: 0.931},
{label: "dandelion", prob: 0.027},
{label: "roses", prob: 0.013},
...
]Advanced usage
Advanced users can access the underlyingGraphModel viamodel.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 topologydict.txt, a newline-separated list of labels- One or more of
*.binfiles 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:
cd demo/object_detection
yarn
yarn watchThis will start a local HTTP server on port 1234 that serves the demo.
Loading the model
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:
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 beHTMLImageElement,HTMLCanvasElement,HTMLVideoElement,ImageData or
a 3D Tensor:
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 thetopkmost 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:
[
{
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 underlyingGraphModel viamodel.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
// 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
<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
<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:
- Move into
tfjs-backend-nodegl(parent directory of this demo folder):
$ cd tfjs-backend-nodegl- Build package and compile TypeScript:
$ yarn && yarn tsc- Move into the demo directory:
$ cd demo- Prep and build demo:
$ yarn- Run demo:
$ node run_mobilenet_inference.js dog.jpgExpected output:
$ 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.75772546708584msFile: 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
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.
- Prunes a specific subgraph from the graph, as defined by the arguments
- 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.
- Schedule the execution of graph operations using kernel implementations
- 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(orMetaGraphDef) 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
mainGraphDef. Defines aFunctionDefin theFunctionDefLibrary
included in aGraphDef. - 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 |
Recommended Approach
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 byOperations) and
corresponds to aTF_Graphin the C API. Mainly used as an argument when
creating newOperationobjects and when starting aSession. Also
supports iterating through the operations in the graph
(TF_GraphNextOperation), looking up operations by name
(TF_GraphOperationByName), and converting to and from aGraphDef
protocol message (TF_GraphToGraphDefandTF_GraphImportGraphDefin the C
API).Operation: Represents a computation node in the graph. Corresponds to aTF_Operationin the C API.Output: Represents one of the outputs of an operation in the graph. Has aDataType(and eventually a shape). May be passed as an input argument to a
function for adding operations to a graph, or to aSession'sRun()
method to fetch that output as a tensor. Corresponds to aTF_Outputin the
C API.Session: Represents a client to a particular instance of the TensorFlow
runtime. Its main job is to be constructed with aGraphand some options
and then field calls toRun()the graph. Corresponds to aTF_Sessionin
the C API.Tensor: Represents an N-dimensional (rectangular) array with elements all
the sameDataType. Gets data into and out of aSession'sRun()call.
Corresponds to aTF_Tensorin the C API.DataType: An enumerant with all the possible tensor types supported by
TensorFlow. Corresponds toTF_DataTypein the C API and often referred to
asdtypein 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_GetAllOpListin the C API retrieves all registeredOpDefprotocol
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 theOpDefmessages.- The C++ function
OpRegistry::Global()->GetRegisteredOps()returns the same
list of all registeredOpDefs (defined intensorflow/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.pbtxtby 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 theTF_OperationDescription*. - Call
TF_AddInput()orTF_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 runningTF_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 likego get for Go and cargo ops for Rust. At the other end of the spectrum, for
some languages the code could be generated dynamically fromtensorflow/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 explicitScope 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 theWithOpName() 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. ASparseTensor 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
Constoperation to a graph typically is a
wrapper since the generated function will typically have redundantDataTypeinputs.
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.
// 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 aBUILD file in the same directory with the following contents:
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):
bazel run -c opt //tensorflow/cc/example:exampleThis 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 atensorflow::Status object which will be used to indicate errors encountered
when constructing operations. The Scope class has value semantics, thus, aScope 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 returnedScope 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:
// Not recommended
MatMul m(scope, a, b);Instead, we recommend the following "functional" style for constructing
operations:
// 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 suchAttrs in multiple ways:
- You can specify a single optional attribute by constructing an
Attrsobject
using thestaticfunctions provided in the C++ class for the operation. For
example:
auto m = MatMul(scope, a, b, MatMul::TransposeA(true));- You can specify multiple optional attributes by chaining together functions
available in theAttrsstruct. For example:
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:
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 theOpDeffor the
operation. For example:
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:
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 thetensorflow::ops::Const function from various kinds of C++ values. For
example:
- Scalars
auto f = Const(scope, 42.0f);
auto s = Const(scope, "hello world!");- Nested initializer lists
// 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
// 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:
// [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 atensorflow::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:
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:
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, andReadOnlyMemoryRegion. - Implement the
FileSysteminterface as a subclass. - Register the
FileSystemimplementation 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 thatFileSystem implementation.
To implement a RandomAccessFile, you must implement a single interface calledRead(), 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.
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.
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 callREGISTER_FILE_SYSTEM::
REGISTER_FILE_SYSTEM("foobar", FooBarFileSystem);When TensorFlow tries to operate on a file whose path starts with foobar://,
it will use the FooBarFileSystem implementation.
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:
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:
- Some familiarity with C++.
- Must have downloaded TensorFlow source, and be
able to build it.
We divide the task of supporting a file format into two pieces:
- File formats: We use a reader
tf.data.Datasetto 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:
tf.data.TFRecordDataset
(source inkernels/data/reader_dataset_ops.cc)tf.data.FixedLengthRecordDataset
(source inkernels/data/reader_dataset_ops.cc)tf.data.TextLineDataset
(source inkernels/data/reader_dataset_ops.cc)
Each of these implementations comprises three related classes:
A
tensorflow::DatasetOpKernelsubclass (e.g.TextLineDatasetOp), which
tells TensorFlow how to construct a dataset object from the inputs to and
attrs of an op, in itsMakeDataset()method.A
tensorflow::GraphDatasetBasesubclass (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 itsMakeIteratorInternal()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 itsGetNextInternal()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 moreTensor objects.
To create a new reader dataset called (for example) MyReaderDataset, you will
need to:
- In C++, define subclasses of
tensorflow::DatasetOpKernel,tensorflow::GraphDatasetBase, andtensorflow::DatasetIterator<Dataset>
that implement the reading logic. - In C++, register a new reader op and kernel with the name
"MyReaderDataset". - In Python, define a subclass of
tf.data.DatasetcalledMyReaderDataset.
You can put all the C++ code in a single file, such asmy_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:
/* 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:
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:
passYou can see some examples of Dataset wrapper classes intensorflow/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(andtf.parse_example)tf.decode_csvtf.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:
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.
with open(FLAGS.graph, "rb") as f:Here we get a file handle for the path we've passed in to the script
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 actuallyParseFromString(), 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:
for node in graph_def.nodeEach 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 aNodeDef, 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 theREGISTER_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 aConst 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:
- 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. - 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). - 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. - Write a function to compute gradients for the op (optional).
- 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
Pythontf.test.compute_gradient_error.
Seerelu_op_test.pyas
an example that tests the forward functions of Relu-like operators and
their gradients.
PREREQUISITES:
- Some familiarity with C++.
- Must have installed the
TensorFlow binary, or must have
downloaded TensorFlow source,
and be able to build it.
[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 ofint32s 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 theREGISTER_OP macro that defines the interface for your op:
#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:
#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:
REGISTER_KERNEL_BUILDER(Name("ZeroOut").Device(DEVICE_CPU), ZeroOutOp);Important: Instances of your OpKernel may be accessed concurrently.
YourComputemethod 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 aResourceMgr
to keep track of op state.
Multi-threaded CPU kernels
To write a multi-threaded CPU kernel, the Shard function inwork_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 inconfig.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:
- Define the OpKernel templated on the Device and the primitive type of the
tensor. - To do the actual computation of the output, the Compute function calls a
templated functor struct. - 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.
// 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_/* Detailed source-code truncated for AI context efficiency. */// 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_CUDABuild 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 theget_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.
$ 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.
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[@]} -O2On macOS, the additional flag "-undefined dynamic_lookup" is required when
building the .so file.
Note on
gccversion>=5: gcc uses the new C++
ABI since version5.
TensorFlow 2.8 and earlier were built withgcc4that uses the older ABI. If
you are using these versions of TensorFlow and are trying to compile your op
library withgcc>=5, add-D_GLIBCXX_USE_CXX11_ABI=0to 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.
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.
$ bazel build --config opt //tensorflow/core/user_ops:zero_out.soAs 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
.sofile) with the
standardcc_libraryrule, we strongly recommend that you use the
tf_custom_op_librarymacro. 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 thetf.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:
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:
import tensorflow as tf
zero_out_module = tf.load_op_library('./zero_out.so')
zero_out = zero_out_module.zero_outVerify 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 filezero_out_op_test.py with the contents:
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):
$ python zero_out_op_test.pyBuilding 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
- Op registration
- GPU support
- Implement the gradient in Python
- Shape functions in C++
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.
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 theInvalidArgument status if it isn't. TheOP_REQUIRES macro takes three arguments:
- The
context, which can either be anOpKernelContextorOpKernelConstructionpointer (seetensorflow/core/framework/op_kernel.h),
for itsSetStatus()method. - The condition. For example, there are functions for validating the shape
of a tensor intensorflow/core/framework/tensor_shape.h - The error itself, which is represented by a
Statusobject, seetensorflow/core/lib/core/status.h. AStatushas both a type (frequentlyInvalidArgument, but see the list of
types) and a message. Functions for constructing an error may be found intensorflow/core/lib/core/errors.h.
Alternatively, if you want to test whether a Status object returned from some
function is an error, and if so return it, useOP_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:
<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:
REGISTER_OP("ZeroOut")
.Attr("preserve_index: int")
.Input("to_zero: int32")
.Output("zeroed: int32");(Note that the set of attribute types is different from thetf.DType used for inputs and outputs.)
Your kernel can then access this attr in its constructor via the context
parameter:
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:
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 ofDataType.shape: ATensorShapeProto.tensor: ATensorProto.list(<type>): A list of<type>, where<type>is one of the above types.
Note thatlist(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 typetype, and must be one of<type1>or<type2>, where<type1>and<type2>are supportedtf.DType. You don't specify
that the type of the attr istype. This is implied when you have a list of
types in{...}. For example, in this case the attrtis a type that must
be anint32, afloat, or abool:c++REGISTER_OP("RestrictedTypeExample") .Attr("t: {int32, float, bool}");There are shortcuts for common type constraints:
numbertype: Typetyperestricted to the numeric (non-string and
non-bool) types.realnumbertype: Likenumbertypewithout complex types.quantizedtype: Likenumbertypebut just the quantized number types.
The specific lists of types allowed by these are defined by the functions
(likeNumberTypes()) in
tensorflow/core/framework/types.h.
In this example the attrtmust be one of the numeric types:c++REGISTER_OP("NumberType") .Attr("t: numbertype");For this op:
pythontf.number_type(t=tf.int32) # Valid tf.number_type(t=tf.bool) # InvalidLists can be combined with other lists and single types. The following
op allows attrtto be any of the numeric types, or the bool type:c++REGISTER_OP("NumberOrBooleanType") .Attr("t: {numbertype, bool}");For this op:
pythontf.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) # Invalidint >= <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
amust
have a value that is at least2: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
ais a
list of types (eitherint32orfloat), 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:
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:
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:
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, orint32, 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:pythondef 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_zerois passed anint32tensor, thenTis automatically set toint32(well, actuallyDT_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:pythondef 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`. """
#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:
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.
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.
#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_KERNELDepending on the list of types you are registering the kernel for, you may be
able to use a macro provided bytensorflow/core/framework/register_types.h:
#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_KERNELList 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).
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).
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:
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:
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:
REGISTER_OP("MinLengthIntListExample")
.Attr("N: int >= 2")
.Input("in: N * int32")
.Output("out: int32");The same syntax works with "list(type)" attrs:
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:
REGISTER_OP("MultipleInsAndOuts")
.Input("y: int32")
.Input("z: float")
.Output("a: string")
.Output("b: int32");Each input or output spec is of the form:
<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 typetypeorlist(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
outis the same as
in the inputin, since both are of typeT.For a sequence of tensors with the same type:
<number> * <type>, where<number>is the name of an Attr with typeint. The<type>can
either be atf.DType,
or the name of an attr with typetype. As an example of the first, this
op accepts a list ofint32tensors: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 (likeTorN).
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, seetensorflow/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.
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:textREGISTER_OP("MyGeneralUnaryOp") .Input("in: float") .Output("out: float");you can make it polymorphic in a backwards-compatible way using:
textREGISTER_OP("MyGeneralUnaryOp") .Input("in: T") .Output("out: T") .Attr("T: numerictype = DT_FLOAT");You can safely make a constraint on an attr less restrictive. For example,
you can change from{int32, int64}to{int32, int64, float}ortype.
Or you may change from{"apple", "orange"}to{"apple", "banana", "orange"}orstring.You can change single inputs / outputs into list inputs / outputs, as long as
the default for the list type matches the old signature.You can add a new list input / output, if it defaults to empty.
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.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 intensorflow/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 theGraphDef 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 intensorflow/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 intensorflow/core/kernels/pad_op_gpu.cu.cc,
and the shared code is a templated class defined intensorflow/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.:
#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. Thetf_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:
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 thetf.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
nvcccompiling step are needed. Add-D_MWAITXINTRIN_H_INCLUDEDto thenvcccommand line to avoid errors frommwaitxintrin.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:
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 inputDetails about registering gradient functions withtf.RegisterGradient:
For an op with one output, the gradient function will take an
tf.Operationopand atf.Tensorgradand build new ops
out of the tensorsop.inputs[i],op.outputs[i], andgrad. Information
about any attrs can be found viatf.Operation.get_attr.If the op has multiple outputs, the gradient function will take
opandgrads, wheregradsis a list of gradients with respect to each output.
The result of the gradient function must be a list ofTensorobjects
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 beNone. For example, for an op taking a floating point tensorxand an
integer indexi, the gradient function wouldreturn [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 useops.NotDifferentiableto 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 theshape_inference::InferenceContext class. For example, in the shape function
for ZeroOut:
.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:
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 ofZeroOut with a vector shape constraint, the shape function
would be as follows:
.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, andMerge 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).
.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 andInferenceContext::WithValue; you can specify that an output dimension is the
sum / product of two input dimensions using InferenceContext::Add andInferenceContext::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, ...)
.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 toparameter_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. Thetf.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:
- Extract: Read data from persistent storage -- either local (e.g. HDD or
SSD) or remote (e.g. GCS or
HDFS). - 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. - 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 totf.estimator.Estimator.train. In code, this might look like the following
(naive, sequential) implementation:
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 datasetThe 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 thetf.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:
dataset = dataset.batch(batch_size=FLAGS.batch_size)
return datasetto:
dataset = dataset.batch(batch_size=FLAGS.batch_size)
dataset = dataset.prefetch(buffer_size=FLAGS.prefetch_buffer_size)
return datasetNote 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 thenum_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:
dataset = dataset.map(map_func=parse_fn)to:
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:
dataset = dataset.map(map_func=parse_fn, num_parallel_calls=FLAGS.num_parallel_calls)
dataset = dataset.batch(batch_size=FLAGS.batch_size)to:
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:
dataset = files.interleave(tf.data.TFRecordDataset)to:
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. (Seetf.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 thebatch 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 fusedtf.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
prefetchtransformation 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
maptransformation by setting thenum_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 fusedmap_and_batchtransformation;
especially if you are using large batch sizes. - If you are working with data stored remotely and / or requiring
deserialization, we recommend using theparallel_interleave
transformation to overlap the reading (and deserialization) of data from
different files. - Vectorize cheap user-defined functions passed in to the
maptransformation
to amortize the overhead associated with scheduling and executing the
function. - If your data can fit into memory, use the
cachetransformation 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 theinterleave,prefetch, andshufflefirst (if possible) to
reduce memory usage. - We recommend applying the
shuffletransformation before therepeat
transformation, ideally using the fusedshuffle_and_repeattransformation.