TensorFlow中tf.assign手动更新权重无效，求正确实现方式

Hey there! I’ve run into this exact issue before, so let’s break down why your weight updates aren’t sticking and the best ways to implement your custom update step (subtracting a fixed value alongside regular weight updates).

Why Your Current Code Isn’t Working #

First, let’s clear up a common misconception: tf.Variable.assign() or assign_sub() does modify the original variable—it doesn’t operate on a copy. The most likely culprits for unapplied updates are:

• In TensorFlow 1.x: Forgetting to run the assign operation in a Session (since TF1 uses static graphs, operations don’t execute unless explicitly run).
• In TensorFlow 2.x: Accidentally working with a disconnected tensor (not the original variable) or not executing the assign operation in a context where it’s tracked (like inside a tf.function or eager execution).

Best Implementation Methods #

1. TensorFlow 2.x (Eager Execution / Keras Custom Loops) – Recommended

TF2’s default eager execution makes this straightforward. Use assign_sub() directly on your variable to subtract the fixed value—this modifies the variable’s internal storage (the end result you want).

Here’s a complete example integrating regular gradient updates and your custom fixed subtraction:

import tensorflow as tf

# Initialize your weight variable
weights = tf.Variable(tf.random.normal(shape=(2, 2)))
fixed_subtraction = tf.constant(0.05)  # Your fixed value to subtract
optimizer = tf.keras.optimizers.SGD(learning_rate=0.01)

# Custom training step
def train_step(inputs, labels):
    with tf.GradientTape() as tape:
        # Forward pass: compute loss
        predictions = tf.matmul(inputs, weights)
        loss = tf.reduce_mean(tf.square(predictions - labels))
    
    # 1. Regular weight update (gradient descent)
    grads = tape.gradient(loss, weights)
    optimizer.apply_gradients([(grads, weights)])
    
    # 2. Custom update: subtract fixed value from weights
    weights.assign_sub(fixed_subtraction)
    
    return loss

# Example usage
sample_inputs = tf.random.normal(shape=(3, 2))
sample_labels = tf.random.normal(shape=(3, 2))
loss = train_step(sample_inputs, sample_labels)

# Verify the update worked
print("Updated weights:\n", weights.numpy())

If you’re using tf.function for performance (converting to a static graph), the code works exactly the same—assign_sub() will be correctly tracked and executed in the graph:

@tf.function
def train_step(inputs, labels):
    # Same code as above
    with tf.GradientTape() as tape:
        predictions = tf.matmul(inputs, weights)
        loss = tf.reduce_mean(tf.square(predictions - labels))
    grads = tape.gradient(loss, weights)
    optimizer.apply_gradients([(grads, weights)])
    weights.assign_sub(fixed_subtraction)
    return loss

2. TensorFlow 1.x (Static Graph)

If you’re stuck on TF1, you need to ensure the assign operation is explicitly run in a Session. Here’s how:

import tensorflow as tf

tf.disable_eager_execution()

weights = tf.Variable(tf.random_normal(shape=(2, 2)))
fixed_subtraction = tf.constant(0.05)
optimizer = tf.train.GradientDescentOptimizer(learning_rate=0.01)

# Define loss and regular update op
inputs = tf.placeholder(tf.float32, shape=(None, 2))
labels = tf.placeholder(tf.float32, shape=(None, 2))
predictions = tf.matmul(inputs, weights)
loss = tf.reduce_mean(tf.square(predictions - labels))
regular_update_op = optimizer.minimize(loss)

# Define custom update op
custom_update_op = weights.assign_sub(fixed_subtraction)

# Run everything in a session
with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    
    # Sample data
    sample_inputs = tf.random_normal(shape=(3, 2)).eval()
    sample_labels = tf.random_normal(shape=(3, 2)).eval()
    
    # Execute both updates
    sess.run(regular_update_op, feed_dict={inputs: sample_inputs, labels: sample_labels})
    sess.run(custom_update_op)
    
    # Verify
    print("Updated weights:\n", sess.run(weights))

Key Takeaways #

• TF2: Always use assign_sub() (or assign_add()) directly on your tf.Variable—no need to reassign the result (e.g., don’t do weights = weights.assign_sub(...), just call weights.assign_sub(...)).
• TF1: Never forget to run the assign operation in your Session—static graphs won’t execute operations unless you explicitly trigger them.
• Avoid creating copies of your variable (e.g., via tf.identity() or unnecessary tensor operations) because this will break the link to the original trainable variable.

内容的提问来源于stack exchange，提问作者user4805479