Hugging Face's logo

Buckets:
hf-doc-build
/
doc-dev

Files
xet
hf-doc-build/doc-dev / transformers /pr_36839 /en /perf_train_tpu_tf.html
rtrm's picture
rtrm
download
raw
68 kB
<meta charset="utf-8" /><meta name="hf:doc:metadata" content="{&quot;title&quot;:&quot;TPU&quot;,&quot;local&quot;:&quot;tpu&quot;,&quot;sections&quot;:[{&quot;title&quot;:&quot;Large datasets&quot;,&quot;local&quot;:&quot;large-datasets&quot;,&quot;sections&quot;:[{&quot;title&quot;:&quot;tf.TFRecord&quot;,&quot;local&quot;:&quot;tftfrecord&quot;,&quot;sections&quot;:[],&quot;depth&quot;:3},{&quot;title&quot;:&quot;Stream from raw data&quot;,&quot;local&quot;:&quot;stream-from-raw-data&quot;,&quot;sections&quot;:[],&quot;depth&quot;:3},{&quot;title&quot;:&quot;Stream with prepare_tf_dataset&quot;,&quot;local&quot;:&quot;stream-with-preparetfdataset&quot;,&quot;sections&quot;:[],&quot;depth&quot;:3}],&quot;depth&quot;:2},{&quot;title&quot;:&quot;TPU types&quot;,&quot;local&quot;:&quot;tpu-types&quot;,&quot;sections&quot;:[],&quot;depth&quot;:2},{&quot;title&quot;:&quot;XLA&quot;,&quot;local&quot;:&quot;xla&quot;,&quot;sections&quot;:[{&quot;title&quot;:&quot;Data dependent conditionals&quot;,&quot;local&quot;:&quot;data-dependent-conditionals&quot;,&quot;sections&quot;:[],&quot;depth&quot;:3},{&quot;title&quot;:&quot;Data dependent shapes&quot;,&quot;local&quot;:&quot;data-dependent-shapes&quot;,&quot;sections&quot;:[],&quot;depth&quot;:3},{&quot;title&quot;:&quot;Recompile different input shapes&quot;,&quot;local&quot;:&quot;recompile-different-input-shapes&quot;,&quot;sections&quot;:[],&quot;depth&quot;:3}],&quot;depth&quot;:2}],&quot;depth&quot;:1}">
 <link href="/docs/transformers/pr_36839/en/_app/immutable/assets/0.e3b0c442.css" rel="modulepreload">
 <link rel="modulepreload" href="/docs/transformers/pr_36839/en/_app/immutable/entry/start.6be8d590.js">
 <link rel="modulepreload" href="/docs/transformers/pr_36839/en/_app/immutable/chunks/scheduler.01eeda35.js">
 <link rel="modulepreload" href="/docs/transformers/pr_36839/en/_app/immutable/chunks/singletons.177df05e.js">
 <link rel="modulepreload" href="/docs/transformers/pr_36839/en/_app/immutable/chunks/index.4862150a.js">
 <link rel="modulepreload" href="/docs/transformers/pr_36839/en/_app/immutable/chunks/paths.517376d1.js">
 <link rel="modulepreload" href="/docs/transformers/pr_36839/en/_app/immutable/entry/app.09748b4b.js">
 <link rel="modulepreload" href="/docs/transformers/pr_36839/en/_app/immutable/chunks/index.6dd51b66.js">
 <link rel="modulepreload" href="/docs/transformers/pr_36839/en/_app/immutable/nodes/0.8897c14d.js">
 <link rel="modulepreload" href="/docs/transformers/pr_36839/en/_app/immutable/chunks/each.e59479a4.js">
 <link rel="modulepreload" href="/docs/transformers/pr_36839/en/_app/immutable/nodes/423.92f95dfd.js">
 <link rel="modulepreload" href="/docs/transformers/pr_36839/en/_app/immutable/chunks/Tip.de9bae2b.js">
 <link rel="modulepreload" href="/docs/transformers/pr_36839/en/_app/immutable/chunks/CodeBlock.864da1b0.js">
 <link rel="modulepreload" href="/docs/transformers/pr_36839/en/_app/immutable/chunks/EditOnGithub.7faefd25.js"><!-- HEAD_svelte-u9bgzb_START --><meta name="hf:doc:metadata" content="{&quot;title&quot;:&quot;TPU&quot;,&quot;local&quot;:&quot;tpu&quot;,&quot;sections&quot;:[{&quot;title&quot;:&quot;Large datasets&quot;,&quot;local&quot;:&quot;large-datasets&quot;,&quot;sections&quot;:[{&quot;title&quot;:&quot;tf.TFRecord&quot;,&quot;local&quot;:&quot;tftfrecord&quot;,&quot;sections&quot;:[],&quot;depth&quot;:3},{&quot;title&quot;:&quot;Stream from raw data&quot;,&quot;local&quot;:&quot;stream-from-raw-data&quot;,&quot;sections&quot;:[],&quot;depth&quot;:3},{&quot;title&quot;:&quot;Stream with prepare_tf_dataset&quot;,&quot;local&quot;:&quot;stream-with-preparetfdataset&quot;,&quot;sections&quot;:[],&quot;depth&quot;:3}],&quot;depth&quot;:2},{&quot;title&quot;:&quot;TPU types&quot;,&quot;local&quot;:&quot;tpu-types&quot;,&quot;sections&quot;:[],&quot;depth&quot;:2},{&quot;title&quot;:&quot;XLA&quot;,&quot;local&quot;:&quot;xla&quot;,&quot;sections&quot;:[{&quot;title&quot;:&quot;Data dependent conditionals&quot;,&quot;local&quot;:&quot;data-dependent-conditionals&quot;,&quot;sections&quot;:[],&quot;depth&quot;:3},{&quot;title&quot;:&quot;Data dependent shapes&quot;,&quot;local&quot;:&quot;data-dependent-shapes&quot;,&quot;sections&quot;:[],&quot;depth&quot;:3},{&quot;title&quot;:&quot;Recompile different input shapes&quot;,&quot;local&quot;:&quot;recompile-different-input-shapes&quot;,&quot;sections&quot;:[],&quot;depth&quot;:3}],&quot;depth&quot;:2}],&quot;depth&quot;:1}"><!-- HEAD_svelte-u9bgzb_END --> <p></p> <h1 class="relative group"><a id="tpu" class="header-link block pr-1.5 text-lg no-hover:hidden with-hover:absolute with-hover:p-1.5 with-hover:opacity-0 with-hover:group-hover:opacity-100 with-hover:right-full" href="#tpu"><span><svg class="" xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" aria-hidden="true" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 256 256"><path d="M167.594 88.393a8.001 8.001 0 0 1 0 11.314l-67.882 67.882a8 8 0 1 1-11.314-11.315l67.882-67.881a8.003 8.003 0 0 1 11.314 0zm-28.287 84.86l-28.284 28.284a40 40 0 0 1-56.567-56.567l28.284-28.284a8 8 0 0 0-11.315-11.315l-28.284 28.284a56 56 0 0 0 79.196 79.197l28.285-28.285a8 8 0 1 0-11.315-11.314zM212.852 43.14a56.002 56.002 0 0 0-79.196 0l-28.284 28.284a8 8 0 1 0 11.314 11.314l28.284-28.284a40 40 0 0 1 56.568 56.567l-28.285 28.285a8 8 0 0 0 11.315 11.314l28.284-28.284a56.065 56.065 0 0 0 0-79.196z" fill="currentColor"></path></svg></span></a> <span>TPU</span></h1> <p data-svelte-h="svelte-wejst4">TPU (Tensor Processing Unit) is a type of hardware designed to accelerate tensor computations for training and inference. TPUs are generally accessed through Google cloud services, but smaller TPUs are also available for free from <a href="https://colab.research.google.com/notebooks/tpu.ipynb" rel="nofollow">Google Colab</a> or <a href="https://www.kaggle.com/docs/tpu" rel="nofollow">Kaggle</a>.</p> <p data-svelte-h="svelte-26fyhs">This guide focuses on training a Keras model for sequence classification on a TPU from Google Colab. Make sure the TPU runtime is enabled by going to <strong>Runtime &gt; Change runtime type</strong> and selecting a TPU.</p> <p data-svelte-h="svelte-ggbuxf">Run the command below to install the latest version of Transformers and <a href="https://huggingface.co/docs/datasets" rel="nofollow">Datasets</a>.</p> <div class="code-block relative"><div class="absolute top-2.5 right-4"><button class="inline-flex items-center relative text-sm focus:text-green-500 cursor-pointer focus:outline-none transition duration-200 ease-in-out opacity-0 mx-0.5 text-gray-600 " title="code excerpt" type="button"><svg class="" xmlns="http://www.w3.org/2000/svg" aria-hidden="true" fill="currentColor" focusable="false" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 32 32"><path d="M28,10V28H10V10H28m0-2H10a2,2,0,0,0-2,2V28a2,2,0,0,0,2,2H28a2,2,0,0,0,2-2V10a2,2,0,0,0-2-2Z" transform="translate(0)"></path><path d="M4,18H2V4A2,2,0,0,1,4,2H18V4H4Z" transform="translate(0)"></path><rect fill="none" width="32" height="32"></rect></svg> <div class="absolute pointer-events-none transition-opacity bg-black text-white py-1 px-2 leading-tight rounded font-normal shadow left-1/2 top-full transform -translate-x-1/2 translate-y-2 opacity-0"><div class="absolute bottom-full left-1/2 transform -translate-x-1/2 w-0 h-0 border-black border-4 border-t-0" style="border-left-color: transparent; border-right-color: transparent; "></div> Copied</div></button></div> <pre class=""><!-- HTML_TAG_START -->!pip install --U transformers datasets<!-- HTML_TAG_END --></pre></div> <p data-svelte-h="svelte-2kj0ub">Create an instance of <a href="https://www.tensorflow.org/api_docs/python/tf/distribute/cluster_resolver/TPUClusterResolver" rel="nofollow">tf.distribute.cluster_resolver.TPUClusterResolver</a>, and then connect to the remote cluster and initialize the TPUs.</p> <div class="code-block relative"><div class="absolute top-2.5 right-4"><button class="inline-flex items-center relative text-sm focus:text-green-500 cursor-pointer focus:outline-none transition duration-200 ease-in-out opacity-0 mx-0.5 text-gray-600 " title="code excerpt" type="button"><svg class="" xmlns="http://www.w3.org/2000/svg" aria-hidden="true" fill="currentColor" focusable="false" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 32 32"><path d="M28,10V28H10V10H28m0-2H10a2,2,0,0,0-2,2V28a2,2,0,0,0,2,2H28a2,2,0,0,0,2-2V10a2,2,0,0,0-2-2Z" transform="translate(0)"></path><path d="M4,18H2V4A2,2,0,0,1,4,2H18V4H4Z" transform="translate(0)"></path><rect fill="none" width="32" height="32"></rect></svg> <div class="absolute pointer-events-none transition-opacity bg-black text-white py-1 px-2 leading-tight rounded font-normal shadow left-1/2 top-full transform -translate-x-1/2 translate-y-2 opacity-0"><div class="absolute bottom-full left-1/2 transform -translate-x-1/2 w-0 h-0 border-black border-4 border-t-0" style="border-left-color: transparent; border-right-color: transparent; "></div> Copied</div></button></div> <pre class=""><!-- HTML_TAG_START --><span class="hljs-keyword">import</span> tensorflow <span class="hljs-keyword">as</span> tf
resolver = tf.distribute.cluster_resolver.TPUClusterResolver()
tf.config.experimental_connect_to_cluster(resolver)
tf.tpu.experimental.initialize_tpu_system(resolver)<!-- HTML_TAG_END --></pre></div> <p data-svelte-h="svelte-i4fn8y">There are various distribution strategies for running your model on multiple TPUs. The <a href="https://www.tensorflow.org/api_docs/python/tf/distribute/TPUStrategy" rel="nofollow">tpu.distribute.TPUStrategy</a> offers synchronized distributed training.</p> <div class="code-block relative"><div class="absolute top-2.5 right-4"><button class="inline-flex items-center relative text-sm focus:text-green-500 cursor-pointer focus:outline-none transition duration-200 ease-in-out opacity-0 mx-0.5 text-gray-600 " title="code excerpt" type="button"><svg class="" xmlns="http://www.w3.org/2000/svg" aria-hidden="true" fill="currentColor" focusable="false" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 32 32"><path d="M28,10V28H10V10H28m0-2H10a2,2,0,0,0-2,2V28a2,2,0,0,0,2,2H28a2,2,0,0,0,2-2V10a2,2,0,0,0-2-2Z" transform="translate(0)"></path><path d="M4,18H2V4A2,2,0,0,1,4,2H18V4H4Z" transform="translate(0)"></path><rect fill="none" width="32" height="32"></rect></svg> <div class="absolute pointer-events-none transition-opacity bg-black text-white py-1 px-2 leading-tight rounded font-normal shadow left-1/2 top-full transform -translate-x-1/2 translate-y-2 opacity-0"><div class="absolute bottom-full left-1/2 transform -translate-x-1/2 w-0 h-0 border-black border-4 border-t-0" style="border-left-color: transparent; border-right-color: transparent; "></div> Copied</div></button></div> <pre class=""><!-- HTML_TAG_START -->strategy = tf.distribute.TPUStrategy(resolver)<!-- HTML_TAG_END --></pre></div> <p data-svelte-h="svelte-1ntv7za">Load and tokenize a dataset - this example uses <a href="https://huggingface.co/datasets/nyu-mll/glue/viewer/cola" rel="nofollow">CoLA</a> from the GLUE benchmark - and pad all samples to the maximum length so it is easier to load as an array and to avoid <a href="#xla">XLA compilation issues</a>.</p> <div class="code-block relative"><div class="absolute top-2.5 right-4"><button class="inline-flex items-center relative text-sm focus:text-green-500 cursor-pointer focus:outline-none transition duration-200 ease-in-out opacity-0 mx-0.5 text-gray-600 " title="code excerpt" type="button"><svg class="" xmlns="http://www.w3.org/2000/svg" aria-hidden="true" fill="currentColor" focusable="false" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 32 32"><path d="M28,10V28H10V10H28m0-2H10a2,2,0,0,0-2,2V28a2,2,0,0,0,2,2H28a2,2,0,0,0,2-2V10a2,2,0,0,0-2-2Z" transform="translate(0)"></path><path d="M4,18H2V4A2,2,0,0,1,4,2H18V4H4Z" transform="translate(0)"></path><rect fill="none" width="32" height="32"></rect></svg> <div class="absolute pointer-events-none transition-opacity bg-black text-white py-1 px-2 leading-tight rounded font-normal shadow left-1/2 top-full transform -translate-x-1/2 translate-y-2 opacity-0"><div class="absolute bottom-full left-1/2 transform -translate-x-1/2 w-0 h-0 border-black border-4 border-t-0" style="border-left-color: transparent; border-right-color: transparent; "></div> Copied</div></button></div> <pre class=""><!-- HTML_TAG_START --><span class="hljs-keyword">from</span> transformers <span class="hljs-keyword">import</span> AutoTokenizer
<span class="hljs-keyword">from</span> datasets <span class="hljs-keyword">import</span> load_dataset
<span class="hljs-keyword">import</span> numpy <span class="hljs-keyword">as</span> np
dataset = load_dataset(<span class="hljs-string">&quot;glue&quot;</span>, <span class="hljs-string">&quot;cola&quot;</span>)[<span class="hljs-string">&quot;train&quot;</span>]
tokenizer = AutoTokenizer.from_pretrained(<span class="hljs-string">&quot;distilbert-base-cased&quot;</span>)
train_data = tokenizer(
 dataset[<span class="hljs-string">&quot;sentence&quot;</span>],
 padding=<span class="hljs-string">&quot;max_length&quot;</span>,
 truncation=<span class="hljs-literal">True</span>,
 max_length=<span class="hljs-number">128</span>,
 return_tensors=<span class="hljs-string">&quot;np&quot;</span>,
)
train_data = <span class="hljs-built_in">dict</span>(train_data)
train_labels = np.array(dataset[<span class="hljs-string">&quot;label&quot;</span>])<!-- HTML_TAG_END --></pre></div> <p data-svelte-h="svelte-6ywacu">The model <strong>must</strong> be created inside <a href="https://www.tensorflow.org/api_docs/python/tf/distribute/MirroredStrategy#scope" rel="nofollow">Strategy.scope</a> in order to replicate the model layers on each TPU device.</p> <div class="code-block relative"><div class="absolute top-2.5 right-4"><button class="inline-flex items-center relative text-sm focus:text-green-500 cursor-pointer focus:outline-none transition duration-200 ease-in-out opacity-0 mx-0.5 text-gray-600 " title="code excerpt" type="button"><svg class="" xmlns="http://www.w3.org/2000/svg" aria-hidden="true" fill="currentColor" focusable="false" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 32 32"><path d="M28,10V28H10V10H28m0-2H10a2,2,0,0,0-2,2V28a2,2,0,0,0,2,2H28a2,2,0,0,0,2-2V10a2,2,0,0,0-2-2Z" transform="translate(0)"></path><path d="M4,18H2V4A2,2,0,0,1,4,2H18V4H4Z" transform="translate(0)"></path><rect fill="none" width="32" height="32"></rect></svg> <div class="absolute pointer-events-none transition-opacity bg-black text-white py-1 px-2 leading-tight rounded font-normal shadow left-1/2 top-full transform -translate-x-1/2 translate-y-2 opacity-0"><div class="absolute bottom-full left-1/2 transform -translate-x-1/2 w-0 h-0 border-black border-4 border-t-0" style="border-left-color: transparent; border-right-color: transparent; "></div> Copied</div></button></div> <pre class=""><!-- HTML_TAG_START --><span class="hljs-keyword">from</span> transformers <span class="hljs-keyword">import</span> TFAutoModelForSequenceClassification
<span class="hljs-keyword">with</span> strategy.scope():
 model = TFAutoModelForSequenceClassification.from_pretrained(model_checkpoint)
 model.<span class="hljs-built_in">compile</span>(optimizer=<span class="hljs-string">&quot;adam&quot;</span>)<!-- HTML_TAG_END --></pre></div> <p data-svelte-h="svelte-wfm0b8">TPUs only accept <a href="https://www.tensorflow.org/api_docs/python/tf/data/Dataset" rel="nofollow">tf.data.Dataset</a> inputs unlike the Keras <a href="https://keras.io/api/models/model_training_apis/#fit-method" rel="nofollow">fit</a> method which accepts a broader range of inputs.</p> <div class="code-block relative"><div class="absolute top-2.5 right-4"><button class="inline-flex items-center relative text-sm focus:text-green-500 cursor-pointer focus:outline-none transition duration-200 ease-in-out opacity-0 mx-0.5 text-gray-600 " title="code excerpt" type="button"><svg class="" xmlns="http://www.w3.org/2000/svg" aria-hidden="true" fill="currentColor" focusable="false" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 32 32"><path d="M28,10V28H10V10H28m0-2H10a2,2,0,0,0-2,2V28a2,2,0,0,0,2,2H28a2,2,0,0,0,2-2V10a2,2,0,0,0-2-2Z" transform="translate(0)"></path><path d="M4,18H2V4A2,2,0,0,1,4,2H18V4H4Z" transform="translate(0)"></path><rect fill="none" width="32" height="32"></rect></svg> <div class="absolute pointer-events-none transition-opacity bg-black text-white py-1 px-2 leading-tight rounded font-normal shadow left-1/2 top-full transform -translate-x-1/2 translate-y-2 opacity-0"><div class="absolute bottom-full left-1/2 transform -translate-x-1/2 w-0 h-0 border-black border-4 border-t-0" style="border-left-color: transparent; border-right-color: transparent; "></div> Copied</div></button></div> <pre class=""><!-- HTML_TAG_START -->BATCH_SIZE = <span class="hljs-number">8</span> * strategy.num_replicas_in_sync
tf_dataset = tf.data.Dataset.from_tensor_slices((train_data, train_labels))
tf_dataset = tf_dataset.shuffle(<span class="hljs-built_in">len</span>(tf_dataset))
tf_dataset = tf_dataset.batch(BATCH_SIZE, drop_remainder=<span class="hljs-literal">True</span>)<!-- HTML_TAG_END --></pre></div> <p data-svelte-h="svelte-1nrwj5f">Finally, call <a href="https://keras.io/api/models/model_training_apis/#fit-method" rel="nofollow">fit</a> to start training.</p> <div class="code-block relative"><div class="absolute top-2.5 right-4"><button class="inline-flex items-center relative text-sm focus:text-green-500 cursor-pointer focus:outline-none transition duration-200 ease-in-out opacity-0 mx-0.5 text-gray-600 " title="code excerpt" type="button"><svg class="" xmlns="http://www.w3.org/2000/svg" aria-hidden="true" fill="currentColor" focusable="false" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 32 32"><path d="M28,10V28H10V10H28m0-2H10a2,2,0,0,0-2,2V28a2,2,0,0,0,2,2H28a2,2,0,0,0,2-2V10a2,2,0,0,0-2-2Z" transform="translate(0)"></path><path d="M4,18H2V4A2,2,0,0,1,4,2H18V4H4Z" transform="translate(0)"></path><rect fill="none" width="32" height="32"></rect></svg> <div class="absolute pointer-events-none transition-opacity bg-black text-white py-1 px-2 leading-tight rounded font-normal shadow left-1/2 top-full transform -translate-x-1/2 translate-y-2 opacity-0"><div class="absolute bottom-full left-1/2 transform -translate-x-1/2 w-0 h-0 border-black border-4 border-t-0" style="border-left-color: transparent; border-right-color: transparent; "></div> Copied</div></button></div> <pre class=""><!-- HTML_TAG_START -->model.fit(tf_dataset)<!-- HTML_TAG_END --></pre></div> <h2 class="relative group"><a id="large-datasets" class="header-link block pr-1.5 text-lg no-hover:hidden with-hover:absolute with-hover:p-1.5 with-hover:opacity-0 with-hover:group-hover:opacity-100 with-hover:right-full" href="#large-datasets"><span><svg class="" xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" aria-hidden="true" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 256 256"><path d="M167.594 88.393a8.001 8.001 0 0 1 0 11.314l-67.882 67.882a8 8 0 1 1-11.314-11.315l67.882-67.881a8.003 8.003 0 0 1 11.314 0zm-28.287 84.86l-28.284 28.284a40 40 0 0 1-56.567-56.567l28.284-28.284a8 8 0 0 0-11.315-11.315l-28.284 28.284a56 56 0 0 0 79.196 79.197l28.285-28.285a8 8 0 1 0-11.315-11.314zM212.852 43.14a56.002 56.002 0 0 0-79.196 0l-28.284 28.284a8 8 0 1 0 11.314 11.314l28.284-28.284a40 40 0 0 1 56.568 56.567l-28.285 28.285a8 8 0 0 0 11.315 11.314l28.284-28.284a56.065 56.065 0 0 0 0-79.196z" fill="currentColor"></path></svg></span></a> <span>Large datasets</span></h2> <p data-svelte-h="svelte-1rmkxi1">The dataset created above pads every sample to the maximum length and loads the whole dataset into memory. This may not be possible if you’re working with larger datasets. When training on large datasets, you may want to create a <a href="https://www.tensorflow.org/tutorials/load_data/tfrecord" rel="nofollow">tf.TFRecord</a> or stream the data.</p> <h3 class="relative group"><a id="tftfrecord" class="header-link block pr-1.5 text-lg no-hover:hidden with-hover:absolute with-hover:p-1.5 with-hover:opacity-0 with-hover:group-hover:opacity-100 with-hover:right-full" href="#tftfrecord"><span><svg class="" xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" aria-hidden="true" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 256 256"><path d="M167.594 88.393a8.001 8.001 0 0 1 0 11.314l-67.882 67.882a8 8 0 1 1-11.314-11.315l67.882-67.881a8.003 8.003 0 0 1 11.314 0zm-28.287 84.86l-28.284 28.284a40 40 0 0 1-56.567-56.567l28.284-28.284a8 8 0 0 0-11.315-11.315l-28.284 28.284a56 56 0 0 0 79.196 79.197l28.285-28.285a8 8 0 1 0-11.315-11.314zM212.852 43.14a56.002 56.002 0 0 0-79.196 0l-28.284 28.284a8 8 0 1 0 11.314 11.314l28.284-28.284a40 40 0 0 1 56.568 56.567l-28.285 28.285a8 8 0 0 0 11.315 11.314l28.284-28.284a56.065 56.065 0 0 0 0-79.196z" fill="currentColor"></path></svg></span></a> <span>tf.TFRecord</span></h3> <p data-svelte-h="svelte-1thxjgi"><a href="https://www.tensorflow.org/tutorials/load_data/tfrecord" rel="nofollow">tf.TFRecord</a> is the standard <a href="https://www.tensorflow.org/guide/data" rel="nofollow">tf.data</a> format for storing training data. For very large training jobs, it’s worth preprocessing your data and storing it in the <code>tf.TFRecord</code> format and building a <code>tf.data</code> pipeline on top. Refer to the table below to help you decide whether <code>tf.TFRecord</code> is helpful for you.</p> <table data-svelte-h="svelte-1bjnyhm"><thead><tr><th>pros</th> <th>cons</th></tr></thead> <tbody><tr><td>works on all TPU instances</td> <td>costs associated with cloud storage</td></tr> <tr><td>supports huge datasets and massive throughput</td> <td>some data types (images) can take a lot of space to store</td></tr> <tr><td>suitable for training on entire TPU pods</td> <td></td></tr> <tr><td>preprocessing is done in advance, maximizing training speed</td> <td></td></tr></tbody></table> <p data-svelte-h="svelte-1pkbze1">Preprocess and tokenize the dataset before writing it to a <code>tf.TFRecord</code> to avoid writing every time the data is loaded.</p> <p data-svelte-h="svelte-1yh9lje">An exception is made for <em>train-time augmentations</em>, because augmentations applied after writing to a <code>tf.TFRecord</code> results in the same augmentation for each epoch. Instead, apply augmentations in the <code>tf.data</code> pipeline that loads the data.</p> <div class="course-tip bg-gradient-to-br dark:bg-gradient-to-r before:border-green-500 dark:before:border-green-800 from-green-50 dark:from-gray-900 to-white dark:to-gray-950 border border-green-50 text-green-700 dark:text-gray-400"><p data-svelte-h="svelte-1njs9h7">In practice, you probably won’t be able to load the entire dataset in memory. Load a chunk of the dataset at a time and convert it to <code>TFRecord</code>, and repeat until the entire dataset is in the <code>TFRecord</code> format. Then you can use a list of all the files to create a <code>TFRecordDataset</code>. The example below demonstrates a single file for simplicity.</p></div> <div class="code-block relative"><div class="absolute top-2.5 right-4"><button class="inline-flex items-center relative text-sm focus:text-green-500 cursor-pointer focus:outline-none transition duration-200 ease-in-out opacity-0 mx-0.5 text-gray-600 " title="code excerpt" type="button"><svg class="" xmlns="http://www.w3.org/2000/svg" aria-hidden="true" fill="currentColor" focusable="false" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 32 32"><path d="M28,10V28H10V10H28m0-2H10a2,2,0,0,0-2,2V28a2,2,0,0,0,2,2H28a2,2,0,0,0,2-2V10a2,2,0,0,0-2-2Z" transform="translate(0)"></path><path d="M4,18H2V4A2,2,0,0,1,4,2H18V4H4Z" transform="translate(0)"></path><rect fill="none" width="32" height="32"></rect></svg> <div class="absolute pointer-events-none transition-opacity bg-black text-white py-1 px-2 leading-tight rounded font-normal shadow left-1/2 top-full transform -translate-x-1/2 translate-y-2 opacity-0"><div class="absolute bottom-full left-1/2 transform -translate-x-1/2 w-0 h-0 border-black border-4 border-t-0" style="border-left-color: transparent; border-right-color: transparent; "></div> Copied</div></button></div> <pre class=""><!-- HTML_TAG_START -->tokenized_data = tokenizer(
 dataset[<span class="hljs-string">&quot;sentence&quot;</span>],
 padding=<span class="hljs-string">&quot;max_length&quot;</span>,
 truncation=<span class="hljs-literal">True</span>,
 max_length=<span class="hljs-number">128</span>,
 return_tensors=<span class="hljs-string">&quot;np&quot;</span>,
)
labels = dataset[<span class="hljs-string">&quot;label&quot;</span>]
<span class="hljs-keyword">with</span> tf.io.TFRecordWriter(<span class="hljs-string">&quot;dataset.tfrecords&quot;</span>) <span class="hljs-keyword">as</span> file_writer:
 <span class="hljs-keyword">for</span> i <span class="hljs-keyword">in</span> <span class="hljs-built_in">range</span>(<span class="hljs-built_in">len</span>(labels)):
 features = {
 <span class="hljs-string">&quot;input_ids&quot;</span>: tf.train.Feature(
 int64_list=tf.train.Int64List(value=tokenized_data[<span class="hljs-string">&quot;input_ids&quot;</span>][i])
 ),
 <span class="hljs-string">&quot;attention_mask&quot;</span>: tf.train.Feature(
 int64_list=tf.train.Int64List(value=tokenized_data[<span class="hljs-string">&quot;attention_mask&quot;</span>][i])
 ),
 <span class="hljs-string">&quot;labels&quot;</span>: tf.train.Feature(
 int64_list=tf.train.Int64List(value=[labels[i]])
 ),
 }
 features = tf.train.Features(feature=features)
 example = tf.train.Example(features=features)
 record_bytes = example.SerializeToString()
 file_writer.write(record_bytes)<!-- HTML_TAG_END --></pre></div> <p data-svelte-h="svelte-1kezdht">Build a <a href="https://www.tensorflow.org/api_docs/python/tf/data/TFRecordDataset" rel="nofollow">TFRecordDataset</a> using the saved filename to load it.</p> <div class="code-block relative"><div class="absolute top-2.5 right-4"><button class="inline-flex items-center relative text-sm focus:text-green-500 cursor-pointer focus:outline-none transition duration-200 ease-in-out opacity-0 mx-0.5 text-gray-600 " title="code excerpt" type="button"><svg class="" xmlns="http://www.w3.org/2000/svg" aria-hidden="true" fill="currentColor" focusable="false" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 32 32"><path d="M28,10V28H10V10H28m0-2H10a2,2,0,0,0-2,2V28a2,2,0,0,0,2,2H28a2,2,0,0,0,2-2V10a2,2,0,0,0-2-2Z" transform="translate(0)"></path><path d="M4,18H2V4A2,2,0,0,1,4,2H18V4H4Z" transform="translate(0)"></path><rect fill="none" width="32" height="32"></rect></svg> <div class="absolute pointer-events-none transition-opacity bg-black text-white py-1 px-2 leading-tight rounded font-normal shadow left-1/2 top-full transform -translate-x-1/2 translate-y-2 opacity-0"><div class="absolute bottom-full left-1/2 transform -translate-x-1/2 w-0 h-0 border-black border-4 border-t-0" style="border-left-color: transparent; border-right-color: transparent; "></div> Copied</div></button></div> <pre class=""><!-- HTML_TAG_START --><span class="hljs-keyword">def</span> <span class="hljs-title function_">decode_fn</span>(<span class="hljs-params">sample</span>):
 features = {
 <span class="hljs-string">&quot;input_ids&quot;</span>: tf.io.FixedLenFeature((<span class="hljs-number">128</span>,), dtype=tf.int64),
 <span class="hljs-string">&quot;attention_mask&quot;</span>: tf.io.FixedLenFeature((<span class="hljs-number">128</span>,), dtype=tf.int64),
 <span class="hljs-string">&quot;labels&quot;</span>: tf.io.FixedLenFeature((<span class="hljs-number">1</span>,), dtype=tf.int64),
 }
 <span class="hljs-keyword">return</span> tf.io.parse_example(sample, features)
<span class="hljs-comment"># TFRecordDataset can handle gs:// paths</span>
tf_dataset = tf.data.TFRecordDataset([<span class="hljs-string">&quot;gs://matt-tf-tpu-tutorial-datasets/cola/dataset.tfrecords&quot;</span>])
tf_dataset = tf_dataset.<span class="hljs-built_in">map</span>(decode_fn)
tf_dataset = tf_dataset.shuffle(<span class="hljs-built_in">len</span>(dataset)).batch(BATCH_SIZE, drop_remainder=<span class="hljs-literal">True</span>)
tf_dataset = tf_dataset.apply(
 tf.data.experimental.assert_cardinality(<span class="hljs-built_in">len</span>(labels) // BATCH_SIZE)
)<!-- HTML_TAG_END --></pre></div> <p data-svelte-h="svelte-1s9k8gm">The dataset can now be passed to the <a href="https://keras.io/api/models/model_training_apis/#fit-method" rel="nofollow">fit</a> method.</p> <div class="code-block relative"><div class="absolute top-2.5 right-4"><button class="inline-flex items-center relative text-sm focus:text-green-500 cursor-pointer focus:outline-none transition duration-200 ease-in-out opacity-0 mx-0.5 text-gray-600 " title="code excerpt" type="button"><svg class="" xmlns="http://www.w3.org/2000/svg" aria-hidden="true" fill="currentColor" focusable="false" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 32 32"><path d="M28,10V28H10V10H28m0-2H10a2,2,0,0,0-2,2V28a2,2,0,0,0,2,2H28a2,2,0,0,0,2-2V10a2,2,0,0,0-2-2Z" transform="translate(0)"></path><path d="M4,18H2V4A2,2,0,0,1,4,2H18V4H4Z" transform="translate(0)"></path><rect fill="none" width="32" height="32"></rect></svg> <div class="absolute pointer-events-none transition-opacity bg-black text-white py-1 px-2 leading-tight rounded font-normal shadow left-1/2 top-full transform -translate-x-1/2 translate-y-2 opacity-0"><div class="absolute bottom-full left-1/2 transform -translate-x-1/2 w-0 h-0 border-black border-4 border-t-0" style="border-left-color: transparent; border-right-color: transparent; "></div> Copied</div></button></div> <pre class=""><!-- HTML_TAG_START -->model.fit(tf_dataset)<!-- HTML_TAG_END --></pre></div> <h3 class="relative group"><a id="stream-from-raw-data" class="header-link block pr-1.5 text-lg no-hover:hidden with-hover:absolute with-hover:p-1.5 with-hover:opacity-0 with-hover:group-hover:opacity-100 with-hover:right-full" href="#stream-from-raw-data"><span><svg class="" xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" aria-hidden="true" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 256 256"><path d="M167.594 88.393a8.001 8.001 0 0 1 0 11.314l-67.882 67.882a8 8 0 1 1-11.314-11.315l67.882-67.881a8.003 8.003 0 0 1 11.314 0zm-28.287 84.86l-28.284 28.284a40 40 0 0 1-56.567-56.567l28.284-28.284a8 8 0 0 0-11.315-11.315l-28.284 28.284a56 56 0 0 0 79.196 79.197l28.285-28.285a8 8 0 1 0-11.315-11.314zM212.852 43.14a56.002 56.002 0 0 0-79.196 0l-28.284 28.284a8 8 0 1 0 11.314 11.314l28.284-28.284a40 40 0 0 1 56.568 56.567l-28.285 28.285a8 8 0 0 0 11.315 11.314l28.284-28.284a56.065 56.065 0 0 0 0-79.196z" fill="currentColor"></path></svg></span></a> <span>Stream from raw data</span></h3> <p data-svelte-h="svelte-19b2fim">Data can be stored in its native format and preprocessed in a <a href="https://www.tensorflow.org/guide/data" rel="nofollow">tf.data</a> pipeline as the data is loaded. This approach isn’t supported for many models with complex tokenization schemes, but some models like BERT are supported because their tokenization can be compiled. Refer to the table below to help you decide whether this approach is helpful for you.</p> <table data-svelte-h="svelte-eplesd"><thead><tr><th>pros</th> <th>cons</th></tr></thead> <tbody><tr><td>suitable for highly compressed big data in native format (images, audio)</td> <td>requires writing a full preprocessing pipeline</td></tr> <tr><td>convenient if raw data is available in a public cloud bucket</td> <td>complex preprocessing on-the-fly can hurt throughput</td></tr> <tr><td>works on all TPU instances if data is stored in Google Cloud</td> <td>must place data in cloud storage if not already there</td></tr> <tr><td></td> <td>not as suitable for text data because writing a tokenization pipeline is hard (use <code>TFRecord</code> for text)</td></tr></tbody></table> <p data-svelte-h="svelte-1yanrvl">The example below demonstrates streaming data for an image model.</p> <p data-svelte-h="svelte-1p9l68o">Load an image dataset and get a list of the underlying image file paths and labels.</p> <div class="code-block relative"><div class="absolute top-2.5 right-4"><button class="inline-flex items-center relative text-sm focus:text-green-500 cursor-pointer focus:outline-none transition duration-200 ease-in-out opacity-0 mx-0.5 text-gray-600 " title="code excerpt" type="button"><svg class="" xmlns="http://www.w3.org/2000/svg" aria-hidden="true" fill="currentColor" focusable="false" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 32 32"><path d="M28,10V28H10V10H28m0-2H10a2,2,0,0,0-2,2V28a2,2,0,0,0,2,2H28a2,2,0,0,0,2-2V10a2,2,0,0,0-2-2Z" transform="translate(0)"></path><path d="M4,18H2V4A2,2,0,0,1,4,2H18V4H4Z" transform="translate(0)"></path><rect fill="none" width="32" height="32"></rect></svg> <div class="absolute pointer-events-none transition-opacity bg-black text-white py-1 px-2 leading-tight rounded font-normal shadow left-1/2 top-full transform -translate-x-1/2 translate-y-2 opacity-0"><div class="absolute bottom-full left-1/2 transform -translate-x-1/2 w-0 h-0 border-black border-4 border-t-0" style="border-left-color: transparent; border-right-color: transparent; "></div> Copied</div></button></div> <pre class=""><!-- HTML_TAG_START --><span class="hljs-keyword">from</span> datasets <span class="hljs-keyword">import</span> load_dataset
image_dataset = load_dataset(<span class="hljs-string">&quot;beans&quot;</span>, split=<span class="hljs-string">&quot;train&quot;</span>)
filenames = image_dataset[<span class="hljs-string">&quot;image_file_path&quot;</span>]
labels = image_dataset[<span class="hljs-string">&quot;labels&quot;</span>]<!-- HTML_TAG_END --></pre></div> <p data-svelte-h="svelte-2il5q0">Convert the local filenames in the dataset into <code>gs://</code> paths in Google Cloud Storage.</p> <div class="code-block relative"><div class="absolute top-2.5 right-4"><button class="inline-flex items-center relative text-sm focus:text-green-500 cursor-pointer focus:outline-none transition duration-200 ease-in-out opacity-0 mx-0.5 text-gray-600 " title="code excerpt" type="button"><svg class="" xmlns="http://www.w3.org/2000/svg" aria-hidden="true" fill="currentColor" focusable="false" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 32 32"><path d="M28,10V28H10V10H28m0-2H10a2,2,0,0,0-2,2V28a2,2,0,0,0,2,2H28a2,2,0,0,0,2-2V10a2,2,0,0,0-2-2Z" transform="translate(0)"></path><path d="M4,18H2V4A2,2,0,0,1,4,2H18V4H4Z" transform="translate(0)"></path><rect fill="none" width="32" height="32"></rect></svg> <div class="absolute pointer-events-none transition-opacity bg-black text-white py-1 px-2 leading-tight rounded font-normal shadow left-1/2 top-full transform -translate-x-1/2 translate-y-2 opacity-0"><div class="absolute bottom-full left-1/2 transform -translate-x-1/2 w-0 h-0 border-black border-4 border-t-0" style="border-left-color: transparent; border-right-color: transparent; "></div> Copied</div></button></div> <pre class=""><!-- HTML_TAG_START --><span class="hljs-comment"># strip everything but the category directory and filenames</span>
base_filenames = [<span class="hljs-string">&#x27;/&#x27;</span>.join(filename.split(<span class="hljs-string">&#x27;/&#x27;</span>)[-<span class="hljs-number">2</span>:]) <span class="hljs-keyword">for</span> filename <span class="hljs-keyword">in</span> filenames]
<span class="hljs-comment"># prepend the Google Cloud base path to everything instead</span>
gs_paths = [<span class="hljs-string">&quot;gs://matt-tf-tpu-tutorial-datasets/beans/&quot;</span>+filename <span class="hljs-keyword">for</span> filename <span class="hljs-keyword">in</span> base_filenames]
<span class="hljs-comment"># create tf_dataset</span>
tf_dataset = tf.data.Dataset.from_tensor_slices(
 {<span class="hljs-string">&quot;filename&quot;</span>: gs_paths, <span class="hljs-string">&quot;labels&quot;</span>: labels}
)
tf_dataset = tf_dataset.shuffle(<span class="hljs-built_in">len</span>(tf_dataset))<!-- HTML_TAG_END --></pre></div> <p data-svelte-h="svelte-pw7lrv">Transformers preprocessing classes like <a href="/docs/transformers/pr_36839/en/model_doc/auto#transformers.AutoImageProcessor">AutoImageProcessor</a> are framework-agnostic and can’t be compiled into a pipeline by <code>tf.data</code>. To get around this, get the normalization values (<code>mean</code> and <code>std</code>) from the <a href="/docs/transformers/pr_36839/en/model_doc/auto#transformers.AutoImageProcessor">AutoImageProcessor</a> and use them in the <code>tf.data</code> pipeline.</p> <div class="code-block relative"><div class="absolute top-2.5 right-4"><button class="inline-flex items-center relative text-sm focus:text-green-500 cursor-pointer focus:outline-none transition duration-200 ease-in-out opacity-0 mx-0.5 text-gray-600 " title="code excerpt" type="button"><svg class="" xmlns="http://www.w3.org/2000/svg" aria-hidden="true" fill="currentColor" focusable="false" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 32 32"><path d="M28,10V28H10V10H28m0-2H10a2,2,0,0,0-2,2V28a2,2,0,0,0,2,2H28a2,2,0,0,0,2-2V10a2,2,0,0,0-2-2Z" transform="translate(0)"></path><path d="M4,18H2V4A2,2,0,0,1,4,2H18V4H4Z" transform="translate(0)"></path><rect fill="none" width="32" height="32"></rect></svg> <div class="absolute pointer-events-none transition-opacity bg-black text-white py-1 px-2 leading-tight rounded font-normal shadow left-1/2 top-full transform -translate-x-1/2 translate-y-2 opacity-0"><div class="absolute bottom-full left-1/2 transform -translate-x-1/2 w-0 h-0 border-black border-4 border-t-0" style="border-left-color: transparent; border-right-color: transparent; "></div> Copied</div></button></div> <pre class=""><!-- HTML_TAG_START --><span class="hljs-keyword">from</span> transformers <span class="hljs-keyword">import</span> AutoImageProcessor
processor = AutoImageProcessor.from_pretrained(<span class="hljs-string">&quot;google/vit-base-patch16-224&quot;</span>)
image_size = (processor.size[<span class="hljs-string">&quot;height&quot;</span>], processor.size[<span class="hljs-string">&quot;width&quot;</span>])
image_mean = processor.image_mean
image_std = processor.image_std<!-- HTML_TAG_END --></pre></div> <p data-svelte-h="svelte-xhmqey">Use these normalization values to create a function to load and preprocess the images.</p> <div class="code-block relative"><div class="absolute top-2.5 right-4"><button class="inline-flex items-center relative text-sm focus:text-green-500 cursor-pointer focus:outline-none transition duration-200 ease-in-out opacity-0 mx-0.5 text-gray-600 " title="code excerpt" type="button"><svg class="" xmlns="http://www.w3.org/2000/svg" aria-hidden="true" fill="currentColor" focusable="false" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 32 32"><path d="M28,10V28H10V10H28m0-2H10a2,2,0,0,0-2,2V28a2,2,0,0,0,2,2H28a2,2,0,0,0,2-2V10a2,2,0,0,0-2-2Z" transform="translate(0)"></path><path d="M4,18H2V4A2,2,0,0,1,4,2H18V4H4Z" transform="translate(0)"></path><rect fill="none" width="32" height="32"></rect></svg> <div class="absolute pointer-events-none transition-opacity bg-black text-white py-1 px-2 leading-tight rounded font-normal shadow left-1/2 top-full transform -translate-x-1/2 translate-y-2 opacity-0"><div class="absolute bottom-full left-1/2 transform -translate-x-1/2 w-0 h-0 border-black border-4 border-t-0" style="border-left-color: transparent; border-right-color: transparent; "></div> Copied</div></button></div> <pre class=""><!-- HTML_TAG_START -->BATCH_SIZE = <span class="hljs-number">8</span> * strategy.num_replicas_in_sync
<span class="hljs-keyword">def</span> <span class="hljs-title function_">decode_fn</span>(<span class="hljs-params">sample</span>):
 image_data = tf.io.read_file(sample[<span class="hljs-string">&quot;filename&quot;</span>])
 image = tf.io.decode_jpeg(image_data, channels=<span class="hljs-number">3</span>)
 image = tf.image.resize(image, image_size)
 array = tf.cast(image, tf.float32)
 array /= <span class="hljs-number">255.0</span>
 array = (array - image_mean) / image_std
 array = tf.transpose(array, perm=[<span class="hljs-number">2</span>, <span class="hljs-number">0</span>, <span class="hljs-number">1</span>])
 <span class="hljs-keyword">return</span> {<span class="hljs-string">&quot;pixel_values&quot;</span>: array, <span class="hljs-string">&quot;labels&quot;</span>: sample[<span class="hljs-string">&quot;labels&quot;</span>]}
tf_dataset = tf_dataset.<span class="hljs-built_in">map</span>(decode_fn)
tf_dataset = tf_dataset.batch(BATCH_SIZE, drop_remainder=<span class="hljs-literal">True</span>)
<span class="hljs-built_in">print</span>(tf_dataset.element_spec)<!-- HTML_TAG_END --></pre></div> <p data-svelte-h="svelte-1s9k8gm">The dataset can now be passed to the <a href="https://keras.io/api/models/model_training_apis/#fit-method" rel="nofollow">fit</a> method.</p> <div class="code-block relative"><div class="absolute top-2.5 right-4"><button class="inline-flex items-center relative text-sm focus:text-green-500 cursor-pointer focus:outline-none transition duration-200 ease-in-out opacity-0 mx-0.5 text-gray-600 " title="code excerpt" type="button"><svg class="" xmlns="http://www.w3.org/2000/svg" aria-hidden="true" fill="currentColor" focusable="false" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 32 32"><path d="M28,10V28H10V10H28m0-2H10a2,2,0,0,0-2,2V28a2,2,0,0,0,2,2H28a2,2,0,0,0,2-2V10a2,2,0,0,0-2-2Z" transform="translate(0)"></path><path d="M4,18H2V4A2,2,0,0,1,4,2H18V4H4Z" transform="translate(0)"></path><rect fill="none" width="32" height="32"></rect></svg> <div class="absolute pointer-events-none transition-opacity bg-black text-white py-1 px-2 leading-tight rounded font-normal shadow left-1/2 top-full transform -translate-x-1/2 translate-y-2 opacity-0"><div class="absolute bottom-full left-1/2 transform -translate-x-1/2 w-0 h-0 border-black border-4 border-t-0" style="border-left-color: transparent; border-right-color: transparent; "></div> Copied</div></button></div> <pre class=""><!-- HTML_TAG_START --><span class="hljs-keyword">from</span> transformers <span class="hljs-keyword">import</span> TFAutoModelForImageClassification
<span class="hljs-keyword">with</span> strategy.scope():
 model = TFAutoModelForImageClassification.from_pretrained(image_model_checkpoint)
 model.<span class="hljs-built_in">compile</span>(optimizer=<span class="hljs-string">&quot;adam&quot;</span>)
model.fit(tf_dataset)<!-- HTML_TAG_END --></pre></div> <h3 class="relative group"><a id="stream-with-preparetfdataset" class="header-link block pr-1.5 text-lg no-hover:hidden with-hover:absolute with-hover:p-1.5 with-hover:opacity-0 with-hover:group-hover:opacity-100 with-hover:right-full" href="#stream-with-preparetfdataset"><span><svg class="" xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" aria-hidden="true" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 256 256"><path d="M167.594 88.393a8.001 8.001 0 0 1 0 11.314l-67.882 67.882a8 8 0 1 1-11.314-11.315l67.882-67.881a8.003 8.003 0 0 1 11.314 0zm-28.287 84.86l-28.284 28.284a40 40 0 0 1-56.567-56.567l28.284-28.284a8 8 0 0 0-11.315-11.315l-28.284 28.284a56 56 0 0 0 79.196 79.197l28.285-28.285a8 8 0 1 0-11.315-11.314zM212.852 43.14a56.002 56.002 0 0 0-79.196 0l-28.284 28.284a8 8 0 1 0 11.314 11.314l28.284-28.284a40 40 0 0 1 56.568 56.567l-28.285 28.285a8 8 0 0 0 11.315 11.314l28.284-28.284a56.065 56.065 0 0 0 0-79.196z" fill="currentColor"></path></svg></span></a> <span>Stream with prepare_tf_dataset</span></h3> <p data-svelte-h="svelte-1nsr6aj"><a href="/docs/transformers/pr_36839/en/main_classes/model#transformers.TFPreTrainedModel.prepare_tf_dataset">prepare_tf_dataset()</a> creates a <code>tf.data</code> pipeline that loads samples from <a href="https://www.tensorflow.org/api_docs/python/tf/data/Dataset" rel="nofollow">tf.data.Dataset</a>. The pipeline uses <a href>tf.numpy_function</a> or <a href="https://huggingface.co/docs/datasets/main/en/package_reference/main_classes#datasets.Dataset.from_generator" rel="nofollow">from_generator</a>, which can’t be compiled by TensorFlow, to access the underlying <code>tf.data.Dataset</code>. It also won’t work on a Colab TPU or TPU Nodes because the pipeline streams data from a local disk. Refer to the table below to help you decide whether this approach is helpful for you.</p> <table data-svelte-h="svelte-1la2fd5"><thead><tr><th>pros</th> <th>cons</th></tr></thead> <tbody><tr><td>simple code</td> <td>only works on TPU VM</td></tr> <tr><td>same approach on TPU/GPU</td> <td>data must be available as a Hugging Face Dataset</td></tr> <tr><td>dataset doesn’t have to fit in memory</td> <td>data must fit on local storage</td></tr> <tr><td>supports variable padding</td> <td>data loading may be a bottleneck on a big TPU pod slice</td></tr></tbody></table> <p data-svelte-h="svelte-x22cj5"><a href="/docs/transformers/pr_36839/en/main_classes/model#transformers.TFPreTrainedModel.prepare_tf_dataset">prepare_tf_dataset()</a> only works on <a href="#tpu-types">TPU VM</a>. Add the tokenizer output as columns in the dataset since the dataset is stored on disk, which means it can handle data larger than the available memory. Use <a href="/docs/transformers/pr_36839/en/main_classes/model#transformers.TFPreTrainedModel.prepare_tf_dataset">prepare_tf_dataset()</a> to stream data from the dataset by wrapping it with a <code>tf.data</code> pipeline.</p> <div class="code-block relative"><div class="absolute top-2.5 right-4"><button class="inline-flex items-center relative text-sm focus:text-green-500 cursor-pointer focus:outline-none transition duration-200 ease-in-out opacity-0 mx-0.5 text-gray-600 " title="code excerpt" type="button"><svg class="" xmlns="http://www.w3.org/2000/svg" aria-hidden="true" fill="currentColor" focusable="false" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 32 32"><path d="M28,10V28H10V10H28m0-2H10a2,2,0,0,0-2,2V28a2,2,0,0,0,2,2H28a2,2,0,0,0,2-2V10a2,2,0,0,0-2-2Z" transform="translate(0)"></path><path d="M4,18H2V4A2,2,0,0,1,4,2H18V4H4Z" transform="translate(0)"></path><rect fill="none" width="32" height="32"></rect></svg> <div class="absolute pointer-events-none transition-opacity bg-black text-white py-1 px-2 leading-tight rounded font-normal shadow left-1/2 top-full transform -translate-x-1/2 translate-y-2 opacity-0"><div class="absolute bottom-full left-1/2 transform -translate-x-1/2 w-0 h-0 border-black border-4 border-t-0" style="border-left-color: transparent; border-right-color: transparent; "></div> Copied</div></button></div> <pre class=""><!-- HTML_TAG_START --><span class="hljs-keyword">def</span> <span class="hljs-title function_">tokenize_function</span>(<span class="hljs-params">examples</span>):
 <span class="hljs-keyword">return</span> tokenizer(
 examples[<span class="hljs-string">&quot;sentence&quot;</span>], padding=<span class="hljs-string">&quot;max_length&quot;</span>, truncation=<span class="hljs-literal">True</span>, max_length=<span class="hljs-number">128</span>
 )
<span class="hljs-comment"># add the tokenizer output to the dataset as new columns</span>
dataset = dataset.<span class="hljs-built_in">map</span>(tokenize_function)
<span class="hljs-comment"># prepare_tf_dataset() chooses columns that match the models input names</span>
tf_dataset = model.prepare_tf_dataset(
 dataset, batch_size=BATCH_SIZE, shuffle=<span class="hljs-literal">True</span>, tokenizer=tokenizer
)<!-- HTML_TAG_END --></pre></div> <p data-svelte-h="svelte-1s9k8gm">The dataset can now be passed to the <a href="https://keras.io/api/models/model_training_apis/#fit-method" rel="nofollow">fit</a> method.</p> <div class="code-block relative"><div class="absolute top-2.5 right-4"><button class="inline-flex items-center relative text-sm focus:text-green-500 cursor-pointer focus:outline-none transition duration-200 ease-in-out opacity-0 mx-0.5 text-gray-600 " title="code excerpt" type="button"><svg class="" xmlns="http://www.w3.org/2000/svg" aria-hidden="true" fill="currentColor" focusable="false" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 32 32"><path d="M28,10V28H10V10H28m0-2H10a2,2,0,0,0-2,2V28a2,2,0,0,0,2,2H28a2,2,0,0,0,2-2V10a2,2,0,0,0-2-2Z" transform="translate(0)"></path><path d="M4,18H2V4A2,2,0,0,1,4,2H18V4H4Z" transform="translate(0)"></path><rect fill="none" width="32" height="32"></rect></svg> <div class="absolute pointer-events-none transition-opacity bg-black text-white py-1 px-2 leading-tight rounded font-normal shadow left-1/2 top-full transform -translate-x-1/2 translate-y-2 opacity-0"><div class="absolute bottom-full left-1/2 transform -translate-x-1/2 w-0 h-0 border-black border-4 border-t-0" style="border-left-color: transparent; border-right-color: transparent; "></div> Copied</div></button></div> <pre class=""><!-- HTML_TAG_START --><span class="hljs-keyword">from</span> transformers <span class="hljs-keyword">import</span> AutoTokenizer, TFAutoModelForSequenceClassification
<span class="hljs-keyword">with</span> strategy.scope():
 model = TFAutoModelForSequenceClassification.from_pretrained(model_checkpoint)
 model.<span class="hljs-built_in">compile</span>(optimizer=<span class="hljs-string">&quot;adam&quot;</span>)
model.fit(tf_dataset)<!-- HTML_TAG_END --></pre></div> <h2 class="relative group"><a id="tpu-types" class="header-link block pr-1.5 text-lg no-hover:hidden with-hover:absolute with-hover:p-1.5 with-hover:opacity-0 with-hover:group-hover:opacity-100 with-hover:right-full" href="#tpu-types"><span><svg class="" xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" aria-hidden="true" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 256 256"><path d="M167.594 88.393a8.001 8.001 0 0 1 0 11.314l-67.882 67.882a8 8 0 1 1-11.314-11.315l67.882-67.881a8.003 8.003 0 0 1 11.314 0zm-28.287 84.86l-28.284 28.284a40 40 0 0 1-56.567-56.567l28.284-28.284a8 8 0 0 0-11.315-11.315l-28.284 28.284a56 56 0 0 0 79.196 79.197l28.285-28.285a8 8 0 1 0-11.315-11.314zM212.852 43.14a56.002 56.002 0 0 0-79.196 0l-28.284 28.284a8 8 0 1 0 11.314 11.314l28.284-28.284a40 40 0 0 1 56.568 56.567l-28.285 28.285a8 8 0 0 0 11.315 11.314l28.284-28.284a56.065 56.065 0 0 0 0-79.196z" fill="currentColor"></path></svg></span></a> <span>TPU types</span></h2> <p data-svelte-h="svelte-1xpcsyo">There are two types of TPUs, a TPU Node and a TPU VM.</p> <p data-svelte-h="svelte-1bu3p95">A TPU Node indirectly accesses a remote TPU. It requires a separate VM to initialize your network and data pipeline, and then forwards it to the remote node. Google Colab TPUs are an example of a TPU Node. You can’t use local data because the TPU is remotely located, and data must be stored in Google Cloud Storage where the data pipeline can access it.</p> <p data-svelte-h="svelte-190h491">TPU VM are connected directly to the machine the TPU is located on, and they are generally easier to work with, especially when it comes to your data pipeline.</p> <div class="course-tip bg-gradient-to-br dark:bg-gradient-to-r before:border-green-500 dark:before:border-green-800 from-green-50 dark:from-gray-900 to-white dark:to-gray-950 border border-green-50 text-green-700 dark:text-gray-400"><p data-svelte-h="svelte-h8mx6j">We recommend avoiding TPU Nodes if possible because it is more difficult to debug than TPU VMs. TPU Nodes may also be unsupported in the future and become a legacy access method.</p></div> <p data-svelte-h="svelte-58j4yw">A single TPU (v2-8, v3-8, v4-8) runs 8 replicas. TPUs can exist in <strong>pods</strong> which run hundreds or even thousands of replicas simultaneously. When you only use a portion of a pod, it is referred to as a <strong>pod slice</strong>. On Google Colab, you’ll typically get a single v2-8 TPU.</p> <h2 class="relative group"><a id="xla" class="header-link block pr-1.5 text-lg no-hover:hidden with-hover:absolute with-hover:p-1.5 with-hover:opacity-0 with-hover:group-hover:opacity-100 with-hover:right-full" href="#xla"><span><svg class="" xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" aria-hidden="true" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 256 256"><path d="M167.594 88.393a8.001 8.001 0 0 1 0 11.314l-67.882 67.882a8 8 0 1 1-11.314-11.315l67.882-67.881a8.003 8.003 0 0 1 11.314 0zm-28.287 84.86l-28.284 28.284a40 40 0 0 1-56.567-56.567l28.284-28.284a8 8 0 0 0-11.315-11.315l-28.284 28.284a56 56 0 0 0 79.196 79.197l28.285-28.285a8 8 0 1 0-11.315-11.314zM212.852 43.14a56.002 56.002 0 0 0-79.196 0l-28.284 28.284a8 8 0 1 0 11.314 11.314l28.284-28.284a40 40 0 0 1 56.568 56.567l-28.285 28.285a8 8 0 0 0 11.315 11.314l28.284-28.284a56.065 56.065 0 0 0 0-79.196z" fill="currentColor"></path></svg></span></a> <span>XLA</span></h2> <p data-svelte-h="svelte-qlxbsa"><a href="https://openxla.org/xla" rel="nofollow">XLA</a> is a linear algebra compiler for high-performance execution and it is used by default to improve performance on TPUs.</p> <p data-svelte-h="svelte-1uul0oc">Before executing your code on a TPU, it’s a good idea to try it first on a CPU or GPU because it is easier to debug. You can train for a few steps to make sure the model and data pipeline work as expected. Set <code>jit_compile=True</code> in the <a href="https://keras.io/api/models/model_training_apis/#compile-method" rel="nofollow">compile</a> method to enable XLA compilation (but remember to remove this line of code before running on a TPU).</p> <p data-svelte-h="svelte-n8bji2">The section below outlines three rules for making your code XLA-compatible. Transformers enforce the first two rules for models and loss functions by default, but don’t forget about them if you’re writing your own models and loss functions.</p> <h3 class="relative group"><a id="data-dependent-conditionals" class="header-link block pr-1.5 text-lg no-hover:hidden with-hover:absolute with-hover:p-1.5 with-hover:opacity-0 with-hover:group-hover:opacity-100 with-hover:right-full" href="#data-dependent-conditionals"><span><svg class="" xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" aria-hidden="true" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 256 256"><path d="M167.594 88.393a8.001 8.001 0 0 1 0 11.314l-67.882 67.882a8 8 0 1 1-11.314-11.315l67.882-67.881a8.003 8.003 0 0 1 11.314 0zm-28.287 84.86l-28.284 28.284a40 40 0 0 1-56.567-56.567l28.284-28.284a8 8 0 0 0-11.315-11.315l-28.284 28.284a56 56 0 0 0 79.196 79.197l28.285-28.285a8 8 0 1 0-11.315-11.314zM212.852 43.14a56.002 56.002 0 0 0-79.196 0l-28.284 28.284a8 8 0 1 0 11.314 11.314l28.284-28.284a40 40 0 0 1 56.568 56.567l-28.285 28.285a8 8 0 0 0 11.315 11.314l28.284-28.284a56.065 56.065 0 0 0 0-79.196z" fill="currentColor"></path></svg></span></a> <span>Data dependent conditionals</span></h3> <p data-svelte-h="svelte-1chl77v">Any <code>if</code> statements cannot depend on values inside a <a href="https://www.tensorflow.org/api_docs/python/tf/Tensor" rel="nofollow">tf.Tensor</a>. The code below can’t be compiled by XLA.</p> <div class="code-block relative"><div class="absolute top-2.5 right-4"><button class="inline-flex items-center relative text-sm focus:text-green-500 cursor-pointer focus:outline-none transition duration-200 ease-in-out opacity-0 mx-0.5 text-gray-600 " title="code excerpt" type="button"><svg class="" xmlns="http://www.w3.org/2000/svg" aria-hidden="true" fill="currentColor" focusable="false" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 32 32"><path d="M28,10V28H10V10H28m0-2H10a2,2,0,0,0-2,2V28a2,2,0,0,0,2,2H28a2,2,0,0,0,2-2V10a2,2,0,0,0-2-2Z" transform="translate(0)"></path><path d="M4,18H2V4A2,2,0,0,1,4,2H18V4H4Z" transform="translate(0)"></path><rect fill="none" width="32" height="32"></rect></svg> <div class="absolute pointer-events-none transition-opacity bg-black text-white py-1 px-2 leading-tight rounded font-normal shadow left-1/2 top-full transform -translate-x-1/2 translate-y-2 opacity-0"><div class="absolute bottom-full left-1/2 transform -translate-x-1/2 w-0 h-0 border-black border-4 border-t-0" style="border-left-color: transparent; border-right-color: transparent; "></div> Copied</div></button></div> <pre class=""><!-- HTML_TAG_START --><span class="hljs-keyword">if</span> tf.reduce_sum(tensor) &gt; <span class="hljs-number">10</span>:
 tensor = tensor / <span class="hljs-number">2.0</span><!-- HTML_TAG_END --></pre></div> <p data-svelte-h="svelte-336eo8">To compile with XLA, use <a href="https://www.tensorflow.org/api_docs/python/tf/cond" rel="nofollow">tf.cond</a> or remove the conditional and use indicator variables instead as shown below.</p> <div class="code-block relative"><div class="absolute top-2.5 right-4"><button class="inline-flex items-center relative text-sm focus:text-green-500 cursor-pointer focus:outline-none transition duration-200 ease-in-out opacity-0 mx-0.5 text-gray-600 " title="code excerpt" type="button"><svg class="" xmlns="http://www.w3.org/2000/svg" aria-hidden="true" fill="currentColor" focusable="false" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 32 32"><path d="M28,10V28H10V10H28m0-2H10a2,2,0,0,0-2,2V28a2,2,0,0,0,2,2H28a2,2,0,0,0,2-2V10a2,2,0,0,0-2-2Z" transform="translate(0)"></path><path d="M4,18H2V4A2,2,0,0,1,4,2H18V4H4Z" transform="translate(0)"></path><rect fill="none" width="32" height="32"></rect></svg> <div class="absolute pointer-events-none transition-opacity bg-black text-white py-1 px-2 leading-tight rounded font-normal shadow left-1/2 top-full transform -translate-x-1/2 translate-y-2 opacity-0"><div class="absolute bottom-full left-1/2 transform -translate-x-1/2 w-0 h-0 border-black border-4 border-t-0" style="border-left-color: transparent; border-right-color: transparent; "></div> Copied</div></button></div> <pre class=""><!-- HTML_TAG_START -->sum_over_10 = tf.cast(tf.reduce_sum(tensor) &gt; <span class="hljs-number">10</span>, tf.float32)
tensor = tensor / (<span class="hljs-number">1.0</span> + sum_over_10)<!-- HTML_TAG_END --></pre></div> <h3 class="relative group"><a id="data-dependent-shapes" class="header-link block pr-1.5 text-lg no-hover:hidden with-hover:absolute with-hover:p-1.5 with-hover:opacity-0 with-hover:group-hover:opacity-100 with-hover:right-full" href="#data-dependent-shapes"><span><svg class="" xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" aria-hidden="true" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 256 256"><path d="M167.594 88.393a8.001 8.001 0 0 1 0 11.314l-67.882 67.882a8 8 0 1 1-11.314-11.315l67.882-67.881a8.003 8.003 0 0 1 11.314 0zm-28.287 84.86l-28.284 28.284a40 40 0 0 1-56.567-56.567l28.284-28.284a8 8 0 0 0-11.315-11.315l-28.284 28.284a56 56 0 0 0 79.196 79.197l28.285-28.285a8 8 0 1 0-11.315-11.314zM212.852 43.14a56.002 56.002 0 0 0-79.196 0l-28.284 28.284a8 8 0 1 0 11.314 11.314l28.284-28.284a40 40 0 0 1 56.568 56.567l-28.285 28.285a8 8 0 0 0 11.315 11.314l28.284-28.284a56.065 56.065 0 0 0 0-79.196z" fill="currentColor"></path></svg></span></a> <span>Data dependent shapes</span></h3> <p data-svelte-h="svelte-eeskwc">The shape of a <a href="https://www.tensorflow.org/api_docs/python/tf/Tensor" rel="nofollow">tf.Tensor</a> cannot depend on their values. For example, <a href="https://www.tensorflow.org/api_docs/python/tf/unique" rel="nofollow">tf.unique</a> can’t be compiled because it returns a tensor containing an instance of each unique value in the input. The shape of this output depends on how repetitive the input <a href="https://www.tensorflow.org/api_docs/python/tf/Tensor" rel="nofollow">tf.Tensor</a> is.</p> <p data-svelte-h="svelte-8wn0zc">This is an issue during <strong>label masking</strong>, where labels are set to a negative value to indicate they should be ignored when computing the loss. The code below can’t be compiled by XLA because the shape of <code>masked_outputs</code> and <code>masked_labels</code> depend on how many positions are masked.</p> <div class="code-block relative"><div class="absolute top-2.5 right-4"><button class="inline-flex items-center relative text-sm focus:text-green-500 cursor-pointer focus:outline-none transition duration-200 ease-in-out opacity-0 mx-0.5 text-gray-600 " title="code excerpt" type="button"><svg class="" xmlns="http://www.w3.org/2000/svg" aria-hidden="true" fill="currentColor" focusable="false" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 32 32"><path d="M28,10V28H10V10H28m0-2H10a2,2,0,0,0-2,2V28a2,2,0,0,0,2,2H28a2,2,0,0,0,2-2V10a2,2,0,0,0-2-2Z" transform="translate(0)"></path><path d="M4,18H2V4A2,2,0,0,1,4,2H18V4H4Z" transform="translate(0)"></path><rect fill="none" width="32" height="32"></rect></svg> <div class="absolute pointer-events-none transition-opacity bg-black text-white py-1 px-2 leading-tight rounded font-normal shadow left-1/2 top-full transform -translate-x-1/2 translate-y-2 opacity-0"><div class="absolute bottom-full left-1/2 transform -translate-x-1/2 w-0 h-0 border-black border-4 border-t-0" style="border-left-color: transparent; border-right-color: transparent; "></div> Copied</div></button></div> <pre class=""><!-- HTML_TAG_START -->label_mask = labels &gt;= <span class="hljs-number">0</span>
masked_outputs = outputs[label_mask]
masked_labels = labels[label_mask]
loss = compute_loss(masked_outputs, masked_labels)
mean_loss = torch.mean(loss)<!-- HTML_TAG_END --></pre></div> <p data-svelte-h="svelte-1wibu7k">To compile with XLA, avoid the data-dependent shapes by computing the loss for every position and zeroing out the masked positions in both the numerator and denominator when calculating the mean. Convert <code>tf.bool</code> to <code>tf.float32</code> as an indicator variable to make your code XLA-compatible.</p> <div class="code-block relative"><div class="absolute top-2.5 right-4"><button class="inline-flex items-center relative text-sm focus:text-green-500 cursor-pointer focus:outline-none transition duration-200 ease-in-out opacity-0 mx-0.5 text-gray-600 " title="code excerpt" type="button"><svg class="" xmlns="http://www.w3.org/2000/svg" aria-hidden="true" fill="currentColor" focusable="false" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 32 32"><path d="M28,10V28H10V10H28m0-2H10a2,2,0,0,0-2,2V28a2,2,0,0,0,2,2H28a2,2,0,0,0,2-2V10a2,2,0,0,0-2-2Z" transform="translate(0)"></path><path d="M4,18H2V4A2,2,0,0,1,4,2H18V4H4Z" transform="translate(0)"></path><rect fill="none" width="32" height="32"></rect></svg> <div class="absolute pointer-events-none transition-opacity bg-black text-white py-1 px-2 leading-tight rounded font-normal shadow left-1/2 top-full transform -translate-x-1/2 translate-y-2 opacity-0"><div class="absolute bottom-full left-1/2 transform -translate-x-1/2 w-0 h-0 border-black border-4 border-t-0" style="border-left-color: transparent; border-right-color: transparent; "></div> Copied</div></button></div> <pre class=""><!-- HTML_TAG_START -->label_mask = tf.cast(labels &gt;= <span class="hljs-number">0</span>, tf.float32)
loss = compute_loss(outputs, labels)
loss = loss * label_mask
mean_loss = tf.reduce_sum(loss) / tf.reduce_sum(label_mask)<!-- HTML_TAG_END --></pre></div> <h3 class="relative group"><a id="recompile-different-input-shapes" class="header-link block pr-1.5 text-lg no-hover:hidden with-hover:absolute with-hover:p-1.5 with-hover:opacity-0 with-hover:group-hover:opacity-100 with-hover:right-full" href="#recompile-different-input-shapes"><span><svg class="" xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" aria-hidden="true" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 256 256"><path d="M167.594 88.393a8.001 8.001 0 0 1 0 11.314l-67.882 67.882a8 8 0 1 1-11.314-11.315l67.882-67.881a8.003 8.003 0 0 1 11.314 0zm-28.287 84.86l-28.284 28.284a40 40 0 0 1-56.567-56.567l28.284-28.284a8 8 0 0 0-11.315-11.315l-28.284 28.284a56 56 0 0 0 79.196 79.197l28.285-28.285a8 8 0 1 0-11.315-11.314zM212.852 43.14a56.002 56.002 0 0 0-79.196 0l-28.284 28.284a8 8 0 1 0 11.314 11.314l28.284-28.284a40 40 0 0 1 56.568 56.567l-28.285 28.285a8 8 0 0 0 11.315 11.314l28.284-28.284a56.065 56.065 0 0 0 0-79.196z" fill="currentColor"></path></svg></span></a> <span>Recompile different input shapes</span></h3> <p data-svelte-h="svelte-1sa5b56">XLA recompiles your model if input shapes are variable which create huge performance problems. It is especially common in text models because input texts have variable lengths after tokenization.</p> <div class="course-tip course-tip-orange bg-gradient-to-br dark:bg-gradient-to-r before:border-orange-500 dark:before:border-orange-800 from-orange-50 dark:from-gray-900 to-white dark:to-gray-950 border border-orange-50 text-orange-700 dark:text-gray-400"><p data-svelte-h="svelte-aao7ck">Execessive padding can also severely slow down training because requires more compute and memory to process.</p></div> <p data-svelte-h="svelte-12qvhff">To avoid different shapes, use padding to pad all your inputs to the same length and use an <code>attention_mask</code>. Try padding batches of samples to a multiple of 32 or 64 tokens. Use the parameters <code>padding=&quot;max_length&quot;</code>, <code>padding=&quot;longest&quot;</code>, or <code>pad_to_multiple_of</code> to help with padding. This often increases the number of tokens by a small amount, but it significantly reduces the number of unique input shapes because every input shape is a multiple of 32 or 64. Fewer unique input shapes requires fewer recompilation.</p> <a class="!text-gray-400 !no-underline text-sm flex items-center not-prose mt-4" href="https://github.com/huggingface/transformers/blob/main/docs/source/en/perf_train_tpu_tf.md" target="_blank"><span data-svelte-h="svelte-1kd6by1">&lt;</span> <span data-svelte-h="svelte-x0xyl0">&gt;</span> <span data-svelte-h="svelte-1dajgef"><span class="underline ml-1.5">Update</span> on GitHub</span></a> <p></p>
<script>
 {
 __sveltekit_1bm5psi = {
 assets: "/docs/transformers/pr_36839/en",
 base: "/docs/transformers/pr_36839/en",
 env: {}
 };
 const element = document.currentScript.parentElement;
 const data = [null,null];
 Promise.all([
 import("/docs/transformers/pr_36839/en/_app/immutable/entry/start.6be8d590.js"),
 import("/docs/transformers/pr_36839/en/_app/immutable/entry/app.09748b4b.js")
 ]).then(([kit, app]) => {
 kit.start(app, element, {
 node_ids: [0, 423],
 data,
 form: null,
 error: null
 });
 });
 }
 </script>

Xet Storage Details

Size:
68 kB
·
Xet hash:
faf1d1b0fff82e8361cc947dbc64960024becac051265eb46db62bb7246f7cb8

Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.