pytorch / executorch

On-device AI across mobile, embedded and edge for PyTorch
https://pytorch.org/executorch/
Other
1.89k stars 309 forks source link

When I use weight_norm, I get an error when exporting pte #3946

Open hgchenkai opened 3 months ago

hgchenkai commented 3 months ago

Hello, When I execute torch._export.capture_pre_autograd_graph(model, inputs) , The the following code throws an error

torch.nn.utils.weight_norm(
                    ConvTranspose1d(
                        hp.gen.upsample_initial_channel // (2 ** i),
                        hp.gen.upsample_initial_channel // (2 ** (i + 1)),
                        k,
                        u,
                        padding=(k - u) // 2)
                )

The error info:

Traceback (most recent call last):
  File "/data/home/demo/cvt_executorch.py", line 121, in <module>
    main(in_args)
  File "/data/home/demo/cvt_executorch.py", line 102, in main
    convert_decoder(generator, args)
  File "/data/home/demo/cvt_executorch.py", line 67, in convert_decoder
    save_model(model, [spk, x, har_source], "decoder", args)
  File "/data/home/demo/cvt_executorch.py", line 40, in save_model
    model = torch._export.capture_pre_autograd_graph(model, inputs)
  File "/root/anaconda3/envs/executorch/lib/python3.10/site-packages/torch/_export/__init__.py", line 152, in capture_pre_autograd_graph
    m = torch._dynamo.export(
  File "/root/anaconda3/envs/executorch/lib/python3.10/site-packages/torch/_dynamo/eval_frame.py", line 1264, in inner
    result_traced = opt_f(*args, **kwargs)
  File "/root/anaconda3/envs/executorch/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1532, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/root/anaconda3/envs/executorch/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1541, in _call_impl
    return forward_call(*args, **kwargs)
  File "/root/anaconda3/envs/executorch/lib/python3.10/site-packages/torch/_dynamo/eval_frame.py", line 416, in _fn
    return fn(*args, **kwargs)
  File "/root/anaconda3/envs/executorch/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1532, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/root/anaconda3/envs/executorch/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1541, in _call_impl
    return forward_call(*args, **kwargs)
  File "/root/anaconda3/envs/executorch/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 981, in catch_errors
    return callback(frame, cache_entry, hooks, frame_state, skip=1)
  File "/root/anaconda3/envs/executorch/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 410, in _convert_frame_assert
    return _compile(
  File "/root/anaconda3/envs/executorch/lib/python3.10/site-packages/torch/_utils_internal.py", line 70, in wrapper_function
    return function(*args, **kwargs)
  File "/root/anaconda3/envs/executorch/lib/python3.10/contextlib.py", line 79, in inner
    return func(*args, **kwds)
  File "/root/anaconda3/envs/executorch/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 703, in _compile
    guarded_code = compile_inner(code, one_graph, hooks, transform)
  File "/root/anaconda3/envs/executorch/lib/python3.10/site-packages/torch/_dynamo/utils.py", line 273, in time_wrapper
    r = func(*args, **kwargs)
  File "/root/anaconda3/envs/executorch/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 570, in compile_inner
    out_code = transform_code_object(code, transform)
  File "/root/anaconda3/envs/executorch/lib/python3.10/site-packages/torch/_dynamo/bytecode_transformation.py", line 1167, in transform_code_object
    transformations(instructions, code_options)
  File "/root/anaconda3/envs/executorch/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 172, in _fn
    return fn(*args, **kwargs)
  File "/root/anaconda3/envs/executorch/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 517, in transform
    tracer.run()
  File "/root/anaconda3/envs/executorch/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 2234, in run
    super().run()
  File "/root/anaconda3/envs/executorch/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 884, in run
    while self.step():
  File "/root/anaconda3/envs/executorch/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 799, in step
    self.dispatch_table[inst.opcode](self, inst)
  File "/root/anaconda3/envs/executorch/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 494, in wrapper
    return inner_fn(self, inst)
  File "/root/anaconda3/envs/executorch/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 1253, in CALL_FUNCTION
    self.call_function(fn, args, {})
  File "/root/anaconda3/envs/executorch/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 737, in call_function
    self.push(fn.call_function(self, args, kwargs))
  File "/root/anaconda3/envs/executorch/lib/python3.10/site-packages/torch/_dynamo/variables/functions.py", line 339, in call_function
    return super().call_function(tx, args, kwargs)
  File "/root/anaconda3/envs/executorch/lib/python3.10/site-packages/torch/_dynamo/variables/functions.py", line 293, in call_function
    return super().call_function(tx, args, kwargs)
  File "/root/anaconda3/envs/executorch/lib/python3.10/site-packages/torch/_dynamo/variables/functions.py", line 90, in call_function
    return tx.inline_user_function_return(self, [*self.self_args(), *args], kwargs)
  File "/root/anaconda3/envs/executorch/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 743, in inline_user_function_return
    return InliningInstructionTranslator.inline_call(self, fn, args, kwargs)
  File "/root/anaconda3/envs/executorch/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 2448, in inline_call
    return cls.inline_call_(parent, func, args, kwargs)
  File "/root/anaconda3/envs/executorch/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 2564, in inline_call_
    tracer.run()
  File "/root/anaconda3/envs/executorch/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 884, in run
    while self.step():
  File "/root/anaconda3/envs/executorch/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 799, in step
    self.dispatch_table[inst.opcode](self, inst)
  File "/root/anaconda3/envs/executorch/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 494, in wrapper
    return inner_fn(self, inst)
  File "/root/anaconda3/envs/executorch/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 1253, in CALL_FUNCTION
    self.call_function(fn, args, {})
  File "/root/anaconda3/envs/executorch/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 737, in call_function
    self.push(fn.call_function(self, args, kwargs))
  File "/root/anaconda3/envs/executorch/lib/python3.10/site-packages/torch/_dynamo/variables/nn_module.py", line 752, in call_function
    return variables.UserFunctionVariable(fn, source=source).call_function(
  File "/root/anaconda3/envs/executorch/lib/python3.10/site-packages/torch/_dynamo/variables/functions.py", line 293, in call_function
    return super().call_function(tx, args, kwargs)
  File "/root/anaconda3/envs/executorch/lib/python3.10/site-packages/torch/_dynamo/variables/functions.py", line 90, in call_function
    return tx.inline_user_function_return(self, [*self.self_args(), *args], kwargs)
  File "/root/anaconda3/envs/executorch/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 743, in inline_user_function_return
    return InliningInstructionTranslator.inline_call(self, fn, args, kwargs)
  File "/root/anaconda3/envs/executorch/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 2448, in inline_call
    return cls.inline_call_(parent, func, args, kwargs)
  File "/root/anaconda3/envs/executorch/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 2564, in inline_call_
    tracer.run()
  File "/root/anaconda3/envs/executorch/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 884, in run
    while self.step():
  File "/root/anaconda3/envs/executorch/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 799, in step
    self.dispatch_table[inst.opcode](self, inst)
  File "/root/anaconda3/envs/executorch/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 494, in wrapper
    return inner_fn(self, inst)
  File "/root/anaconda3/envs/executorch/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 1253, in CALL_FUNCTION
    self.call_function(fn, args, {})
  File "/root/anaconda3/envs/executorch/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 737, in call_function
    self.push(fn.call_function(self, args, kwargs))
  File "/root/anaconda3/envs/executorch/lib/python3.10/site-packages/torch/_dynamo/variables/user_defined.py", line 786, in call_function
    return self.call_method(tx, "__call__", args, kwargs)
  File "/root/anaconda3/envs/executorch/lib/python3.10/site-packages/torch/_dynamo/variables/user_defined.py", line 638, in call_method
    return UserMethodVariable(method, self, source=source).call_function(
  File "/root/anaconda3/envs/executorch/lib/python3.10/site-packages/torch/_dynamo/variables/functions.py", line 339, in call_function
    return super().call_function(tx, args, kwargs)
  File "/root/anaconda3/envs/executorch/lib/python3.10/site-packages/torch/_dynamo/variables/functions.py", line 293, in call_function
    return super().call_function(tx, args, kwargs)
  File "/root/anaconda3/envs/executorch/lib/python3.10/site-packages/torch/_dynamo/variables/functions.py", line 90, in call_function
    return tx.inline_user_function_return(self, [*self.self_args(), *args], kwargs)
  File "/root/anaconda3/envs/executorch/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 743, in inline_user_function_return
    return InliningInstructionTranslator.inline_call(self, fn, args, kwargs)
  File "/root/anaconda3/envs/executorch/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 2448, in inline_call
    return cls.inline_call_(parent, func, args, kwargs)
  File "/root/anaconda3/envs/executorch/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 2564, in inline_call_
    tracer.run()
  File "/root/anaconda3/envs/executorch/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 884, in run
    while self.step():
  File "/root/anaconda3/envs/executorch/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 799, in step
    self.dispatch_table[inst.opcode](self, inst)
  File "/root/anaconda3/envs/executorch/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 494, in wrapper
    return inner_fn(self, inst)
  File "/root/anaconda3/envs/executorch/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 1253, in CALL_FUNCTION
    self.call_function(fn, args, {})
  File "/root/anaconda3/envs/executorch/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 737, in call_function
    self.push(fn.call_function(self, args, kwargs))
  File "/root/anaconda3/envs/executorch/lib/python3.10/site-packages/torch/_dynamo/variables/builtin.py", line 946, in call_function
    return handler(tx, args, kwargs)
  File "/root/anaconda3/envs/executorch/lib/python3.10/site-packages/torch/_dynamo/variables/builtin.py", line 821, in builtin_dipatch
    rv = handler(tx, args, kwargs)
  File "/root/anaconda3/envs/executorch/lib/python3.10/site-packages/torch/_dynamo/variables/builtin.py", line 750, in call_self_handler
    result = self_handler(tx, *args, **kwargs)
  File "/root/anaconda3/envs/executorch/lib/python3.10/site-packages/torch/_dynamo/variables/builtin.py", line 1603, in call_setattr
    return obj.call_method(tx, "__setattr__", [name_var, val], {})
  File "/root/anaconda3/envs/executorch/lib/python3.10/site-packages/torch/_dynamo/variables/nn_module.py", line 851, in call_method
    return super().call_method(tx, name, args, kwargs)
  File "/root/anaconda3/envs/executorch/lib/python3.10/site-packages/torch/_dynamo/variables/user_defined.py", line 569, in call_method
    return self.method_setattr_standard(tx, *args, **kwargs)
  File "/root/anaconda3/envs/executorch/lib/python3.10/site-packages/torch/_dynamo/variables/user_defined.py", line 654, in method_setattr_standard
    unimplemented(f"setattr({self}, {name}, ...)")
  File "/root/anaconda3/envs/executorch/lib/python3.10/site-packages/torch/_dynamo/exc.py", line 212, in unimplemented
    raise Unsupported(msg)
torch._dynamo.exc.Unsupported: setattr(UnspecializedNNModuleVariable(ConvTranspose1d), weight, ...)

from user code:
   File "/data/home/demo/cvt_utils.py", line 53, in forward
    return self.model.inference(spk, x, har_source)
  File "/data/home/demo/generator.py", line 181, in inference
    x = self.ups[i](x)
  File "/root/anaconda3/envs/executorch/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1571, in _call_impl
    args_result = hook(self, args)
  File "/root/anaconda3/envs/executorch/lib/python3.10/site-packages/torch/nn/utils/weight_norm.py", line 65, in __call__
    setattr(module, self.name, self.compute_weight(module))

Set TORCH_LOGS="+dynamo" and TORCHDYNAMO_VERBOSE=1 for more information

When I replaced torch.nn.utils.weight_norm with torch.nn.utils.parametrizations.weight_norm, torch._export.capture_pre_autograd_graph executed successfully.But when I execute to edge.to_backend(XnnpackPartitioner()), I get a new error:

Traceback (most recent call last):
  File "/root/anaconda3/envs/executorch/lib/python3.10/site-packages/torch/fx/passes/infra/pass_manager.py", line 270, in __call__
    res = fn(module)
  File "/root/anaconda3/envs/executorch/lib/python3.10/site-packages/torch/fx/passes/infra/pass_base.py", line 40, in __call__
    res = self.call(graph_module)
  File "/root/anaconda3/envs/executorch/lib/python3.10/site-packages/executorch/backends/xnnpack/passes/conv1d_unsqueeze_pass.py", line 123, in call
    raise AssertionError(
AssertionError: Expected op for convolution weight node to be a get_attr node or a parameter

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/data/home/demo/cvt_executorch.py", line 121, in <module>
    main(in_args)
  File "/data/home/demo/cvt_executorch.py", line 102, in main
    convert_decoder(generator, args)
  File "/data/home/demo/cvt_executorch.py", line 67, in convert_decoder
    save_model(model, [spk, x, har_source], "decoder", args)
  File "/data/home/demo/cvt_executorch.py", line 52, in save_model
    edge = edge.to_backend(XnnpackPartitioner())
  File "/root/anaconda3/envs/executorch/lib/python3.10/site-packages/executorch/exir/program/_program.py", line 1168, in to_backend
    new_edge_programs[name] = to_backend(program, partitioner)
  File "/root/anaconda3/envs/executorch/lib/python3.10/functools.py", line 878, in wrapper
    return dispatch(args[0].__class__)(*args, **kw)
  File "/root/anaconda3/envs/executorch/lib/python3.10/site-packages/executorch/exir/backend/backend_api.py", line 384, in _
    tagged_graph_module = _partition_and_lower(
  File "/root/anaconda3/envs/executorch/lib/python3.10/site-packages/executorch/exir/backend/backend_api.py", line 299, in _partition_and_lower
    partitioned_module = _partition_and_lower_one_graph_module(
  File "/root/anaconda3/envs/executorch/lib/python3.10/site-packages/executorch/exir/backend/backend_api.py", line 230, in _partition_and_lower_one_graph_module
    lowered_submodule = to_backend(
  File "/root/anaconda3/envs/executorch/lib/python3.10/functools.py", line 878, in wrapper
    return dispatch(args[0].__class__)(*args, **kw)
  File "/root/anaconda3/envs/executorch/lib/python3.10/site-packages/executorch/exir/backend/backend_api.py", line 114, in _
    preprocess_result: PreprocessResult = cls.preprocess(
  File "/root/anaconda3/envs/executorch/lib/python3.10/site-packages/executorch/backends/xnnpack/xnnpack_preprocess.py", line 122, in preprocess
    ep = XNNPACKPassManager(ep, passes=passes).transform()
  File "/root/anaconda3/envs/executorch/lib/python3.10/site-packages/executorch/backends/xnnpack/passes/__init__.py", line 86, in transform
    ep = _transform(ep, transform_pass)
  File "/root/anaconda3/envs/executorch/lib/python3.10/site-packages/executorch/exir/program/_program.py", line 179, in _transform
    res = pm(self.graph_module)
  File "/root/anaconda3/envs/executorch/lib/python3.10/site-packages/torch/fx/passes/infra/pass_manager.py", line 296, in __call__
    raise Exception(msg) from e  # noqa: TRY002
Exception: An error occurred when running the 'Conv1dUnsqueezePass' pass after the following passes: []

How should I use weight_norm?Is there something wrong with the way I use it? Please help me,thanks

### Tasks
digantdesai commented 3 months ago

Not sure if we support parameterized weights for conv in XNNPACK delegate. Is it possible to normalize weights ahead of time before export, I guess?

Alternatively, can you export and run successfully without xnnpackpartitioner?

hgchenkai commented 3 months ago

Not sure if we support parameterized weights for conv in XNNPACK delegate. Is it possible to normalize weights ahead of time before export, I guess?

Alternatively, can you export and run successfully without xnnpackpartitioner?

Thanks for the reply. Yes, I can run sucessfully without xnnpackpartitioner