How to Use TensorFlow in Python (Step by Step)

How to use TensorFlow in Python in five steps: create a virtual environment on Python 3.10 to 3.13, run pip install tensorflow, import tensorflow as tf, build a model with the bundled Keras API, then train it with model.fit() and save it as a .keras file. That covers most real projects. Under the hood you are working with tensors, automatic gradients and optional graph compilation, and knowing those three ideas makes the rest of TensorFlow much easier to read.

This guide is a hands on walkthrough. Every code block below was run on TensorFlow 2.21.0 with Keras 3.15.1 and Python 3.13 on October 15, 2026, so you can paste it as is.

What you need before you start

TensorFlow is a Python library, so the main thing to get right is the Python version. The current release on PyPI is TensorFlow 2.21.0, and it ships wheels for Python 3.10, 3.11, 3.12 and 3.13 only.

If you just want to try code without installing anything, the official beginner quickstart runs as a notebook in Google Colab, and Kaggle notebooks are another hosted option if you'd rather not set up Python locally.

Step 1: Install TensorFlow in a virtual environment

The official pip install guide recommends a venv, which keeps TensorFlow and its dependencies away from your system Python:

python3 -m venv tf-env
source tf-env/bin/activate        # Windows: tf-env\Scripts\activate
pip install --upgrade pip
pip install tensorflow

Then check that it imports and can do some maths:

python3 -c "import tensorflow as tf; print(tf.reduce_sum(tf.random.normal([1000, 1000])))"

A few platform notes from that install page are worth knowing. TensorFlow is only officially released to PyPI, so use pip rather than conda. There is also a smaller CPU only package called tensorflow-cpu. On Windows, the native package is CPU only from version 2.11 onward; for an NVIDIA GPU you need WSL2. If you want GPU acceleration, follow our separate guide on how to use TensorFlow with GPU, which covers the [and-cuda] extra, Docker and the Mac Metal plugin. Nothing in your Python code changes when you move from CPU to GPU.

You'll see a few lines of log output when TensorFlow loads (oneDNN notices, CPU instruction hints). They are informational. Setting the environment variable TF_CPP_MIN_LOG_LEVEL=2 hides most of them.

Step 2: Work with tensors

A tensor is TensorFlow's array type: like a NumPy array, with a fixed shape and data type, but able to live on a GPU and take part in gradient calculations. The tensor guide covers them in depth. Here are the basics:

import tensorflow as tf

print("TensorFlow", tf.__version__)

a = tf.constant([[1.0, 2.0], [3.0, 4.0]])
b = tf.ones((2, 2))

print(a + b)            # element-wise add
print(tf.matmul(a, b))  # matrix multiply
print(a.shape, a.dtype) # (2, 2) float32
print(a.numpy())        # back to a NumPy array

TensorFlow 2 runs in eager mode by default, so each line runs immediately and prints a real value, just like ordinary Python. Tensors are immutable. When you need a value that changes, such as a model weight, you use a tf.Variable.

Step 3: Compute gradients with GradientTape

Training a neural network means nudging weights in the direction that lowers the loss, and that needs gradients. The autodiff guide explains that tf.GradientTape records operations as they run and then plays them back in reverse to compute gradients:

w = tf.Variable(3.0)

with tf.GradientTape() as tape:
    y = w ** 2

print(tape.gradient(y, w))  # 6.0, because dy/dw = 2w

You rarely call this by hand when you use Keras, because model.fit() does it for you. It matters when you write a custom training loop or a research model, and it is the same idea PyTorch calls autograd.

Step 4: Build and train a model with Keras

For almost every project, the easiest way to use TensorFlow in Python is through Keras, the high level API that ships with it. In TensorFlow 2.21, tf.keras and a plain import keras point at the same Keras 3 package, so either style works. This example follows the official quickstart: a small network that classifies handwritten digits from the MNIST dataset.

import keras

(x_train, y_train), (x_test, y_test) = keras.datasets.mnist.load_data()
x_train, x_test = x_train / 255.0, x_test / 255.0   # scale pixels to 0-1

model = keras.Sequential([
    keras.Input(shape=(28, 28)),
    keras.layers.Flatten(),
    keras.layers.Dense(128, activation="relu"),
    keras.layers.Dropout(0.2),
    keras.layers.Dense(10),
])

model.compile(
    optimizer="adam",
    loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
    metrics=["accuracy"],
)

model.summary()
model.fit(x_train, y_train, epochs=5, validation_split=0.1)
model.evaluate(x_test, y_test, verbose=2)

What each part does:

• Sequential stacks layers in order. Flatten turns each 28 by 28 image into a list of 784 numbers, Dense layers learn the patterns, and Dropout randomly switches off some units during training to reduce overfitting.
• compile picks the optimizer, the loss function and the metrics to report. from_logits=True tells the loss that the last layer outputs raw scores rather than probabilities.
• fit runs the training loop, and validation_split holds back 10% of the training data so you can watch for overfitting.
• evaluate scores the model on test data it has never seen.

We used keras.Input(shape=...) as the first layer instead of passing input_shape to Flatten, because current Keras prints a warning for the older style that some tutorials still show. In our test run on a CPU, even a single epoch reached about 96% validation accuracy on MNIST; your exact numbers will vary from run to run.

To get predictions, turn the raw scores into probabilities with softmax:

logits = model.predict(x_test[:5])
probs = tf.nn.softmax(logits).numpy()
print(probs.argmax(axis=1), y_test[:5])

Step 5: Feed data efficiently with tf.data

Passing NumPy arrays to fit() is fine for small datasets. For larger ones, or data that doesn't fit in memory, the tf.data guide recommends building an input pipeline that shuffles, batches and prefetches while the model trains:

ds = (
    tf.data.Dataset.from_tensor_slices((x_train, y_train))
    .shuffle(10_000)
    .batch(64)
    .prefetch(tf.data.AUTOTUNE)
)

model.fit(ds, epochs=1)

When you pass a tf.data.Dataset, shuffle it in the pipeline as shown. Keras ignores its own shuffle setting for datasets and prints a warning saying so.

Step 6: Speed up custom code with tf.function

Eager mode is easy to debug, but plain Python adds overhead. The tf.function guide explains that decorating a function with @tf.function traces it into a graph that TensorFlow can optimise and run without the Python interpreter in the loop:

@tf.function
def scaled_sum(x, scale):
    return tf.reduce_sum(x) * scale

print(scaled_sum(a, 2.0))  # 20.0

The guide's advice is simple: debug in eager mode first, then add the decorator. Inside a tf.function, stick to TensorFlow operations and avoid relying on Python side effects such as appending to a list, because those only happen while the function is being traced. Keras already compiles fit() this way, so you only need tf.function for your own training steps or preprocessing.

Step 7: Save and load your model

The Keras saving guide says the native format in Keras 3 is a single .keras file that holds the architecture, the weights, the compile settings and the optimizer state:

model.save("mnist.keras")
restored = keras.models.load_model("mnist.keras")

The extension matters. In our test, model.save("mnist") with no extension raised a ValueError asking for .keras (or legacy .h5). If you need a TensorFlow SavedModel for serving, use model.export("mnist_savedmodel") instead, which writes a SavedModel folder with a serve endpoint.

Using TensorFlow in Python 3.14

As of October 15, 2026, TensorFlow 2.21.0 on PyPI has no wheels for Python 3.14, so pip install tensorflow in a 3.14 environment fails with a "no matching distribution" style error. The fix is to create the virtual environment with an older interpreter, for example python3.13 -m venv tf-env, or use a tool such as uv or pyenv to install 3.12 or 3.13 alongside 3.14. Check the PyPI page before upgrading, because new wheels appear there first.

TensorFlow alternatives in Python

TensorFlow is still a solid choice for Keras workflows, existing production pipelines and TensorFlow Serving or LiteRT deployments. For open weight LLMs and most new research code, though, PyTorch is the default, and JAX is popular for TPU work and functional style research.

If speed and compilers are what you care about, our Tinygrad vs PyTorch performance guide explains how graph compilation trades startup time for faster steady state runs, which is the same trade tf.function makes. And if your end goal is adapting a language model rather than training a classifier, start with how to fine tune an LLM locally.

FAQ

Do I need to learn Keras to use TensorFlow?

Not strictly, but it is the recommended path. Keras is TensorFlow's high level API and the official quickstart uses it. Lower level tools such as tensors, tf.GradientTape and tf.function are there when you need custom behaviour.

Is tf.keras different from keras?

In TensorFlow 2.21 they are the same Keras 3 package. TensorFlow installs Keras 3.12 or newer as a dependency, and in our test tf.keras.Sequential and keras.Sequential were the same object.

Can I use TensorFlow in Python 3.14?

Not with TensorFlow 2.21, which only ships wheels for Python 3.10 to 3.13. Use a 3.12 or 3.13 virtual environment until a release with 3.14 wheels appears on PyPI.

Does TensorFlow work in Jupyter notebooks?

Yes. Install TensorFlow in your environment, install Jupyter in the same environment, and import tensorflow as tf in a cell. Google Colab is the zero install option, and the official tutorials link straight to it. Our guide on how to use TensorFlow in Jupyter Notebook covers kernels and common notebook errors.

Do I need a GPU to learn TensorFlow?

No. The MNIST example in this guide trains in seconds per epoch on a modern CPU. A GPU helps once models and datasets get larger; see how to use TensorFlow with GPU when you get there.

Is TensorFlow free?

Yes. TensorFlow is open source software released under the Apache 2.0 licence, and installing it from PyPI costs nothing.