打开求导

a = torch.tensor([1.0, 2.0])
w = torch.tensor([3.0, 4.0], requires_grad=True)
aa = a * a
aw = a * w

let ad = Device::flex().autodiff();
let a = Tensor::<1>::from_floats([1.0, 2.0], &ad);
let w = Tensor::<1>::from_floats([3.0, 4.0], &ad).require_grad();
let aa = a.clone() * a.clone();
let aw = a.clone() * w.clone();

let lifted = Tensor::<1>::from_floats([1.0, 2.0], &plain).autodiff().require_grad();

def flags(name, t):
  print(f"{name:10} requires_grad={t.requires_grad!s:5} is_leaf={t.is_leaf!s:5} grad_fn={type(t.grad_fn).__name__ if t.grad_fn else None}")

pub fn flags<const D: usize>(name: &str, t: &Tensor<D>) {
  println!(
    "{name:10} is_autodiff={:5} is_tracked={:5} is_require_grad={}",
    t.is_autodiff(),
    t.is_tracked(),
    t.is_require_grad()
  );
}

a          requires_grad=False is_leaf=True  grad_fn=None
w          requires_grad=True  is_leaf=True  grad_fn=None
a*a        requires_grad=False is_leaf=True  grad_fn=None
a*w        requires_grad=True  is_leaf=False grad_fn=MulBackward0

a          is_autodiff=true  is_tracked=false is_require_grad=false
w          is_autodiff=true  is_tracked=true  is_require_grad=true
a*a        is_autodiff=true  is_tracked=false is_require_grad=false
a*w        is_autodiff=true  is_tracked=true  is_require_grad=false

attempt("int requires_grad", lambda: torch.tensor([1, 2], requires_grad=True))
attempt("backward on constant", lambda: (a * 2).sum().backward())
attempt("requires_grad_(True) on a*w", lambda: aw.requires_grad_(True))
attempt("requires_grad_(False) on a*w", lambda: aw.requires_grad_(False))
attempt("requires_grad_() on a*a", lambda: flags("a*a after", aa.requires_grad_()))

int requires_grad: RuntimeError: Only Tensors of floating point and complex dtype can require gradients
backward on constant: RuntimeError: element 0 of tensors does not require grad and does not have a grad_fn
requires_grad_(True) on a*w: ok
requires_grad_(False) on a*w: RuntimeError: you can only change requires_grad flags of leaf variables. If you want to use a computed variable in a subgraph that doesn't require differentiation use var_no_grad = var.detach().
a*a after  requires_grad=True  is_leaf=True  grad_fn=None
requires_grad_() on a*a: ok

let plain = Device::flex();
attempt("require_grad on Device::flex()", || {
  Tensor::<1>::from_floats([1.0, 2.0], &plain).require_grad()
});
attempt("backward on constant", || (a.clone() * 2.0).sum().backward());
attempt("require_grad on a*w", || aw.clone().require_grad());
attempt("set_require_grad(false) on a*w", || {
  flags("a*w off", &aw.clone().set_require_grad(false))
});
attempt("require_grad on a*a", || flags("a*a after", &aa.clone().require_grad()));

thread 'main' (8533721) panicked at .../burn-tensor-0.22.1/src/tensor/api/float.rs:1185:13:
Tensor::require_grad requires autodiff; call Tensor::autodiff first
note: run with `RUST_BACKTRACE=1` environment variable to display a backtrace
require_grad on Device::flex(): panicked

thread 'main' (8533721) panicked at .../burn-tensor-0.22.1/src/tensor/api/autodiff.rs:75:9:
Tensor::backward requires a tracked autodiff tensor; call Tensor::autodiff().require_grad() on the source leaf before computing the output
backward on constant: panicked

thread 'main' (8533721) panicked at .../burn-autodiff-0.22.1/src/tensor.rs:190:17:
Can't convert a non leaf tensor into a tracked tensor
require_grad on a*w: panicked
a*w off    is_autodiff=true  is_tracked=false is_require_grad=false
set_require_grad(false) on a*w: ok
a*a after  is_autodiff=true  is_tracked=true  is_require_grad=true
require_grad on a*a: ok

let ad = Device::flex().autodiff();
let i = Tensor::<1, Int>::from_ints([1, 2], &ad).require_grad();

error[E0599]: no method named `require_grad` found for struct `Tensor<1, burn::tensor::Int>` in the current scope
 --> src/main.rs:5:52
  |
5 |   let i = Tensor::<1, Int>::from_ints([1, 2], &ad).require_grad();
  |                                                    ^^^^^^^^^^^^
  |
help: there is a method `is_require_grad` with a similar name
  |
5 |   let i = Tensor::<1, Int>::from_ints([1, 2], &ad).is_require_grad();
  |                                                    +++

x = torch.tensor([1.0, 2.0])
used = x * w
x.requires_grad_()
used.sum().backward()
print("late x.grad =", x.grad, "| w.grad =", w.grad)

let x = Tensor::<1>::from_floats([1.0, 2.0], &ad);
let used = x.clone() * w.clone();
let x = x.require_grad();
let grads = used.sum().backward();
println!("late x.grad = {:?} | w.grad = {:?}", v(x.grad(&grads)), v(w.grad(&grads)));

late x.grad = None | w.grad = tensor([1., 2.])

late x.grad = None | w.grad = Some([1.0, 2.0])

fc = torch.nn.Linear(2, 2)
flags("fc.weight", fc.weight)

let fc = burn::nn::LinearConfig::new(2, 2).init(&ad);
flags("fc.weight", &fc.weight.val());
let fc_plain = burn::nn::LinearConfig::new(2, 2).init(&plain);
flags("plain fc.w", &fc_plain.weight.val());

fc.weight  requires_grad=True  is_leaf=True  grad_fn=None

fc.weight  is_autodiff=true  is_tracked=true  is_require_grad=true
plain fc.w is_autodiff=false is_tracked=false is_require_grad=false