TensorFlow2中多网络梯度合并失败问题：梯度为None及AttributeError报错排查

Let's break down your issues step by step and fix the code properly:

1. Why You're Getting None Gradients #

The biggest culprit here is using np.log in your neg_log_likelihood function. NumPy operations break TensorFlow's computation graph gradient tracking chain, which is exactly why tape.gradient returns None. You need to replace this with TensorFlow's native tf.math.log.

Additionally, your dimension handling in the loss function was off: you were indexing mus[0] which grabs the first sample's mean, but you need to index across the batch dimension (e.g., mus[:, 0]) to get the first Gaussian component's mean for all samples in the batch.

2. Fixing the AttributeError: 'list' object has no attribute 'device' #

This error comes from passing a nested list of weights to apply_gradients. Each Dense layer's trainable_weights is already a list (containing the kernel and bias), so [m_phis.trainable_weights, m_mus.trainable_weights, ...] creates a list of lists. apply_gradients expects a flat list of (gradient, variable) pairs, so you need to combine all weights into a single flattened list.

Corrected Full Code #

Here's the fixed version with key changes highlighted:

import tensorflow as tf
import tensorflow_probability as tfp
from tensorflow import keras
import numpy as np

random_seed = 42
n_samples = 1000
phis_true = [0.3, 0.7]
mus_true = [-2.0, 1.5]
sigmas_true = [0.5, 0.8]

univariate_gmm = tfp.distributions.MixtureSameFamily(
    mixture_distribution=tfp.distributions.Categorical(probs=phis_true),
    components_distribution=tfp.distributions.Normal(loc=mus_true, scale=sigmas_true)
)
x = univariate_gmm.sample(n_samples, seed=random_seed).numpy()
dataset = tf.data.Dataset.from_tensor_slices(x)
dataset = dataset.shuffle(buffer_size=1024).batch(64)

m_phis = keras.layers.Dense(2, activation=tf.nn.softmax)
m_mus = keras.layers.Dense(2)
m_sigmas = keras.layers.Dense(2, activation=tf.nn.softplus)

def neg_log_likelihood(y, phis, mus, sigmas):
    # Use TF's log instead of NumPy's to keep gradient tracking intact
    a = tfp.distributions.Normal(loc=mus[:, 0], scale=sigmas[:, 0]).prob(y)
    b = tfp.distributions.Normal(loc=mus[:, 1], scale=sigmas[:, 1]).prob(y)
    # Index across batch dimension to get component-wise values
    mixture_prob = phis[:, 0] * a + phis[:, 1] * b
    c = tf.math.log(mixture_prob)
    return tf.reduce_sum(-c, axis=-1)

loss_fn = neg_log_likelihood
optimizer = tf.keras.optimizers.SGD(learning_rate=1e-3)

# Combine all trainable weights into a single flat list
all_trainable_weights = m_phis.trainable_weights + m_mus.trainable_weights + m_sigmas.trainable_weights

for step, y in enumerate(dataset):
    # Use tf.expand_dims to keep operations within TF's graph
    yy = tf.expand_dims(y, axis=1)
    with tf.GradientTape() as tape:
        phis = m_phis(yy)
        mus = m_mus(yy)
        sigmas = m_sigmas(yy)
        loss = loss_fn(yy, phis, mus, sigmas)
        # Average loss over the batch for more stable training
        loss = tf.reduce_mean(loss)
    
    # Calculate gradients against the flat weight list
    gradients = tape.gradient(loss, all_trainable_weights)
    # Apply gradients with matching flat pairs
    optimizer.apply_gradients(zip(gradients, all_trainable_weights))
    
    if step % 100 == 0:
        print("Step:", step, "Loss:", float(loss))

Key Changes Recap #

• Replaced np.log with tf.math.log: Preserves gradient tracking in the computation graph.
• Fixed dimension indexing: Ensures we're computing probabilities correctly across the entire batch.
• Flattened trainable weights: Eliminates the nested list issue that caused the device attribute error.
• Added batch loss averaging: Makes training more stable compared to summing all sample losses.
• Used tf.expand_dims: Keeps tensor operations within TensorFlow's ecosystem (better practice than NumPy here).

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