backward 与 Gradients

static STATE: Mutex<Option<GraphLocator>> = Mutex::new(None);

#[derive(Default)]
pub struct AutodiffServer {
    steps: HashMap<NodeId, StepBoxed>,
    actions_builder: HashMap<NodeId, CheckpointerBuilder>,
    memory_management: GraphMemoryManagement,
}

        let root = steps.get(&root_id).ok_or(root_id)?;
        let mut nodes = vec![root_id];
        let mut visited = HashSet::new();
        let mut parents = root.parents().to_vec();
        visited.insert(root_id);

        while let Some(parent) = parents.pop() {
            if !visited.insert(parent.id) {
                continue;
            }
            match steps.get(&parent.id) {
                Some(step) => {
                    nodes.push(parent.id);
                    parents.extend_from_slice(step.parents());
                }
                None if parent.is_leaf => {}
                None => return Err(parent.id),
            }
        }

    fn execute_steps(
        tape: Vec<Vec<StepBoxed>>,
        mut grads: Gradients,
        mut checkpointer: Checkpointer,
    ) -> Gradients {
        tape.into_iter().rev().for_each(|steps| {
            steps
                .into_iter()
                .for_each(|step| step.step(&mut grads, &mut checkpointer))
        });
        // ...
        grads
    }

    pub fn consume<B: Backend>(&mut self, node: &NodeRef) -> FloatTensor<B> {
        match node.requirement {
            Requirement::Grad => self
                .container
                .get::<TensorPrimitive<B>>(&node.id.value)
                .map(|tensor| tensor.tensor())
                .expect("Can't consume the gradients before they are registered at least once."),
            Requirement::GradInBackward => self
                .container
                .remove::<TensorPrimitive<B>>(&node.id.value)
                .map(|tensor| tensor.tensor())
                .expect("Can't consume the gradients before they are registered at least once."),
            Requirement::None => panic!("Trying to consume the gradients for an untracked tensor"),
        }
    }

  let x = Tensor::<1>::from_floats([1.0, 2.0], &device).require_grad();
  let unused = Tensor::<1>::from_floats([5.0, 5.0], &device).require_grad();

  // x appears on two paths: x*x and x*3
  let h = x.clone() * x.clone();
  let loss = (h.clone() + x.clone() * 3.0).sum();
  let mut grads = loss.backward();

  println!("x.grad       = {:?}", v(x.grad(&grads)));
  println!("h.grad       = {:?}", v(h.grad(&grads)));
  println!("loss.grad    = {:?}", loss.grad(&grads).map(|t| t.dims()));
  println!("unused.grad  = {:?}", v(unused.grad(&grads)));
  println!("x.grad_remove= {:?}", v(x.grad_remove(&mut grads)));
  println!("x.grad again = {:?}", v(x.grad(&grads)));

x.grad       = Some([5.0, 7.0])
h.grad       = None
loss.grad    = None
unused.grad  = None
x.grad_remove= Some([5.0, 7.0])
x.grad again = None

thread 'main' (3168867) panicked at ~/.cargo/registry/src/index.crates.io-1949cf8c6b5b557f/burn-autodiff-0.22.1/src/runtime/server.rs:74:17:
Cannot run backward: graph tape has already been consumed at NodeId(5). Burn does not support retain_graph. Combine losses before backward or recompute the forward pass. Use detach() only to intentionally sever gradients through the previous computation.
note: run with `RUST_BACKTRACE=1` environment variable to display a backtrace
backward twice on the same loss: panicked
l1 -> x.grad = Some([5.4365635, 14.778112])

thread 'main' (3168867) panicked at ~/.cargo/registry/src/index.crates.io-1949cf8c6b5b557f/burn-autodiff-0.22.1/src/runtime/server.rs:74:17:
Cannot run backward: graph tape has already been consumed at NodeId(6). Burn does not support retain_graph. Combine losses before backward or recompute the forward pass. Use detach() only to intentionally sever gradients through the previous computation.
l2.backward after l1 consumed `shared`: panicked
l3 -> x.grad = Some([4.0, 4.0])
l1+l2 -> x.grad = Some([13.591409, 36.945282])
detached: is_tracked=false is_require_grad=false

fn forward(x: &Tensor<2>) -> Tensor<1> {
  (x.clone().exp() * x.clone()).matmul(x.clone()).sum()
}

fn main() {
  let gpu = Device::metal(DeviceKind::DefaultDevice);
  let ad = gpu.clone().autodiff();

  let x = Tensor::<2>::random([512, 512], Distribution::Default, &ad).require_grad();
  usage(&ad, "after creating x (1 MiB)");

  for i in 1..=20 {
    let loss = forward(&x);
    drop(loss);
    if i % 5 == 0 {
      usage(&ad, &format!("forward-only x{i}, loss dropped"));
    }
  }

  let grads = forward(&x).backward();
  usage(&ad, "after one backward");
  drop(grads);
  usage(&ad, "grads dropped");

  let xi = x.clone().without_autodiff();
  // ... 同样的前向 20 次
}

after creating x (1 MiB)           in_use=   1.0 MiB allocs=2
forward-only x5, loss dropped      in_use=  11.0 MiB allocs=18
forward-only x10, loss dropped     in_use=  21.0 MiB allocs=28
forward-only x15, loss dropped     in_use=  31.0 MiB allocs=38
forward-only x20, loss dropped     in_use=  41.0 MiB allocs=48
after one backward                 in_use=   2.0 MiB allocs=18
grads dropped                      in_use=   1.0 MiB allocs=17
without_autodiff forward x10       in_use=   1.0 MiB allocs=18
without_autodiff forward x20       in_use=   1.0 MiB allocs=18

5 abandoned graphs                 in_use=  15.0 MiB allocs=24
control: untracked op + read #1    in_use=  15.0 MiB allocs=25
control: untracked op + read #2    in_use=  15.0 MiB allocs=25
after unrelated backward #1        in_use=   0.0 MiB allocs=14