diff --git a/tests/torchinfo_test.py b/tests/torchinfo_test.py index e91ff82..73a9477 100644 --- a/tests/torchinfo_test.py +++ b/tests/torchinfo_test.py @@ -467,6 +467,37 @@ def test_multiple_input_tensor_list() -> None: assert metrics.input_size == [torch.Size([1, 300]), torch.Size([1, 300])] +def test_sparse_input() -> None: + # https://github.com/TylerYep/torchinfo/issues/150 + class GraphConvolution(nn.Module): + def __init__(self, in_features: int, out_features: int) -> None: + super().__init__() + self.weight = nn.Parameter(torch.randn(in_features, out_features)) + + def forward(self, x: torch.Tensor, adj: torch.Tensor) -> torch.Tensor: + return torch.spmm(adj, x @ self.weight) + + class GCN(nn.Module): + def __init__(self) -> None: + super().__init__() + self.gc1 = GraphConvolution(16, 8) + self.gc2 = GraphConvolution(8, 4) + + def forward(self, x: torch.Tensor, adj: torch.Tensor) -> torch.Tensor: + out: torch.Tensor = self.gc2(torch.relu(self.gc1(x, adj)), adj) + return out + + features = torch.randn(10, 16) + adj = torch.sparse_coo_tensor( + torch.tensor([[0, 1, 2], [1, 2, 0]]), torch.ones(3), (10, 10) + ).coalesce() + + metrics = summary(GCN(), input_data=[features, adj], verbose=0) + + assert metrics.total_params == 16 * 8 + 8 * 4 == 160 + assert metrics.trainable_params == 160 + + def test_namedtuple() -> None: model = NamedTuple() input_size = [(2, 1, 28, 28), (2, 1, 28, 28)] diff --git a/torchinfo/torchinfo.py b/torchinfo/torchinfo.py index 46b4502..682d0fb 100644 --- a/torchinfo/torchinfo.py +++ b/torchinfo/torchinfo.py @@ -507,15 +507,45 @@ def get_input_data_sizes(data: Any) -> Any: ) +def get_tensor_memory(data: torch.Tensor) -> int: + """Calculates the memory of a single tensor, including sparse layouts.""" + try: + storage = ( + data.untyped_storage() + if hasattr(data, "untyped_storage") + else data.storage() # type: ignore[no-untyped-call] + ) + return sys.getsizeof(storage) + except NotImplementedError: + # Sparse tensors (COO, CSR, ...) expose no single storage, + # so sum the member tensors instead. Accessors differ per + # layout, so probe the known ones and keep what succeeds. + members: list[torch.Tensor] = [data.values()] + for name in ( + "indices", + "crow_indices", + "col_indices", + "ccol_indices", + "row_indices", + ): + getter = getattr(data, name, None) + if getter is None: + continue + try: + member = getter() + except NotImplementedError: + continue + except RuntimeError: + continue + members.append(member) + return sum(get_tensor_memory(member) for member in members) + + def get_total_memory_used(data: CORRECTED_INPUT_DATA_TYPE) -> int: """Calculates the total memory of all tensors stored in data.""" result = traverse_input_data( data, - action_fn=lambda data: sys.getsizeof( - data.untyped_storage() - if hasattr(data, "untyped_storage") - else data.storage() - ), + action_fn=get_tensor_memory, aggregate_fn=( # We don't need the dictionary keys in this case lambda data: (