torch._dynamo.exc.Unsupported: Tensor.backward

gengrill: <p>Hi! When seeing the PT 2.0 announcement blog I got really interested in Dynamo. I installed the nightlies using <code>pip3 install numpy --pre torch torchvision torchaudio --force-reinstall --extra-index-url https://download.pytorch.org/whl/nightly/cpu</code> and I’m following the Tutorial here: <a href="https://pytorch.org/tutorials/intermediate/dynamo_tutorial.html" rel="noopener nofollow ugc">https://pytorch.org/tutorials/intermediate/dynamo_tutorial.html</a></p>
<p>My goal is to export the ATen graph for the function below using Dynamo, but I get an error. Am I doing something wrong?</p>
<pre><code class="lang-auto">&gt;&gt;&gt; def train(model, data):                                                                                                                                                                                                                    
    pred = model(data[0])                                                                                                                                                                                                                  
    loss = nn.CrossEntropyLoss()(pred, data[1])                                                                                                                                                                                            
    loss.backward()                                                                                                                                                                                                                        
    return loss
...
&gt;&gt;&gt; train_exp = dynamo.export(train, model, generate_data(16), aten_graph=True)`
Traceback (most recent call last):
  File "&lt;stdin&gt;", line 1, in &lt;module&gt;
  File "/home/david/projects/ml-compiler-talk/pytorch-examples/pytorch-2.0/lib/python3.10/site-packages/torch/_dynamo/eval_frame.py", line 622, in export
    result_traced = opt_f(*args, **kwargs)
  File "/home/david/projects/ml-compiler-talk/pytorch-examples/pytorch-2.0/lib/python3.10/site-packages/torch/_dynamo/eval_frame.py", line 212, in _fn
    return fn(*args, **kwargs)
  File "/home/david/projects/ml-compiler-talk/pytorch-examples/pytorch-2.0/lib/python3.10/site-packages/torch/_dynamo/eval_frame.py", line 333, in catch_errors
    return callback(frame, cache_size, hooks)
  File "/home/david/projects/ml-compiler-talk/pytorch-examples/pytorch-2.0/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 103, in _fn
    return fn(*args, **kwargs)
  File "/home/david/projects/ml-compiler-talk/pytorch-examples/pytorch-2.0/lib/python3.10/site-packages/torch/_dynamo/utils.py", line 88, in time_wrapper
    r = func(*args, **kwargs)
  File "/home/david/projects/ml-compiler-talk/pytorch-examples/pytorch-2.0/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 339, in _convert_frame_assert
    return _compile(
  File "/home/david/projects/ml-compiler-talk/pytorch-examples/pytorch-2.0/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 400, in _compile
    out_code = transform_code_object(code, transform)
  File "/home/david/projects/ml-compiler-talk/pytorch-examples/pytorch-2.0/lib/python3.10/site-packages/torch/_dynamo/bytecode_transformation.py", line 341, in transform_code_object
    transformations(instructions, code_options)
  File "/home/david/projects/ml-compiler-talk/pytorch-examples/pytorch-2.0/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 387, in transform
    tracer.run()
  File "/home/david/projects/ml-compiler-talk/pytorch-examples/pytorch-2.0/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 1684, in run
    super().run()
  File "/home/david/projects/ml-compiler-talk/pytorch-examples/pytorch-2.0/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 538, in run
    and self.step()
  File "/home/david/projects/ml-compiler-talk/pytorch-examples/pytorch-2.0/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 501, in step
    getattr(self, inst.opname)(inst)
  File "/home/david/projects/ml-compiler-talk/pytorch-examples/pytorch-2.0/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 307, in wrapper
    return inner_fn(self, inst)
  File "/home/david/projects/ml-compiler-talk/pytorch-examples/pytorch-2.0/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 966, in CALL_FUNCTION
    self.call_function(fn, args, {})
  File "/home/david/projects/ml-compiler-talk/pytorch-examples/pytorch-2.0/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 435, in call_function
    self.push(fn.call_function(self, args, kwargs))
  File "/home/david/projects/ml-compiler-talk/pytorch-examples/pytorch-2.0/lib/python3.10/site-packages/torch/_dynamo/variables/misc.py", line 654, in call_function
    return self.obj.call_method(tx, self.name, args, kwargs).add_options(self)
  File "/home/david/projects/ml-compiler-talk/pytorch-examples/pytorch-2.0/lib/python3.10/site-packages/torch/_dynamo/variables/tensor.py", line 266, in call_method
    unimplemented(f"Tensor.{name}")
  File "/home/david/projects/ml-compiler-talk/pytorch-examples/pytorch-2.0/lib/python3.10/site-packages/torch/_dynamo/exc.py", line 67, in unimplemented
    raise Unsupported(msg)
torch._dynamo.exc.Unsupported: Tensor.backward
</code></pre>

ptrblck: <p>Could you post a minimal, executable code snippet to reproduce the issue, please?</p>

gengrill: <p>Sure, thanks for looking!</p>
<pre><code class="lang-auto">$ python3
import torch
import torch._dynamo as dynamo
from torch import nn

class Simple(nn.Module):
    def __init__(self, H=28, W=28, C=10):
        super(Simple, self).__init__()
        self.linear = nn.Linear(H*W, C)
    def forward(self, x):
        x = torch.flatten(x, start_dim=1)
        x = self.linear(x)
        return nn.functional.relu(x)

def generate_data(b):
    return (torch.randn(b,28,28).to(torch.float32), torch.randint(10, (b,)))

def train(model, data):
    pred = model(data[0])
    loss = nn.CrossEntropyLoss()(pred, data[1])
    loss.backward()
    return loss

model = Simple()
model_exp = dynamo.export(train, model, generate_data(16), aten_graph=True)
</code></pre>

ptrblck: <p>Thanks for the code snippet!<br>
This seems to be a known limitation as seen <a href="https://github.com/pytorch/pytorch/blob/6f034dc0b09a96c50421cf92ddb8709c59d95edf/torch/_dynamo/variables/tensor.py#L265-L266">here</a>. I couldn’t find a tracking issue for its support on GitHub, so don’t know about its current status.</p>

gengrill: <p>Okay, makes sense. Is there any other currently supported way to export the ATen training graph in Dynamo (i.e., when there is a backward pass)?</p>

a1m33: <p>Hey, I get this error as well but for numpy and using Inductor as the backend compiler. Is there an tracking issue or workaround for this by any chance?</p>

ptrblck: <p>Using <code>numpy</code> works for me if I disable the <code>backward</code> call and reuse the same code:</p>
<pre data-code-wrap="python"><code class="lang-python">import torch
import torch._dynamo as dynamo
from torch import nn

class Simple(nn.Module):
    def __init__(self, H=28, W=28, C=10):
        super(Simple, self).__init__()
        self.linear = nn.Linear(H*W, C)
    def forward(self, x):
        x = torch.flatten(x, start_dim=1)
        x = self.linear(x)
        x = np.exp(x.detach().cpu().numpy())
        return nn.functional.relu(torch.from_numpy(x))

def generate_data(b):
    return (torch.randn(b,28,28).to(torch.float32), torch.randint(10, (b,)))

def train(model, data):
    pred = model(data[0])
    loss = nn.CrossEntropyLoss()(pred, data[1])
    #loss.backward()
    return loss

model = Simple()
model_exp = dynamo.export(train, model, generate_data(16), aten_graph=True)
</code></pre>