Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

torchview fails with neuraloperator library #112

Open
nachigokh opened this issue Sep 23, 2024 · 0 comments
Open

torchview fails with neuraloperator library #112

nachigokh opened this issue Sep 23, 2024 · 0 comments

Comments

@nachigokh
Copy link

Describe the bug

Using torchview with the neuralop library gives the following error

TypeError: can't assign a numpy.ndarray to a torch.ComplexFloatTensor

To Reproduce

  1. Install neuraloperator, install tag 0.3.0
  2. Run the following script. You will see that the model completes a forward pass in the line y=model(x), but draw_graph fails when the same tensor is provided as input
  3. See error below

Script which reproduces error

from torchview import draw_graph
import torch
from neuralop.models import FNO3d

modesx, modesy, modesz = 8,8,8
n_layers = 4
in_channels = 3
out_channels = 3
hidden_channels = 20;
use_mlp=True;
device='cpu'

x = torch.randn(1,in_channels,32,32,10).to(device)

model =  FNO3d(modesx, modesy, modesz,n_layers=n_layers, in_channels=in_channels,out_channels=out_channels,hidden_channels=hidden_channels,
               use_mlp=use_mlp).to(device)


y = model(x)
print('Executed FNO forward pass')
#device='meta' -> no memory is consumed for visualization
model_graph = draw_graph(model, input_data=x, device=device, save_graph=True,filename='torchview.dot')
#model_graph.visual_graph

Expected behavior

torchview does not fail and produces a graph of the model

Screenshots / Text

Entire error message is below.

Executed FNO forward pass
Traceback (most recent call last):
  File "/mnt/d/virtualenv/neuralop/lib/python3.10/site-packages/torchview/torchview.py", line 256, in forward_prop
    _ = model.to(device)(*x, **kwargs)
  File "/mnt/d/virtualenv/neuralop/lib/python3.10/site-packages/torchview/recorder_tensor.py", line 146, in _module_forward_wrapper
    out = _orig_module_forward(mod, *args, **kwargs)
  File "/mnt/d/virtualenv/neuralop/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1553, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/mnt/d/virtualenv/neuralop/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1562, in _call_impl
    return forward_call(*args, **kwargs)
  File "/mnt/d/FNO/neuraloperator/neuralop/models/fno.py", line 253, in forward
    x = self.fno_blocks(x, layer_idx, output_shape=output_shape[layer_idx])
  File "/mnt/d/virtualenv/neuralop/lib/python3.10/site-packages/torchview/recorder_tensor.py", line 146, in _module_forward_wrapper
    out = _orig_module_forward(mod, *args, **kwargs)
  File "/mnt/d/virtualenv/neuralop/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1553, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/mnt/d/virtualenv/neuralop/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1562, in _call_impl
    return forward_call(*args, **kwargs)
  File "/mnt/d/FNO/neuraloperator/neuralop/layers/fno_block.py", line 195, in forward
    return self.forward_with_postactivation(x, index, output_shape)
  File "/mnt/d/FNO/neuraloperator/neuralop/layers/fno_block.py", line 208, in forward_with_postactivation
    x_fno = self.convs(x, index, output_shape=output_shape)
  File "/mnt/d/virtualenv/neuralop/lib/python3.10/site-packages/torchview/recorder_tensor.py", line 146, in _module_forward_wrapper
    out = _orig_module_forward(mod, *args, **kwargs)
  File "/mnt/d/virtualenv/neuralop/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1553, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/mnt/d/virtualenv/neuralop/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1562, in _call_impl
    return forward_call(*args, **kwargs)
  File "/mnt/d/FNO/neuraloperator/neuralop/layers/spectral_convolution.py", line 459, in forward
    out_fft[slices_x] = self._contract(x[slices_x], weight, separable=False)
  File "/mnt/d/virtualenv/neuralop/lib/python3.10/site-packages/torchview/recorder_tensor.py", line 241, in __torch_function__
    out = super().__torch_function__(func, types, args, kwargs)
  File "/mnt/d/virtualenv/neuralop/lib/python3.10/site-packages/torch/_tensor.py", line 1437, in __torch_function__
    ret = func(*args, **kwargs)
TypeError: can't assign a numpy.ndarray to a torch.ComplexFloatTensor

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

Traceback (most recent call last):
  File "/mnt/d/FNO/temp/do_torchview_fno.py", line 24, in <module>
    model_graph = draw_graph(model, input_data=x, device=device, save_graph=True,filename='torchview.dot')
  File "/mnt/d/virtualenv/neuralop/lib/python3.10/site-packages/torchview/torchview.py", line 220, in draw_graph
    forward_prop(
  File "/mnt/d/virtualenv/neuralop/lib/python3.10/site-packages/torchview/torchview.py", line 264, in forward_prop
    raise RuntimeError(
RuntimeError: Failed to run torchgraph see error message

Additional context

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

No branches or pull requests

1 participant