python - How do I make a copy of a GradientTape object in tensorflow? (Monte Carlo App)

Question

I have an expensive initialization op before I perform my calculation in tensorflow.

My code looks something like this:

x = tf.Variable(2.0)
w = tf.Variable(5.0)

with tf.GradientTape() as tape:
  tape.watch(x)
  tape.watch(w)
  y = x ** 2
  z = w ** 3
  o = tf.math.log(y*z) # note that this step is the arbitrarily complex init code

# now i need to run a loop n times (here n is 10)
res = []
for i in range(10):
  with tape:
    z =tf.random.normal([1,10])
    f = tf.reduce_sum(x*z, axis=1)*o+w
  df = tape.gradient(f, {'x':x, 'w':w})
  res.append(df)

Basically I'm trying to run a monte carlo simulation and need the gradients without having to run the initialization code on every loop. This code works fine if n==1 but gives the wrong answers if n>=2.

What i need is a way to copy the state of tape before I start the monte carlo loop. so instead of saying "with tape" something like:

  with tf.GradientTape(tape) as tape2:
    ...
  df = tape2.gradient(f, {'x':x, 'w':w})

Is this possible? How can I achieve something similar?

As a second part to the question, I've noticed that even if I recalculate the value of o in the main loop, tensorflow only works if the tape is not persistent. If it is - i run out of memory on the GPU after several iterations of the loop. This is not ideal as I'd like to also define other functions that depend on x and w and record their gradients also.

i.e. if i do this:

res = []
for i in range(10):
  with tf.GradientTape(persistent=True) as tape:
    z =tf.random.normal([1,10])

    # rerun init for every loop
    y = x ** 2
    t = w ** 3
    o = tf.math.log(y*t)

    f = tf.reduce_sum(x*z, axis=1)*o+w
    g = tf.reduce_sum(x*z+w, axis=1)*o

  df = tape.gradient(f, {'x':x, 'w':w})
  dg = tape.gradient(g, {'x':x, 'w':w})
  res.append([df,dg])

I don't understand this behavior - surely the tape is being discarded after each iteration of the loop (and thus it shouldn't matter if it was persistent or not)?

Thanks

Answer 1

等待大神答复