Is it possible to visualize model?

[JunsuAndrewLee]

I tried to visualize the model to export as a png file using various pytorch model visualization tools such as torchviz, hiddenlayer, approaching backward and forward propagation. I also failed to export the model to onnx file to use netron. Is there anybody who succeeded visualizing or has some idea with the following error?

1. tried to export the model to an onnx file

data = AtomicData.from_ase(atoms=atoms, r_max=self.r_max)
for k in AtomicDataDict.ALL_ENERGY_KEYS:
        if k in data:
            del data[k]
data = self.transform(data)
data = data.to(self.device)
data = AtomicData.to_AtomicDataDict(data)
# predict + extract data
out = self.model(data)

input_names = ['Sentence']
output_names = ['yhat']
torch.onnx.export(self.model, data, 'nequip.onnx.onnx', input_names=input_names, output_names=output_names)

/home/users/user/anaconda3/envs/nequip1/lib/python3.10/site-packages/torch/onnx/utils.py:847: UserWarning: no signature found for <torch.ScriptMethod object at 0x1502230723e0>, skipping _decide_input_format
  warnings.warn(f"{e}, skipping _decide_input_format")
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/users/user/anaconda3/envs/nequip1/lib/python3.10/site-packages/nequip/ase/nequip_calculator.py", line 125, in calculate
    torch.onnx.export(self.model, data, 'nequip.onnx.onnx', input_names=input_names, output_names=output_names)
  File "/home/users/user/anaconda3/envs/nequip1/lib/python3.10/site-packages/torch/onnx/utils.py", line 516, in export
    _export(
  File "/home/users/user/anaconda3/envs/nequip1/lib/python3.10/site-packages/torch/onnx/utils.py", line 1613, in _export
    graph, params_dict, torch_out = _model_to_graph(
  File "/home/users/user/anaconda3/envs/nequip1/lib/python3.10/site-packages/torch/onnx/utils.py", line 1135, in _model_to_graph
    graph, params, torch_out, module = _create_jit_graph(model, args)
  File "/home/users/user/anaconda3/envs/nequip1/lib/python3.10/site-packages/torch/onnx/utils.py", line 979, in _create_jit_graph
    _check_flatten_did_not_remove(args, flattened_args)
  File "/home/users/user/anaconda3/envs/nequip1/lib/python3.10/site-packages/torch/onnx/utils.py", line 967, in _check_flatten_did_not_remove
    raise ValueError(
ValueError: args contained 6 None's after flattening. When exporting a ScriptModule or ScriptFunction, no args may be None because that breaks type propagation.

2. hiddenlayer

import hiddenlayer as hl
graph = hl.build_graph(self.model, data)
graph.theme = hl.graph.THEMES['blue'].copy()
graph.save('rnn_hiddenlayer', format='png')


Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/users/user/anaconda3/envs/nequip1/lib/python3.10/site-packages/nequip/ase/nequip_calculator.py", line 125, in calculate
    graph = hl.build_graph(self.model, data)
  File "/home/users/user/anaconda3/envs/nequip1/lib/python3.10/site-packages/hiddenlayer/graph.py", line 143, in build_graph
    import_graph(g, model, args)
  File "/home/users/user/anaconda3/envs/nequip1/lib/python3.10/site-packages/hiddenlayer/pytorch_builder.py", line 70, in import_graph
    trace, out = torch.jit._get_trace_graph(model, args)
  File "/home/users/user/anaconda3/envs/nequip1/lib/python3.10/site-packages/torch/jit/_trace.py", line 1296, in _get_trace_graph
    outs = ONNXTracedModule(
  File "/home/users/user/anaconda3/envs/nequip1/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1511, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/users/user/anaconda3/envs/nequip1/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1520, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/users/user/anaconda3/envs/nequip1/lib/python3.10/site-packages/torch/jit/_trace.py", line 138, in forward
    graph, out = torch._C._create_graph_by_tracing(
  File "/home/users/user/anaconda3/envs/nequip1/lib/python3.10/site-packages/torch/jit/_trace.py", line 129, in wrapper
    outs.append(self.inner(*trace_inputs))
  File "/home/users/user/anaconda3/envs/nequip1/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1511, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/users/user/anaconda3/envs/nequip1/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1520, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/users/user/anaconda3/envs/nequip1/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1490, in _slow_forward
    return self.forward(*input, **kwargs)
RuntimeError: Tried to trace <__torch__.nequip.nn._rescale.___torch_mangle_0.RescaleOutput object at 0x13729e60> but it is not part of the active trace. Modules that are called during a trace must be registered as submodules of the thing being traced.

3. torchviz

from torchviz import make_dot
make_dot(out, params=dict(list(self.model.named_parameters()))).render("rnn_torchviz", format="png")

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/users/user/anaconda3/envs/nequip1/lib/python3.10/site-packages/nequip/ase/nequip_calculator.py", line 123, in calculate
    make_dot(out, params=dict(list(self.model.named_parameters()))).render("rnn_torchviz", format="png")
  File "/home/users/user/anaconda3/envs/nequip1/lib/python3.10/site-packages/torchviz/dot.py", line 163, in make_dot
    add_base_tensor(var)
  File "/home/users/user/anaconda3/envs/nequip1/lib/python3.10/site-packages/torchviz/dot.py", line 146, in add_base_tensor
    if var in seen:
TypeError: unhashable type: 'dict'

from torchviz import make_dot
        make_dot(data, params=dict(list(self.model.named_parameters()))).render("rnn_torchviz", format="png")

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/users/user/anaconda3/envs/nequip1/lib/python3.10/site-packages/nequip/ase/nequip_calculator.py", line 123, in calculate
    make_dot(data, params=dict(list(self.model.named_parameters()))).render("rnn_torchviz", format="png")
  File "/home/users/user/anaconda3/envs/nequip1/lib/python3.10/site-packages/torchviz/dot.py", line 163, in make_dot
    add_base_tensor(var)
  File "/home/users/user/anaconda3/envs/nequip1/lib/python3.10/site-packages/torchviz/dot.py", line 146, in add_base_tensor
    if var in seen:
TypeError: unhashable type: 'dict'

Thank you for reading this question.