torchinfo
torchinfo copied to clipboard
View model summaries in PyTorch!
My model requires the following inputs in the forward method: ```python def forward( self, input_ids=None, attention_mask=None, position_ids=None, head_mask=None, video_query_tokens=None, frame_hidden_state=None, frame_atts=None, past_key_values=None, use_cache=None, output_attentions=None, output_hidden_states=None, return_dict=None, is_decoder=False, ): ``` However,...
Currently, our CI test suite only tests CPU-only code. Any models that require a GPU/cuda need to be tested manually via Google Colab. A much better way to run these...
Imagine I have a module like this: ``` class Model(nn.Module): def __init__(self, generator, discriminator): super(Model, self).__init__() # Define Generator self.generator = generator # Define Discriminator self.discriminator = discriminator def forward(self,...
https://github.com/TylerYep/torchinfo/blob/73ed5687acfd6199b77fa1dcb65aa54762c1b720/torchinfo/torchinfo.py#L501 ``` >>> get_total_memory_used(["abc", "def"]) Traceback (most recent call last): File "", line 1, in File "/usr/local/lib/python3.9/dist-packages/torchinfo/torchinfo.py", line 503, in get_total_memory_used result = traverse_input_data( File "/usr/local/lib/python3.9/dist-packages/torchinfo/torchinfo.py", line 447, in traverse_input_data...
Is possible add a "write bandwidth", "read bandwidth" into `summary` function's col_names with specific input. the net inference performance also depond on not only the computation opertaion also the memory...
**Describe the bug** If I try to use summary on a model that returns a list, it will only print the output shape of the first element in the list....
**Describe the bug** Model takes into account unused layers when counting number of parameters **To Reproduce** Steps to reproduce the behavior: ``` class VanillaLeNet5(nn.Module): """ Original implementation of LeNet5 paper...
**Is your feature request related to a problem? Please describe.** Every time I want to put the specific structure of the model I use into a thesis, I will use...
**Describe the bug** For ViT, the returned total mult-adds from torchinfo.summary is much smaller than that reported in other websites. **To Reproduce** Code snippet: ```python from torchinfo import summary from...
**The bug** Hi, I see a problem with the code calculating the size of the layers: In layer_info.py line 109: ``` if hasattr(inputs[0], "size") and callable(inputs[0].size): return list(inputs[0].size()), inputs[0].element_size() ```...