mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-06 12:20:52 +01:00
8c13e96be2
Log tensor size information with the `graph_sizes` logging artifact, as part of the model x-ray feature requests. Typically can be combined with `graph_code`.
Sample:
```python
import torch
def fn(a, b, c, d):
return (a + b) @ (c + d)
opt_fn = torch.compile(fn, backend="eager", dynamic=False)
opt_fn(torch.randn(10, 20), torch.randn(1, 20), torch.randn(20, 15), torch.randn(1, 15))
opt_fn(torch.randn(5, 2), torch.randn(1, 2), torch.randn(2, 4), torch.randn(1, 4))
```
Output:
```shell
$ TORCH_LOGS="graph_sizes,graph_code" python playground8.py
[2023-07-06 01:42:39,093] torch._dynamo.output_graph.__graph_code: [DEBUG] TRACED GRAPH
===== __compiled_fn_0 =====
<eval_with_key>.0 class GraphModule(torch.nn.Module):
def forward(self, L_a_ : torch.Tensor, L_b_ : torch.Tensor, L_c_ : torch.Tensor, L_d_ : torch.Tensor):
l_a_ = L_a_
l_b_ = L_b_
l_c_ = L_c_
l_d_ = L_d_
# File: playground8.py:66, code: return (a + b) @ (c + d)
add = l_a_ + l_b_; l_a_ = l_b_ = None
add_1 = l_c_ + l_d_; l_c_ = l_d_ = None
matmul = add @ add_1; add = add_1 = None
return (matmul,)
[2023-07-06 01:42:39,093] torch._dynamo.output_graph.__graph_sizes: [DEBUG] TRACED GRAPH TENSOR SIZES
===== __compiled_fn_0 =====
l_a_: (10, 20)
l_b_: (1, 20)
l_c_: (20, 15)
l_d_: (1, 15)
add: (10, 20)
add_1: (20, 15)
matmul: (10, 15)
[2023-07-06 01:42:39,198] torch._dynamo.output_graph.__graph_code: [DEBUG] TRACED GRAPH
===== __compiled_fn_1 =====
<eval_with_key>.1 class GraphModule(torch.nn.Module):
def forward(self, s0 : torch.SymInt, s1 : torch.SymInt, L_a_ : torch.Tensor, L_b_ : torch.Tensor, s4 : torch.SymInt, L_c_ : torch.Tensor, L_d_ : torch.Tensor):
l_a_ = L_a_
l_b_ = L_b_
l_c_ = L_c_
l_d_ = L_d_
# File: playground8.py:66, code: return (a + b) @ (c + d)
add = l_a_ + l_b_; l_a_ = l_b_ = None
add_1 = l_c_ + l_d_; l_c_ = l_d_ = None
matmul = add @ add_1; add = add_1 = None
return (matmul,)
[2023-07-06 01:42:39,198] torch._dynamo.output_graph.__graph_sizes: [DEBUG] TRACED GRAPH TENSOR SIZES
===== __compiled_fn_1 =====
l_a_: (s0, s1)
l_a_ (concrete): (5, 2)
l_b_: (1, s1)
l_b_ (concrete): (1, 2)
l_c_: (s1, s4)
l_c_ (concrete): (2, 4)
l_d_: (1, s4)
l_d_ (concrete): (1, 4)
add: (s0, s1)
add (concrete): (5, 2)
add_1: (s1, s4)
add_1 (concrete): (2, 4)
matmul: (s0, s4)
matmul (concrete): (5, 4)
```
Pull Request resolved: https://github.com/pytorch/pytorch/pull/104672
Approved by: https://github.com/ezyang
|
||
|---|---|---|
| .. | ||
| _awaits | ||
| _C | ||
| _C_flatbuffer | ||
| _custom_op | ||
| _decomp | ||
| _dispatch | ||
| _dynamo | ||
| _export | ||
| _functorch | ||
| _higher_order_ops | ||
| _inductor | ||
| _lazy | ||
| _logging | ||
| _prims | ||
| _prims_common | ||
| _refs | ||
| _subclasses | ||
| amp | ||
| ao | ||
| autograd | ||
| backends | ||
| compiler | ||
| contrib | ||
| cpu | ||
| csrc | ||
| cuda | ||
| distributed | ||
| distributions | ||
| fft | ||
| func | ||
| futures | ||
| fx | ||
| jit | ||
| legacy | ||
| lib | ||
| linalg | ||
| masked | ||
| monitor | ||
| mps | ||
| multiprocessing | ||
| nested | ||
| nn | ||
| onnx | ||
| optim | ||
| package | ||
| profiler | ||
| quantization | ||
| signal | ||
| sparse | ||
| special | ||
| testing | ||
| utils | ||
| __config__.py | ||
| __future__.py | ||
| __init__.py | ||
| _appdirs.py | ||
| _classes.py | ||
| _compile.py | ||
| _deploy.py | ||
| _guards.py | ||
| _jit_internal.py | ||
| _linalg_utils.py | ||
| _lobpcg.py | ||
| _lowrank.py | ||
| _meta_registrations.py | ||
| _namedtensor_internals.py | ||
| _ops.py | ||
| _python_dispatcher.py | ||
| _sources.py | ||
| _storage_docs.py | ||
| _tensor_docs.py | ||
| _tensor_str.py | ||
| _tensor.py | ||
| _torch_docs.py | ||
| _utils_internal.py | ||
| _utils.py | ||
| _VF.py | ||
| _vmap_internals.py | ||
| _weights_only_unpickler.py | ||
| abi-check.cpp | ||
| CMakeLists.txt | ||
| custom_class_detail.h | ||
| custom_class.h | ||
| extension.h | ||
| functional.py | ||
| hub.py | ||
| library.h | ||
| library.py | ||
| overrides.py | ||
| py.typed | ||
| quasirandom.py | ||
| random.py | ||
| README.txt | ||
| return_types.py | ||
| script.h | ||
| serialization.py | ||
| storage.py | ||
| torch_version.py | ||
| types.py | ||
| version.py.tpl | ||
Note [TH abstraction violation] ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ TH/THC provide some hpp headers, which are proper C++ headers rather than C headers. These headers serve double duty as *internal implementation detail* headers, whose contents should largely not be used by external clients. Ideally, we would not install these headers at all; instead, you should use public functions (in headers like `THTensor.h`, NOT `THTensor.hpp`) to manipulate these structs. However, there are a few places in torch/csrc where we violate this abstraction. They are marked with a pointer to this note. Each of those sites will have to be refactored when we refactor the guts of THTensor and related structures.