在 TensorFlow2 中結合來自不同“網絡”的梯度

Question

我試圖將一些“網絡”組合成一個最終的損失函數。 我想知道我所做的是否“合法”，到目前為止我似乎無法完成這項工作。 我正在使用張量流概率：

主要問題在這里：

# Get gradients of the loss wrt the weights.
gradients = tape.gradient(loss, [m_phis.trainable_weights, m_mus.trainable_weights, m_sigmas.trainable_weights])

# Update the weights of our linear layer.
optimizer.apply_gradients(zip(gradients, [m_phis.trainable_weights, m_mus.trainable_weights, m_sigmas.trainable_weights])

這給了我 None 漸變並拋出應用漸變：

AttributeError: 'list' 對象沒有屬性 'device'

完整代碼：

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):
    a = tfp.distributions.Normal(loc=mus[0],scale=sigmas[0]).prob(y)
    b = tfp.distributions.Normal(loc=mus[1],scale=sigmas[1]).prob(y)
    c = np.log(phis[0]*a + phis[1]*b)
    return tf.reduce_sum(-c, axis=-1)

# Instantiate a logistic loss function that expects integer targets.
loss_fn = neg_log_likelihood

# Instantiate an optimizer.
optimizer = tf.keras.optimizers.SGD(learning_rate=1e-3)

# Iterate over the batches of the dataset.
for step, y in enumerate(dataset):
    
    yy = np.expand_dims(y, axis=1)

    # Open a GradientTape.
    with tf.GradientTape() as tape:
        
        # Forward pass.
        phis = m_phis(yy)
        mus = m_mus(yy)
        sigmas = m_sigmas(yy)

        # Loss value for this batch.
        loss = loss_fn(yy, phis, mus, sigmas)

    # Get gradients of the loss wrt the weights.
    gradients = tape.gradient(loss, [m_phis.trainable_weights, m_mus.trainable_weights, m_sigmas.trainable_weights])

    # Update the weights of our linear layer.
    optimizer.apply_gradients(zip(gradients, [m_phis.trainable_weights, m_mus.trainable_weights, m_sigmas.trainable_weights]))

    # Logging.
    if step % 100 == 0:
        print("Step:", step, "Loss:", float(loss))

Answer 1

有兩個單獨的問題需要考慮。

1. 梯度是None ：

通常，如果在GradientTape監視的代碼中執行非張量流操作，就會發生這種情況。 具體而言，這種擔憂的計算np.log在neg_log_likelihood功能。 如果您要更換np.log與tf.math.log ，梯度應該計算。 盡量不要在“內部”tensorflow 組件中使用 numpy 可能是一個好習慣，因為這可以避免此類錯誤。 對於大多數 numpy 操作，有一個很好的 tensorflow 替代品。

2. apply_gradients用於多個可訓練對象：

這主要與apply_gradients期望的輸入有關。 你有兩個選擇：

第一個選項：調用apply_gradients三次，每次使用不同的可訓練數據

optimizer.apply_gradients(zip(m_phis_gradients, m_phis.trainable_weights))
optimizer.apply_gradients(zip(m_mus_gradients, m_mus.trainable_weights))
optimizer.apply_gradients(zip(m_sigmas_gradients, m_sigmas.trainable_weights))

另一種方法是創建一個元組列表，如tensorflow 文檔中所示（引用：“grads_and_vars：（梯度，變量）對列表。”）。 這意味着調用類似的東西

optimizer.apply_gradients(
   [
      zip(m_phis_gradients, m_phis.trainable_weights),
      zip(m_mus_gradients, m_mus.trainable_weights),
      zip(m_sigmas_gradients, m_sigmas.trainable_weights),
   ]
)

這兩個選項都要求您拆分漸變。 您可以通過計算梯度並分別為它們編制索引（ gradients[0],... ），或者您可以簡單地單獨計算梯度。 請注意，這可能需要您的GradientTape persistent=True 。

    # [...]
    # Open a GradientTape.
    with tf.GradientTape(persistent=True) as tape:
        # Forward pass.
        phis = m_phis(yy)
        mus = m_mus(yy)
        sigmas = m_sigmas(yy)

        # Loss value for this batch.
        loss = loss_fn(yy, phis, mus, sigmas)

    # Get gradients of the loss wrt the weights.
    m_phis_gradients = tape.gradient(loss, m_phis.trainable_weights)
    m_mus_gradients = tape.gradient(loss, m_mus.trainable_weights)
    m_sigmas_gradients = tape.gradient(loss, m_sigmas .trainable_weights)

    # Update the weights of our linear layer.
    optimizer.apply_gradients(
        [
            zip(m_phis_gradients, m_phis.trainable_weights),
            zip(m_mus_gradients, m_mus.trainable_weights),
            zip(m_sigmas_gradients, m_sigmas.trainable_weights),
       ]
   )
   # [...]