iwantcoding.com
🔥 Daily 👥 Rooms 🏆 Top Log in Sign up

Variables

tf.Variable holds trainable state. Different from tensors: persistent, mutable, tracked by autograd.

TensorFlow — variables

EXAMPLE
import tensorflow as tf

# ===== Create =====
w = tf.Variable(tf.random.normal([4, 3]))
b = tf.Variable(tf.zeros([3]))
counter = tf.Variable(0, trainable=False)         # not trained, just tracked

# Initial value sets the dtype:
v = tf.Variable(1.0)              # float32
v2 = tf.Variable([1, 2, 3])       # int32

# ===== Read =====
print(w)
print(w.numpy())                  # to numpy
print(w.shape, w.dtype, w.device)

# ===== Update =====
w.assign(w * 0.99)               # in-place
w.assign_add(tf.ones_like(w))    # +=
w.assign_sub(tf.ones_like(w))    # -=
# These do NOT create a new Variable; they update the existing one.

# Rebinding the Python variable does NOT update the tracked Variable:
# w = tf.Variable(...)            # this creates a new one; old graph holds the old

# ===== Trainable =====
non_trained = tf.Variable(0, trainable=False)
# trainable=False prevents inclusion in tape.watched_variables() automatically.

# ===== Within a tf.Module =====
class Linear(tf.Module):
    def __init__(self, in_features, out_features, name=None):
        super().__init__(name=name)
        self.w = tf.Variable(tf.random.normal([in_features, out_features]), name='w')
        self.b = tf.Variable(tf.zeros([out_features]), name='b')
    def __call__(self, x):
        return x @ self.w + self.b

lin = Linear(4, 3)
print(lin.trainable_variables)   # list including w + b

# ===== Use in autograd =====
x = tf.constant([[1., 2., 3., 4.]])
with tf.GradientTape() as tape:
    y = lin(x)
    loss = tf.reduce_sum(y ** 2)
grads = tape.gradient(loss, lin.trainable_variables)

# ===== With an optimizer =====
opt = tf.optimizers.Adam(1e-3)
opt.apply_gradients(zip(grads, lin.trainable_variables))

# ===== Persistence =====
ckpt = tf.train.Checkpoint(lin=lin, opt=opt)
ckpt.save('./ckpt/m')
# Restore:
ckpt2 = tf.train.Checkpoint(lin=lin, opt=opt)
ckpt2.restore(tf.train.latest_checkpoint('./ckpt'))

# ===== Keras Variables =====
# tf.keras.layers.Layer holds tf.Variables in .weights / .trainable_weights / .non_trainable_weights.
class MyLayer(tf.keras.layers.Layer):
    def build(self, input_shape):
        self.kernel = self.add_weight('kernel', shape=(input_shape[-1], 3), initializer='glorot_uniform')
    def call(self, x):
        return x @ self.kernel

# ===== Device + naming =====
# Variables can be pinned:
with tf.device('/CPU:0'):
    v3 = tf.Variable(tf.zeros([1000]))
# Naming helps in TensorBoard graphs:
w_named = tf.Variable(tf.random.normal([3]), name='w/layer1')

# ===== Patterns to internalise =====
# - Variables for trainable params; tensors for everything else
# - assign / assign_add / assign_sub — never rebind the Python name
# - Group Variables in tf.Module or Keras Layer for cleanliness
# - Use Checkpoint for save / restore; SavedModel for deploy

# ===== Pitfalls =====
# - Rebinding a Python variable instead of assign() -> tracked Variable forgotten
# - trainable=True on a Variable that should be a constant (slows training, bigger checkpoints)
# - Mixing dtype between a Variable and an update tensor -> TF refuses silently
# - Forgetting to include Variables in optimizer.apply_gradients (they will not learn)

Why it matters

Variables are TensorFlow trainable state. Create with tf.Variable, mutate with assign / assign_add, group inside tf.Module or Keras layer, save with Checkpoint. Never rebind the Python name; always assign. That single discipline saves a lot of confused autograd bugs.

Tip: Tweak the snippet with Try it Yourself », then sit the quiz at the bottom of the page.

Example

Example
import tensorflow as tf
w = tf.Variable(tf.random.normal((3, 1)))
w.assign_sub(0.01 * w)  # gradient-descent-ish update
Try it Yourself »

Discussion

Loading…