有没有一种方法可以在每个总和上获得带有stop_gradient的张量流积？

Question

所以基本上像

w[t+1] = stop_gradient(w[t]) + dw[t]

但是有效。我假设for循环与cumsum相比不是很好。

我也可以只重组代码，但是可能有一种巧妙的方法可以完成这项工作，并且更容易阅读。

更新：

这是一个简单的示例，用于调整问题以避免累加：

import pandas as pd
from pylab import *
ion()
import numpy as np
import tensorflow as tf
def get_loss(w, lr=1.0):
    n = tf.range(w[0].shape[0] - 1, dtype=tf.float32)
    loss = 0
    dw = w[1:] - tf.stop_gradient(w[:-1])
    temp = lr * tf.reduce_sum(dw ** 2, axis=range(1, len(w.shape)))
    loss += tf.reduce_sum(temp) / 2
    return loss

opt = tf.keras.optimizers.SGD(learning_rate=0.1)
w = tf.keras.backend.variable(np.sin(np.linspace(0, 10 * np.pi, 1000)[:, None]))

w_ = [w.numpy()]

for i in range(1000):
    with tf.GradientTape() as tape:
        loss = get_loss(w) 
    g = tape.gradient(loss, [w])
    opt.apply_gradients(zip(g, [w]))
    if i % 200 == 0:
        w_.append(w.numpy())

ww = np.hstack(w_)

figure(1)
clf()
plot(ww, alpha=0.5)
show()