├── docs ├── images │ ├── graph_drawing.png.md │ ├── toy_resnet.png │ ├── draw_graph1.png │ ├── draw_graph2.png │ ├── draw_graph3.png │ ├── draw_graph4.png │ ├── draw_graph5.png │ ├── draw_graph6.png │ ├── draw_graph7.png │ ├── draw_graph8.png │ └── many_linear_layers.png ├── reference │ ├── model_tools.md │ ├── symbolic_model.md │ ├── graph_algorithms.md │ ├── optimize_module_calls.md │ ├── useful_layers.md │ ├── add_to_graph.md │ └── symbolic_data.md ├── benchmarks.md ├── quick_start.md └── advanced_topics.md ├── .gitignore ├── examples ├── common.py ├── encoder_decoder.py ├── vgg.py ├── resnet.py └── lstm.py ├── pytorch_symbolic ├── config.py ├── __init__.py ├── model_tools.py ├── symbolic_api_2.py ├── functions_utility.py ├── useful_layers.py ├── code_generator.py ├── graph_algorithms.py ├── symbolic_model.py └── symbolic_data.py ├── check.sh ├── .github └── workflows │ ├── check-code-quality.yaml │ ├── run-notebook.yaml │ ├── install-from-pypi-run.yaml │ ├── check-with-pytest.yaml │ └── install-local-and-run.yaml ├── LICENSE.txt ├── benchmarks ├── bench1 │ ├── define_models.py │ ├── run_one.py │ └── draw_results.py ├── README.md ├── bench2 │ ├── define_models.py │ ├── run_one.py │ └── draw_results.py └── tagged_collection.py ├── tests ├── test_raising_assertions.py ├── test_iterating_and_slicing.py ├── test_nontorch_graphs.py ├── test_creating_extreme.py ├── test_docs_automated.py ├── test_creating_examples.py ├── test_creating_examples_detached.py ├── test_reusing_nodes.py ├── test_docs_manual.py ├── test_creating_multi_in_out.py ├── test_add_to_graph.py ├── test_creating_models.py └── test_basic_ops.py ├── pyproject.toml └── README.md /docs/images/graph_drawing.png.md: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /docs/images/toy_resnet.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/szmikler/pytorch-symbolic/HEAD/docs/images/toy_resnet.png -------------------------------------------------------------------------------- /docs/images/draw_graph1.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/szmikler/pytorch-symbolic/HEAD/docs/images/draw_graph1.png -------------------------------------------------------------------------------- /docs/images/draw_graph2.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/szmikler/pytorch-symbolic/HEAD/docs/images/draw_graph2.png -------------------------------------------------------------------------------- /docs/images/draw_graph3.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/szmikler/pytorch-symbolic/HEAD/docs/images/draw_graph3.png -------------------------------------------------------------------------------- /docs/images/draw_graph4.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/szmikler/pytorch-symbolic/HEAD/docs/images/draw_graph4.png -------------------------------------------------------------------------------- /docs/images/draw_graph5.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/szmikler/pytorch-symbolic/HEAD/docs/images/draw_graph5.png -------------------------------------------------------------------------------- /docs/images/draw_graph6.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/szmikler/pytorch-symbolic/HEAD/docs/images/draw_graph6.png -------------------------------------------------------------------------------- /docs/images/draw_graph7.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/szmikler/pytorch-symbolic/HEAD/docs/images/draw_graph7.png -------------------------------------------------------------------------------- /docs/images/draw_graph8.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/szmikler/pytorch-symbolic/HEAD/docs/images/draw_graph8.png -------------------------------------------------------------------------------- /docs/images/many_linear_layers.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/szmikler/pytorch-symbolic/HEAD/docs/images/many_linear_layers.png -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | **/__pycache__ 2 | **/build/ 3 | **/dist/ 4 | **/*.egg-info 5 | 6 | .ipynb_checkpoints 7 | *.sh 8 | !check.sh 9 | 10 | .idea 11 | .coverage 12 | .obsidian 13 | devel 14 | benchmarks/logs 15 | 16 | MNIST 17 | -------------------------------------------------------------------------------- /docs/reference/model_tools.md: -------------------------------------------------------------------------------- 1 | # model_tools 2 | 3 | ::: pytorch_symbolic.model_tools 4 | options: 5 | show_source: false 6 | heading_level: 2 7 | show_root_heading: true 8 | members_order: source 9 | show_object_full_path: false 10 | docstring_section_style: table 11 | show_signature_annotations: true 12 | separate_signature: true 13 | annotations_path: brief 14 | merge_init_into_class: true 15 | show_root_full_path: true 16 | -------------------------------------------------------------------------------- /docs/reference/symbolic_model.md: -------------------------------------------------------------------------------- 1 | # SymbolicModel 2 | 3 | ::: pytorch_symbolic.SymbolicModel 4 | options: 5 | show_source: false 6 | heading_level: 2 7 | show_root_heading: true 8 | members_order: source 9 | show_object_full_path: false 10 | docstring_section_style: table 11 | show_signature_annotations: true 12 | separate_signature: true 13 | annotations_path: brief 14 | merge_init_into_class: true 15 | show_root_full_path: true 16 | -------------------------------------------------------------------------------- /docs/reference/graph_algorithms.md: -------------------------------------------------------------------------------- 1 | # graph_algorithms 2 | 3 | ::: pytorch_symbolic.graph_algorithms 4 | options: 5 | show_source: false 6 | heading_level: 2 7 | show_root_heading: true 8 | members_order: source 9 | show_object_full_path: false 10 | docstring_section_style: table 11 | show_signature_annotations: true 12 | separate_signature: true 13 | annotations_path: brief 14 | merge_init_into_class: true 15 | show_root_full_path: true 16 | -------------------------------------------------------------------------------- /docs/reference/optimize_module_calls.md: -------------------------------------------------------------------------------- 1 | # optimize_module_calls 2 | 3 | ::: pytorch_symbolic.optimize_module_calls 4 | options: 5 | show_source: false 6 | heading_level: 2 7 | show_root_heading: true 8 | members_order: source 9 | show_object_full_path: false 10 | docstring_section_style: table 11 | show_signature_annotations: true 12 | separate_signature: true 13 | annotations_path: brief 14 | merge_init_into_class: true 15 | show_root_full_path: true 16 | -------------------------------------------------------------------------------- /docs/reference/useful_layers.md: -------------------------------------------------------------------------------- 1 | # useful_layers 2 | 3 | Collection of simple ``torch.nn.Modules``. 4 | 5 | Some of them might be useful when building models using Pytorch Symbolic. 6 | 7 | ::: pytorch_symbolic.useful_layers 8 | options: 9 | show_source: false 10 | heading_level: 2 11 | show_root_heading: true 12 | members_order: source 13 | show_object_full_path: false 14 | docstring_section_style: table 15 | show_signature_annotations: true 16 | separate_signature: true 17 | annotations_path: brief 18 | merge_init_into_class: true 19 | show_root_full_path: true 20 | -------------------------------------------------------------------------------- /examples/common.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2022 Szymon Mikler 2 | 3 | from torch import nn 4 | 5 | from pytorch_symbolic import useful_layers 6 | 7 | 8 | def classifier(flow, n_classes, pooling="avgpool"): 9 | if pooling == "catpool": 10 | maxp = flow(nn.MaxPool2d(kernel_size=(flow.H, flow.W))) 11 | avgp = flow(nn.AvgPool2d(kernel_size=(flow.H, flow.W))) 12 | flow = maxp(useful_layers.ConcatLayer(dim=1), avgp)(nn.Flatten()) 13 | if pooling == "avgpool": 14 | flow = flow(nn.AvgPool2d(kernel_size=(flow.H, flow.W)))(nn.Flatten()) 15 | if pooling == "maxpool": 16 | flow = flow(nn.MaxPool2d(kernel_size=(flow.H, flow.W)))(nn.Flatten()) 17 | return flow(nn.Linear(flow.features, n_classes)) 18 | -------------------------------------------------------------------------------- /pytorch_symbolic/config.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2022 Szymon Mikler 2 | 3 | import os 4 | 5 | from .symbolic_api_2 import enable_symbolic_api_2_for_new_modules 6 | 7 | 8 | def read_from_env(name, default): 9 | if name in os.environ: 10 | value = eval(os.environ[name], {}, {}) 11 | else: 12 | value = default 13 | return value 14 | 15 | 16 | # Constants 17 | CODEGEN_BY_DEFAULT = read_from_env("PYTORCH_SYMBOLIC_CODEGEN_BY_DEFAULT", True) 18 | CODEGEN_MIN_LOOP_LENGTH = read_from_env("PYTORCH_SYMBOLIC_CODEGEN_MIN_LOOP_LENGTH", 50) 19 | API_2_ENABLED_BY_DEFAULT = read_from_env("PYTORCH_SYMBOLIC_API_2_ENABLED_BY_DEFAULT", True) 20 | 21 | if API_2_ENABLED_BY_DEFAULT: 22 | enable_symbolic_api_2_for_new_modules() 23 | -------------------------------------------------------------------------------- /check.sh: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2022 Szymon Mikler 2 | 3 | echo "Running test suite for pytorch-symbolic!" 4 | echo "Optional dependencies are required. Install using 'pip install pytorch-symbolic[full]'" 5 | echo "" 6 | 7 | echo "Running isort..." 8 | isort . 9 | 10 | echo "Running black..." 11 | black . 12 | 13 | echo "Running flake8..." 14 | flake8 . 15 | 16 | echo "Running mypy..." 17 | mypy . 18 | 19 | echo "Running pytest with the default settings..." 20 | pytest --no-header . || exit 21 | 22 | echo "Running pytest with codegen active..." 23 | export PYTORCH_SYMBOLIC_CODEGEN_MIN_LOOP_LENGTH=2 24 | pytest --no-header . || exit 25 | 26 | echo "Running pytest with codegen disabled..." 27 | export PYTORCH_SYMBOLIC_CODEGEN_BY_DEFAULT=False 28 | pytest --no-header . || exit 29 | -------------------------------------------------------------------------------- /pytorch_symbolic/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2022 Szymon Mikler 2 | 3 | 4 | from . import config, graph_algorithms, model_tools, symbolic_data, useful_layers 5 | from .functions_utility import add_to_graph 6 | from .symbolic_api_2 import SymbolicAPI2ContextManager, optimize_module_calls 7 | from .symbolic_data import CustomInput, Input 8 | from .symbolic_model import SymbolicModel 9 | 10 | __license__ = "MIT" 11 | __version__ = "1.1.1" 12 | __author__ = "Szymon Mikler" 13 | 14 | __all__ = [ 15 | "Input", 16 | "CustomInput", 17 | "SymbolicModel", 18 | "add_to_graph", 19 | "config", 20 | "graph_algorithms", 21 | "model_tools", 22 | "optimize_module_calls", 23 | "symbolic_data", 24 | "useful_layers", 25 | "SymbolicAPI2ContextManager", 26 | ] 27 | -------------------------------------------------------------------------------- /.github/workflows/check-code-quality.yaml: -------------------------------------------------------------------------------- 1 | name: check code quality 2 | 3 | on: 4 | push: 5 | branches: [ "main" ] 6 | pull_request: 7 | branches: [ "main" ] 8 | 9 | permissions: 10 | contents: read 11 | 12 | jobs: 13 | build: 14 | 15 | runs-on: ubuntu-22.04 16 | 17 | steps: 18 | - uses: actions/checkout@v4 19 | - name: Set up Python 3.10 20 | uses: actions/setup-python@v5 21 | with: 22 | python-version: "3.10" 23 | - name: Install dependencies 24 | run: | 25 | python -m pip install --upgrade pip 26 | pip install ruff pyright 27 | - name: Run ruff linting 28 | run: ruff check . 29 | - name: Run ruff formatting 30 | run: ruff format . --check 31 | # - name: Run pyright 32 | # run: pyright pytorch_symbolic 33 | 34 | -------------------------------------------------------------------------------- /.github/workflows/run-notebook.yaml: -------------------------------------------------------------------------------- 1 | name: run notebook 2 | 3 | on: 4 | push: 5 | branches: [ "main" ] 6 | pull_request: 7 | branches: [ "main" ] 8 | 9 | permissions: 10 | contents: read 11 | 12 | jobs: 13 | build: 14 | 15 | runs-on: ubuntu-22.04 16 | strategy: 17 | matrix: 18 | python-version: [ "3.7", "3.8", "3.9", "3.10", "3.11", "3.12", "3.13" ] 19 | 20 | steps: 21 | - uses: actions/checkout@v4 22 | - name: Set up Python ${{ matrix.python-version }} 23 | uses: actions/setup-python@v5 24 | with: 25 | python-version: ${{ matrix.python-version }} 26 | - name: Install dependencies 27 | run: | 28 | python -m pip install --upgrade pip 29 | pip install torch numpy jupyter 30 | - name: Convert notebook to python 31 | run: | 32 | jupyter nbconvert introduction.ipynb --to python 33 | - name: Run script 34 | run: | 35 | ipython introduction.py 36 | -------------------------------------------------------------------------------- /LICENSE.txt: -------------------------------------------------------------------------------- 1 | Copyright 2022 Szymon Mikler 2 | 3 | Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions: 4 | 5 | The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software. 6 | 7 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. -------------------------------------------------------------------------------- /.github/workflows/install-from-pypi-run.yaml: -------------------------------------------------------------------------------- 1 | name: install from pypi 2 | 3 | on: 4 | push: 5 | branches: [ "main" ] 6 | pull_request: 7 | branches: [ "main" ] 8 | 9 | jobs: 10 | build: 11 | strategy: 12 | matrix: 13 | python-version: [ 14 | "3.7", 15 | "3.8", 16 | "3.9", 17 | "3.10", 18 | "3.11", 19 | "3.12", 20 | "3.13" , 21 | ] 22 | os: [ 23 | "ubuntu-22.04", 24 | "windows-latest", 25 | ] 26 | 27 | runs-on: ${{ matrix.os }} 28 | 29 | steps: 30 | - name: Set up Python ${{ matrix.python-version }} 31 | uses: actions/setup-python@v5 32 | with: 33 | python-version: ${{ matrix.python-version }} 34 | - name: Install the library and test basic imports 35 | run: | 36 | pip install pytorch-symbolic 37 | pip show pytorch-symbolic 38 | 39 | python -c "from pytorch_symbolic import SymbolicModel, Input, useful_layers, graph_algorithms" 40 | -------------------------------------------------------------------------------- /.github/workflows/check-with-pytest.yaml: -------------------------------------------------------------------------------- 1 | name: run tests 2 | 3 | on: 4 | push: 5 | branches: [ "main" ] 6 | pull_request: 7 | branches: [ "main" ] 8 | 9 | jobs: 10 | build: 11 | strategy: 12 | matrix: 13 | python-version: [ 14 | "3.7", 15 | "3.8", 16 | "3.9", 17 | "3.10", 18 | "3.11", 19 | "3.12", 20 | "3.13" , 21 | ] 22 | os: [ 23 | "ubuntu-22.04", 24 | "windows-latest", 25 | ] 26 | 27 | runs-on: ${{ matrix.os }} 28 | 29 | steps: 30 | - uses: actions/checkout@v4 31 | - name: Set up Python ${{ matrix.python-version }} 32 | uses: actions/setup-python@v5 33 | with: 34 | python-version: ${{ matrix.python-version }} 35 | - name: Install dependencies 36 | run: | 37 | pip install . 38 | pip install pytest 39 | pip install networkx matplotlib scipy 40 | pip install torch numpy 41 | 42 | - name: Test with pytest 43 | run: pytest . 44 | -------------------------------------------------------------------------------- /benchmarks/bench1/define_models.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2022 Szymon Mikler 2 | 3 | from torch import nn 4 | 5 | from pytorch_symbolic import Input, SymbolicModel 6 | 7 | 8 | class Vanilla(nn.Module): 9 | def __init__(self, layers): 10 | super().__init__() 11 | self.layers = layers 12 | for idx, layer in enumerate(layers): 13 | self.add_module(str(idx), layer) 14 | 15 | def forward(self, x): 16 | for layer in self.layers: 17 | x = layer(x) 18 | return x 19 | 20 | 21 | def create_sequential_multi_layer(n_layers, n_features): 22 | """Very thin, unrealistic network whose bootleneck is the number of subsequent calls.""" 23 | layers = [nn.Linear(n_features, n_features) for _ in range(n_layers)] 24 | 25 | x = inputs = Input(shape=(n_features,)) 26 | for layer in layers: 27 | x = x(layer) 28 | 29 | models = [ 30 | (("functional",), SymbolicModel(inputs, x)), 31 | (("sequential",), nn.Sequential(*layers)), 32 | (("vanilla",), Vanilla(layers)), 33 | ] 34 | return models 35 | -------------------------------------------------------------------------------- /.github/workflows/install-local-and-run.yaml: -------------------------------------------------------------------------------- 1 | name: install local 2 | 3 | on: 4 | push: 5 | branches: [ "main" ] 6 | pull_request: 7 | branches: [ "main" ] 8 | 9 | permissions: 10 | contents: read 11 | 12 | jobs: 13 | build: 14 | strategy: 15 | matrix: 16 | python-version: [ 17 | "3.7", 18 | "3.8", 19 | "3.9", 20 | "3.10", 21 | "3.11", 22 | "3.12", 23 | "3.13" , 24 | ] 25 | os: [ 26 | "ubuntu-22.04", 27 | "windows-latest", 28 | ] 29 | 30 | runs-on: ${{ matrix.os }} 31 | 32 | steps: 33 | - uses: actions/checkout@v4 34 | - name: Set up Python ${{ matrix.python-version }} 35 | uses: actions/setup-python@v5 36 | with: 37 | python-version: ${{ matrix.python-version }} 38 | - name: Install the library and test basic imports 39 | run: | 40 | pip install . 41 | cd .. 42 | pip show pytorch-symbolic 43 | rm -r pytorch-symbolic/* 44 | 45 | python -c "from pytorch_symbolic import SymbolicModel, Input, useful_layers, graph_algorithms" 46 | 47 | -------------------------------------------------------------------------------- /tests/test_raising_assertions.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2022 Szymon Mikler 2 | 3 | from torch import nn 4 | 5 | from pytorch_symbolic import Input, SymbolicModel 6 | 7 | 8 | def test_cycle1(): 9 | x1 = nn.Linear(2, 2)(Input(shape=(2,))) 10 | x2 = nn.Linear(2, 2)(x1) 11 | x3 = nn.Linear(2, 2)(x2) 12 | x4 = nn.Linear(2, 2)(x3) 13 | x5 = nn.Linear(2, 2)(x4) 14 | 15 | x5._children.append(x1) 16 | x1._parents = (x5,) 17 | 18 | try: 19 | _ = SymbolicModel(inputs=x2, outputs=x4) 20 | raise UserWarning("Created model with a cycle! Assertion should have been raised!") 21 | except AssertionError: 22 | pass 23 | 24 | 25 | def test_cycle2(): 26 | x1 = nn.Linear(2, 2)(Input(shape=(2,))) 27 | x2 = nn.Linear(2, 2)(x1) 28 | x3 = nn.Linear(2, 2)(x2) 29 | x4 = x3 + Input(shape=(2,)) 30 | x5 = nn.Linear(2, 2)(x4) 31 | x6 = nn.Linear(2, 2)(x5) 32 | x7 = nn.Linear(2, 2)(x6) 33 | 34 | x4._parents = (x3, x5) # introduce a cycle 35 | 36 | try: 37 | _ = SymbolicModel(inputs=x2, outputs=x7) 38 | raise UserWarning("Created model with a cycle! Assertion should have been raised!") 39 | except AssertionError: 40 | pass 41 | -------------------------------------------------------------------------------- /benchmarks/README.md: -------------------------------------------------------------------------------- 1 | # Benchmarks 2 | 3 | Fair benchmarking can be very difficult. 4 | 5 | This is not a guide how to benchmark stuff correctly, but we'll share what we used to produce our charts. 6 | 7 | We're using Linux with Intel CPU, some of the commands are different for AMD CPUs. 8 | 9 | Before running any benchmarking script on thread `$THREAD`, we make sure to run the following: 10 | 11 | ```bash 12 | echo 1 > /sys/devices/system/cpu/intel_pstate/no_turbo 13 | echo 0 | tee /sys/kernel/randomize_va_space 14 | 15 | echo performance > /sys/devices/system/cpu/cpu"$THREAD"/cpufreq/scaling_governor 16 | echo performance > /sys/devices/system/cpu/cpu"$THREAD_NEXT"/cpufreq/scaling_governor 17 | 18 | echo 1 > /sys/devices/system/cpu/cpu"$THREAD"/online 19 | echo 0 > /sys/devices/system/cpu/cpu"$THREAD_NEXT"/online 20 | ``` 21 | 22 | # Requirement 23 | 24 | NVIDIA's dllogger is required for logging. It can be downloaded from [https://github.com/NVIDIA/dllogger](https://github.com/NVIDIA/dllogger). 25 | 26 | # Sources 27 | 28 | * [https://llvm.org/docs/Benchmarking.html](https://llvm.org/docs/Benchmarking.html) 29 | 30 | * [https://easyperf.net/blog/2019/08/02/Perf-measurement-environment-on-Linux](https://easyperf.net/blog/2019/08/02/Perf-measurement-environment-on-Linux) 31 | -------------------------------------------------------------------------------- /examples/encoder_decoder.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2022 Szymon Mikler 2 | 3 | """ 4 | CORRESPONDES TO THE EXAMPLE FROM TENSORFLOW GUIDE 5 | https://www.tensorflow.org/guide/keras/functional 6 | """ 7 | 8 | from torch import nn 9 | 10 | from pytorch_symbolic import Input, SymbolicModel, useful_layers 11 | 12 | 13 | def simple_encoder_decoder(input_shape=(1, 28, 28)): 14 | relu = nn.ReLU() 15 | 16 | encoder_input = x = Input(shape=input_shape) 17 | x = x(nn.Conv2d(x.channels, 16, 3))(relu) 18 | x = x(nn.Conv2d(x.channels, 32, 3))(relu) 19 | x = x(nn.MaxPool2d(3)) 20 | x = x(nn.Conv2d(x.channels, 32, 3))(relu) 21 | x = x(nn.Conv2d(x.channels, 16, 3))(relu) 22 | encoder_output = x(nn.MaxPool2d(kernel_size=(x.H, x.W)))(nn.Flatten()) 23 | encoder = SymbolicModel(inputs=encoder_input, outputs=encoder_output) 24 | 25 | decoder_input = x = Input(shape=(encoder_output.features,)) 26 | x = x(useful_layers.ReshapeLayer((1, 4, 4))) 27 | x = x(nn.ConvTranspose2d(x.channels, 16, 3))(relu) 28 | x = x(nn.ConvTranspose2d(x.channels, 32, 3))(relu) 29 | x = x(nn.Upsample(3)) 30 | x = x(nn.ConvTranspose2d(x.channels, 16, 3))(relu) 31 | decoder_output = x(nn.ConvTranspose2d(x.channels, 1, 3))(relu) 32 | decoder = SymbolicModel(inputs=decoder_input, outputs=decoder_output) 33 | 34 | autoencoder_input = Input(shape=input_shape) 35 | encoded_img = autoencoder_input(encoder) 36 | decoded_img = encoded_img(decoder) 37 | autoencoder = SymbolicModel(autoencoder_input, decoded_img) 38 | return autoencoder 39 | -------------------------------------------------------------------------------- /tests/test_iterating_and_slicing.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2022 Szymon Mikler 2 | 3 | import torch 4 | from torch import nn 5 | 6 | from pytorch_symbolic import CustomInput, Input, SymbolicModel, add_to_graph 7 | 8 | 9 | def test_iter_sum(): 10 | x = Input(batch_shape=(10, 20)) 11 | 12 | def tower(x): 13 | for _ in range(10): 14 | x = nn.Identity()(x) 15 | return x 16 | 17 | summed = y = Input(batch_shape=(20,)) 18 | for i, row in enumerate(x): 19 | summed += tower(row * (i + 1)) 20 | model1 = SymbolicModel(inputs=(x, y), outputs=summed) 21 | 22 | summed = y = Input(batch_shape=(20,)) 23 | for i, row in enumerate(x): 24 | summed += row * (i + 1) 25 | model2 = SymbolicModel(inputs=(x, y), outputs=summed) 26 | 27 | torch.manual_seed(0) 28 | x = torch.rand(10, 20) 29 | real_result = torch.zeros(size=(20,)) 30 | 31 | for i, row in enumerate(x): 32 | real_result += row * (i + 1) 33 | 34 | y = torch.zeros(size=(20,)) 35 | our_result1 = model1(x, y) 36 | our_result2 = model2(x, y) 37 | assert torch.equal(real_result, our_result1) 38 | assert torch.equal(real_result, our_result2) 39 | 40 | 41 | def test_indexing_basic1(): 42 | x = Input((10,)) 43 | 44 | def f(x): 45 | return x, [1, 2] 46 | 47 | outs1, outs2 = add_to_graph(f, x) 48 | _ = outs2[0] 49 | _ = outs2[1] 50 | 51 | 52 | def test_unpacking_non_tensor(): 53 | x = CustomInput(data=[5, 6, 7]) 54 | 55 | model = SymbolicModel(inputs=x, outputs=(*x,)) 56 | 57 | real_x = [9, 0, 1] 58 | outs = model(real_x) 59 | assert outs == (9, 0, 1) 60 | -------------------------------------------------------------------------------- /examples/vgg.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2022 Szymon Mikler 2 | 3 | """ 4 | This is a flexible implementation of VGG architecture. 5 | """ 6 | 7 | from torch import nn 8 | 9 | from pytorch_symbolic import Input, SymbolicModel 10 | 11 | from .common import classifier 12 | 13 | relu = nn.ReLU() 14 | 15 | 16 | def VGG( 17 | input_shape, 18 | n_classes, 19 | version=None, 20 | group_sizes=(1, 1, 2, 2, 2), 21 | channels=(64, 128, 256, 512, 512), 22 | pools=(2, 2, 2, 2, 2), 23 | activation=relu, 24 | final_pooling="avgpool", 25 | **kwargs, 26 | ): 27 | if kwargs: 28 | print(f"VGG: unknown parameters: {kwargs.keys()}") 29 | if version: 30 | if version == 11: 31 | group_sizes = (1, 1, 2, 2, 2) 32 | elif version == 13: 33 | group_sizes = (2, 2, 2, 2, 2) 34 | elif version == 16: 35 | group_sizes = (2, 2, 3, 3, 3) 36 | elif version == 19: 37 | group_sizes = (2, 2, 4, 4, 4) 38 | else: 39 | raise NotImplementedError(f"Unkown version={version}!") 40 | 41 | inputs = Input(shape=input_shape) 42 | flow = inputs 43 | 44 | iteration = 0 45 | for group_size, width, pool in zip(group_sizes, channels, pools): 46 | if iteration == 0: 47 | iteration = 1 48 | else: 49 | flow = flow(nn.MaxPool2d(pool)) 50 | 51 | for _ in range(group_size): 52 | flow = flow(nn.Conv2d(flow.channels, width, 3, 1, 1, bias=False)) 53 | flow = flow(nn.BatchNorm2d(flow.channels))(activation) 54 | 55 | outs = classifier(flow, n_classes, final_pooling) 56 | model = SymbolicModel(inputs=inputs, outputs=outs) 57 | return model 58 | -------------------------------------------------------------------------------- /tests/test_nontorch_graphs.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2022 Szymon Mikler 2 | 3 | import random 4 | 5 | from pytorch_symbolic import CustomInput, SymbolicModel 6 | 7 | 8 | def test_set_opts(): 9 | x = CustomInput(set()) 10 | y = CustomInput(set()) 11 | z = CustomInput(set()) 12 | u = {i * 2 for i in range(50)} # constant value 13 | 14 | outs = ((x & y ^ z) | (x ^ z - y)) & u 15 | model = SymbolicModel((x, y, z), outs) 16 | 17 | for _case in range(10): 18 | x = set(random.randint(0, 100) for _ in range(20)) 19 | y = set(random.randint(0, 100) for _ in range(20)) 20 | z = set(random.randint(0, 100) for _ in range(20)) 21 | 22 | outs = model(x, y, z) 23 | real_outs = ((x & y ^ z) | (x ^ z - y)) & u 24 | assert outs == real_outs 25 | 26 | 27 | def test_illegal_inplace(): 28 | """This test performs in-place operations which are discuraged! 29 | 30 | A lot might break when using them. Even if they output None, they are required in the outputs. 31 | If they won't be here, they won't be replayed. Even if they are there, other problems might happend 32 | when multiple models share the same nodes. 33 | """ 34 | x = CustomInput([]) 35 | o1 = x.append(1) 36 | o2 = x.append(2) 37 | o3 = x.append(3) 38 | model = SymbolicModel(inputs=x, outputs=(x, o1, o2, o3)) 39 | 40 | outs, _, _, _ = model([3, 4, 5]) 41 | assert outs == [3, 4, 5, 1, 2, 3] 42 | 43 | 44 | def test_int_and_constants(): 45 | x = CustomInput(5) 46 | y = CustomInput(1) 47 | z = 14 48 | 49 | a = x + 14 50 | b = z + x 51 | c = x + y 52 | 53 | model = SymbolicModel((x, y), outputs=(a, b, c)) 54 | 55 | outs = model(22, 21) 56 | assert outs == (36, 36, 43) 57 | -------------------------------------------------------------------------------- /pytorch_symbolic/model_tools.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2022 Szymon Mikler 2 | 3 | from __future__ import annotations 4 | 5 | import torch 6 | from torch import nn 7 | 8 | 9 | def get_parameter_count(model: nn.Module, only_trainable=False): 10 | """Get the number of parameters of a model.""" 11 | cnt = 0 12 | for param in model.parameters(): 13 | if only_trainable and not param.requires_grad: 14 | continue 15 | cnt += param.shape.numel() 16 | return cnt 17 | 18 | 19 | def get_parameter_shapes(model: nn.Module): 20 | """Get the shapes of parameters of a model.""" 21 | shapes = [] 22 | for param in model.parameters(): 23 | shapes.append(tuple(param.shape)) 24 | return shapes 25 | 26 | 27 | def model_similar(a: nn.Module, b: nn.Module): 28 | """Check whether two models have the same number of parameters and the same shapes of parameters.""" 29 | if get_parameter_count(a) != get_parameter_count(b): 30 | return False 31 | 32 | if sorted(get_parameter_shapes(a)) != sorted(get_parameter_shapes(b)): 33 | return False 34 | return True 35 | 36 | 37 | def hash_torch_tensor(tensor: torch.Tensor): 38 | """Interpret the tensor as a string and return its hash.""" 39 | tensor_as_string = str(tensor.flatten().tolist()) 40 | return hash(tensor_as_string) 41 | 42 | 43 | def models_have_corresponding_parameters(a: nn.Module, b: nn.Module): 44 | """Check whether two models' parameters have identical hash values. 45 | 46 | Parameter order does not matter. 47 | So if two models have identical parameters but in different order, this will still return True. 48 | """ 49 | hashes_a = [hash_torch_tensor(p) for p in a.parameters()] 50 | hashes_b = [hash_torch_tensor(p) for p in b.parameters()] 51 | return set(hashes_a) == set(hashes_b) 52 | -------------------------------------------------------------------------------- /docs/reference/add_to_graph.md: -------------------------------------------------------------------------------- 1 | # add_to_graph 2 | 3 | This module provides a way to add custom functions to the graph. 4 | 5 | To register a function in your graph, instead of doing: 6 | 7 | - ``x = function(*args, **kwds)`` 8 | 9 | do this: 10 | 11 | - ``x = add_to_graph(function, *args, **kwds)``. 12 | 13 | For this to work, there must be at least one Symbolic Data among ``*args, **kwds``, 14 | but other data types are allowed to be there as well. 15 | 16 | Example for using ``torch.concat``: 17 | 18 | ```python 19 | from pytorch_symbolic import Input 20 | from pytorch_symbolic.functions_utility import add_to_graph 21 | import torch 22 | 23 | v1 = Input((10,)) 24 | v2 = Input((20,)) 25 | output = add_to_graph(torch.concat, tensors=(v1, v2), dim=1) 26 | output 27 | ``` 28 | 29 | ``` 30 | 31 | ``` 32 | 33 | This will work for most of the user custom functions, even if Symbolic Tensors 34 | are hidden in nested tuples, lists or dicts. You should also know that there is 35 | a small time overhead for `__call__` during runtime for every function registered this way. 36 | This overhead _should not_ be present when dealing with large models on GPU, 37 | because then CPU does its work before GPU finishes previous kernel computation. 38 | 39 | Recommended, overhead-free way to use custom functions is to write yourself an ``nn.Module`` that does the same as the function of choice. 40 | Then you can use the model without sacrificing performance. 41 | 42 | ::: pytorch_symbolic.add_to_graph 43 | options: 44 | show_source: false 45 | heading_level: 2 46 | show_root_heading: true 47 | members_order: source 48 | show_object_full_path: false 49 | docstring_section_style: table 50 | show_signature_annotations: true 51 | separate_signature: true 52 | annotations_path: brief 53 | merge_init_into_class: true 54 | show_root_full_path: true 55 | -------------------------------------------------------------------------------- /pyproject.toml: -------------------------------------------------------------------------------- 1 | [build-system] 2 | requires = ["hatchling"] 3 | build-backend = "hatchling.build" 4 | 5 | [project] 6 | name = "pytorch-symbolic" 7 | dynamic = ["version"] 8 | description = "Symbolic API for model creation in PyTorch." 9 | dependencies = ["torch"] 10 | license = { text = "MIT" } 11 | requires-python = ">=3.7" 12 | authors = [ 13 | { name = "Szymon Mikler", email = "sjmikler@gmail.com" } 14 | ] 15 | classifiers = [ 16 | "License :: OSI Approved :: MIT License", 17 | "Programming Language :: Python :: 3", 18 | "Programming Language :: Python :: 3.7", 19 | "Programming Language :: Python :: 3.8", 20 | "Programming Language :: Python :: 3.9", 21 | "Programming Language :: Python :: 3.10", 22 | "Programming Language :: Python :: 3.11", 23 | "Programming Language :: Python :: 3.12", 24 | "Programming Language :: Python :: 3.13", 25 | ] 26 | readme = "README.md" 27 | 28 | [project.optional-dependencies] 29 | dev = ["ruff", "pyright", "pytest", "pytest-cov", "hatch", "twine"] 30 | full = ["networkx", "matplotlib", "scipy"] 31 | 32 | [project.urls] 33 | Home = "https://github.com/gahaalt/pytorch-symbolic" 34 | Documentation = "https://github.com/sjmikler/pytorch-symbolic/blob/main/docs" 35 | 36 | [tool.hatch.version] 37 | path = "pytorch_symbolic/__init__.py" 38 | 39 | [tool.hatch.build.targets.sdist] 40 | include = ["pytorch_symbolic", "examples", "docs", "README.md"] 41 | 42 | [tool.hatch.build.targets.wheel] 43 | packages = ["pytorch_symbolic"] 44 | 45 | # %% 46 | 47 | [tool.pyright] 48 | typeCheckingMode = "standard" 49 | exclude = ["devel", "build", "dist"] 50 | 51 | [tool.ruff] 52 | line-length = 120 53 | target-version = "py39" 54 | 55 | [tool.ruff.lint] 56 | select = ["E", "F", "I", "B"] 57 | 58 | # %% 59 | 60 | [tool.pytest.ini_options] 61 | pythonpath = ["."] 62 | 63 | [tool.mypy] 64 | ignore_missing_imports = true 65 | exclude = ["devel", "build", "dist"] 66 | 67 | [tool.isort] 68 | profile = "black" 69 | line_length = 120 70 | 71 | [tool.black] 72 | line_length = 120 -------------------------------------------------------------------------------- /tests/test_creating_extreme.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2022 Szymon Mikler 2 | 3 | import random 4 | 5 | import torch 6 | from torch import nn 7 | 8 | from pytorch_symbolic import Input, SymbolicModel, model_tools 9 | 10 | 11 | def test_deep_linear(): 12 | n_layers = 2500 13 | layers = [nn.Linear(4, 4) for _ in range(n_layers)] 14 | 15 | x = inputs = Input(shape=(4,)) 16 | for layer in layers: 17 | x = layer(x) 18 | symbolic_model = SymbolicModel(inputs, x) 19 | model = nn.Sequential(*layers) 20 | 21 | x = torch.rand(size=(4, 4)) 22 | assert torch.equal(model(x), symbolic_model(x)) 23 | assert model_tools.model_similar(model, symbolic_model) 24 | assert model_tools.models_have_corresponding_parameters(model, symbolic_model) 25 | 26 | 27 | def test_detached_deep_linear(): 28 | n_layers = 2500 29 | layers = [nn.Linear(4, 4) for _ in range(n_layers)] 30 | 31 | x = inputs = Input(shape=(4,)) 32 | for layer in layers: 33 | x = layer(x) 34 | symbolic_model = SymbolicModel(inputs, x).detach_from_graph() 35 | model = nn.Sequential(*layers) 36 | 37 | x = torch.rand(size=(4, 4)) 38 | assert torch.equal(model(x), symbolic_model(x)) 39 | assert model_tools.model_similar(model, symbolic_model) 40 | assert model_tools.models_have_corresponding_parameters(model, symbolic_model) 41 | 42 | 43 | def test_random_models_from_large_graph(): 44 | inputs = Input(shape=(4,)) 45 | nodes = [inputs] 46 | 47 | n_layers = 2500 48 | for _ in range(n_layers): 49 | parent = random.choice(nodes) 50 | child = nn.Linear(4, 4)(parent) 51 | nodes.append(child) 52 | 53 | for _ in range(10): 54 | idx = random.randint(0, n_layers - 1) 55 | model = SymbolicModel(inputs=inputs, outputs=nodes[idx]) 56 | _ = model(torch.rand(4, 4)) 57 | 58 | counts = model_tools.get_parameter_count(model) 59 | num_layers = len(model._execution_order_layers) 60 | assert counts == num_layers * (16 + 4) 61 | -------------------------------------------------------------------------------- /tests/test_docs_automated.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2022 Szymon Mikler 2 | 3 | import logging 4 | import os 5 | import pathlib 6 | import re 7 | 8 | 9 | def get_header(msg): 10 | msg_len = len(msg) 11 | header = [ 12 | "#" * (msg_len + 6), 13 | "## " + msg + " ##", 14 | "#" * (msg_len + 6), 15 | ] 16 | return "\n".join(header) + "\n" 17 | 18 | 19 | def scan_for_code_blobs(text): 20 | # Capture code blocks starting with py or python 21 | blobs = re.findall(r"```(py|python)\n([\s\S]+?)\n```", text) 22 | 23 | new_blobs = [] 24 | for mode, blob in blobs: 25 | if "..." in blob: 26 | continue 27 | header = get_header(f"Generated ({mode})") 28 | blob = header + blob 29 | 30 | # For `py` code blocks, we append them to previous existing block 31 | # But `python` block starts a new scope and finishes the previous block 32 | # They usually need to include imports 33 | if mode == "py" and new_blobs: 34 | new_blobs[-1] = new_blobs[-1] + "\n" + blob 35 | else: 36 | new_blobs.append(blob) 37 | return new_blobs 38 | 39 | 40 | def test_all_code_blobs(): 41 | assert os.path.exists("docs") 42 | all_code_blobs = [] 43 | 44 | for root, _dirs, files in os.walk("."): 45 | for file in files: 46 | path = pathlib.Path(os.path.join(root, file)) 47 | if path.suffix == ".md": 48 | code_blobs = scan_for_code_blobs(path.open("r").read()) 49 | for blob in code_blobs: 50 | all_code_blobs.append(blob) 51 | 52 | logging.warning(f"Detected {len(all_code_blobs)} code examples!") 53 | 54 | for idx, blob in enumerate(all_code_blobs): 55 | try: 56 | globals_temp = {} 57 | exec(blob, globals_temp) 58 | except Exception as e: 59 | print(f"Exception during automated documentation testing {idx}/{len(all_code_blobs)}:") 60 | print(blob) 61 | raise e 62 | -------------------------------------------------------------------------------- /pytorch_symbolic/symbolic_api_2.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2022 Szymon Mikler 2 | 3 | """ 4 | Enables Symbolic API 2, which allows for registering layers by calling 5 | layer(*symbolic) instead of symbolic_1(layer, *other_symbolic). 6 | """ 7 | 8 | from __future__ import annotations 9 | 10 | import logging 11 | 12 | from torch import nn 13 | 14 | from .symbolic_data import SymbolicData 15 | 16 | __nn_module_old_new__ = nn.Module.__new__ 17 | __nn_module_old_call__ = nn.Module.__call__ 18 | 19 | 20 | def call_wrapper_for_api_2(self, *args, custom_name: str | None = None, **kwds): 21 | if not any(isinstance(x, SymbolicData) for x in args): 22 | return __nn_module_old_call__(self, *args, **kwds) 23 | elif len(kwds) == 0 and all(isinstance(x, SymbolicData) for x in args): 24 | node = args[0] 25 | return node(self, *args[1:], custom_name=custom_name) 26 | else: 27 | msg = "Only *args are allowed as arguments! If you need to use **kwds, try `functions_utility.add_to_graph`!" 28 | raise UserWarning(msg) 29 | 30 | 31 | class SymbolicAPI2ContextManager: 32 | def __enter__(self): 33 | logging.debug("Symbolic API has been enabled!") 34 | nn.Module.__call__ = call_wrapper_for_api_2 35 | 36 | def __exit__(self, exit_type, value, traceback): 37 | logging.debug("Symbolic API has been disabled!") 38 | nn.Module.__call__ = __nn_module_old_call__ 39 | 40 | 41 | def optimize_module_calls(): 42 | """Remove the call wrapper from `torch.nn.Module`.""" 43 | msg = "Optimizing module calls! Reusing existing layers will not work with layer(*symbols) notation!" 44 | logging.warning(msg) 45 | SymbolicAPI2ContextManager().__exit__(None, None, None) 46 | 47 | 48 | def is_symbolic_api_2_enabled(): 49 | """Whether symbolic API 2 is enabled.""" 50 | return nn.Module.__call__ is call_wrapper_for_api_2 51 | 52 | 53 | def enable_symbolic_api_2_for_new_modules(): 54 | """Add a __new__ wrapper for `torch.nn.Module`.""" 55 | 56 | def wrapped_new(self, *args, **kwds): 57 | SymbolicAPI2ContextManager().__enter__() 58 | obj = super(nn.Module, self).__new__(self) 59 | return obj 60 | 61 | nn.Module.__new__ = wrapped_new 62 | 63 | 64 | def disable_symbolic_api_2_for_new_modules(): 65 | """Remove the __new__ wrapper for `torch.nn.Module`.""" 66 | 67 | nn.Module.__new__ = __nn_module_old_new__ 68 | -------------------------------------------------------------------------------- /tests/test_creating_examples.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2022 Szymon Mikler 2 | 3 | import torch 4 | from torch import nn 5 | 6 | import examples.encoder_decoder 7 | import examples.lstm 8 | import examples.resnet 9 | import examples.vgg 10 | from pytorch_symbolic import Input, SymbolicModel, model_tools 11 | 12 | 13 | def test_example_toy_resnet(): 14 | model = examples.resnet.ToyResNet((3, 32, 32), n_classes=10) 15 | inputs = torch.rand(16, 3, 32, 32) 16 | outputs = model(inputs) 17 | assert list(outputs.shape) == [16, 10] 18 | assert model_tools.get_parameter_count(model) == 175530 19 | 20 | 21 | def test_example_wrn(): 22 | model = examples.resnet.ResNet((3, 32, 32), n_classes=10, version=("WRN", 16, 4)) 23 | inputs = torch.rand(16, 3, 32, 32) 24 | outputs = model(inputs) 25 | assert list(outputs.shape) == [16, 10] 26 | assert model_tools.get_parameter_count(model) == 2750698 27 | 28 | 29 | def test_example_vgg(): 30 | model = examples.vgg.VGG((3, 32, 32), n_classes=10, version=13) 31 | inputs = torch.rand(16, 3, 32, 32) 32 | outputs = model(inputs) 33 | assert list(outputs.shape) == [16, 10] 34 | assert model_tools.get_parameter_count(model) == 9413066 35 | 36 | 37 | def test_example_enc_dec(): 38 | model = examples.encoder_decoder.simple_encoder_decoder((3, 32, 32)) 39 | inputs = torch.rand(16, 3, 32, 32) 40 | outputs = model(inputs) 41 | assert list(outputs.shape) == [16, 1, 7, 7] 42 | assert model_tools.get_parameter_count(model) == 28529 43 | 44 | 45 | def test_resnet(): 46 | inputs = Input(shape=(3, 32, 32)) 47 | x = nn.Conv2d(inputs.channels, 32, 3)(inputs)(nn.ReLU()) 48 | x = nn.Conv2d(x.channels, 64, 3)(x)(nn.ReLU()) 49 | block_1_output = nn.MaxPool2d(3)(x) 50 | 51 | x = nn.Conv2d(block_1_output.channels, 64, 3, padding=1)(block_1_output)(nn.ReLU()) 52 | x = nn.Conv2d(x.channels, 64, 3, padding=1)(x)(nn.ReLU()) 53 | block_2_output = x + block_1_output 54 | 55 | x = nn.Conv2d(block_2_output.channels, 64, 3, padding=1)(block_2_output)(nn.ReLU()) 56 | x = nn.Conv2d(x.channels, 64, 3, padding=1)(x)(nn.ReLU()) 57 | block_3_output = x + block_2_output 58 | 59 | x = nn.Conv2d(x.channels, 64, 3)(block_3_output)(nn.ReLU()) 60 | x = nn.AvgPool2d(kernel_size=(x.H, x.W))(x)(nn.Flatten()) 61 | x = nn.Linear(x.features, 256)(x)(nn.ReLU()) 62 | x = nn.Dropout(0.5)(x) 63 | outputs = nn.Linear(x.features, 10)(x) 64 | 65 | model = SymbolicModel(inputs, outputs) 66 | assert model_tools.get_parameter_count(model) == 223242 67 | 68 | 69 | def test_lstm(): 70 | examples.lstm.run() 71 | -------------------------------------------------------------------------------- /tests/test_creating_examples_detached.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2022 Szymon Mikler 2 | 3 | import torch 4 | from torch import nn 5 | 6 | import examples.encoder_decoder 7 | import examples.lstm 8 | import examples.resnet 9 | import examples.vgg 10 | from pytorch_symbolic import Input, SymbolicModel, model_tools 11 | 12 | 13 | def test_detached_example_toy_resnet(): 14 | model = examples.resnet.ToyResNet((3, 32, 32), n_classes=10) 15 | model_detached = model.detach_from_graph() 16 | inputs = torch.rand(16, 3, 32, 32) 17 | assert torch.equal(model(inputs), model_detached(inputs)) 18 | 19 | 20 | def test_detached_example_wrn(): 21 | model = examples.resnet.ResNet((3, 32, 32), n_classes=10, version=("WRN", 16, 4)) 22 | model_detached = model.detach_from_graph() 23 | inputs = torch.rand(16, 3, 32, 32) 24 | assert torch.equal(model(inputs), model_detached(inputs)) 25 | 26 | 27 | def test_detached_example_vgg(): 28 | model = examples.vgg.VGG((3, 32, 32), n_classes=10, version=13) 29 | model_detached = model.detach_from_graph() 30 | inputs = torch.rand(16, 3, 32, 32) 31 | assert torch.equal(model(inputs), model_detached(inputs)) 32 | 33 | 34 | def test_detached_example_enc_dec(): 35 | model = examples.encoder_decoder.simple_encoder_decoder((3, 32, 32)) 36 | model_detached = model.detach_from_graph() 37 | inputs = torch.rand(16, 3, 32, 32) 38 | assert torch.equal(model(inputs), model_detached(inputs)) 39 | 40 | 41 | def test_detached_resnet(): 42 | inputs = Input(shape=(3, 32, 32)) 43 | x = nn.Conv2d(inputs.channels, 32, 3)(inputs)(nn.ReLU()) 44 | x = nn.Conv2d(x.channels, 64, 3)(x)(nn.ReLU()) 45 | block_1_output = nn.MaxPool2d(3)(x) 46 | 47 | x = nn.Conv2d(block_1_output.channels, 64, 3, padding=1)(block_1_output)(nn.ReLU()) 48 | x = nn.Conv2d(x.channels, 64, 3, padding=1)(x)(nn.ReLU()) 49 | block_2_output = x + block_1_output 50 | 51 | x = nn.Conv2d(block_2_output.channels, 64, 3, padding=1)(block_2_output)(nn.ReLU()) 52 | x = nn.Conv2d(x.channels, 64, 3, padding=1)(x)(nn.ReLU()) 53 | block_3_output = x + block_2_output 54 | 55 | x = nn.Conv2d(x.channels, 64, 3)(block_3_output)(nn.ReLU()) 56 | x = nn.AvgPool2d(kernel_size=(x.H, x.W))(x)(nn.Flatten()) 57 | x = nn.Linear(x.features, 256)(x)(nn.ReLU()) 58 | x = nn.Dropout(0.5)(x) 59 | outputs = nn.Linear(x.features, 10)(x) 60 | 61 | model = SymbolicModel(inputs, outputs) 62 | model_detached = model.detach_from_graph() 63 | 64 | assert model_tools.model_similar(model, model_detached) 65 | assert model_tools.models_have_corresponding_parameters(model, model_detached) 66 | -------------------------------------------------------------------------------- /tests/test_reusing_nodes.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2022 Szymon Mikler 2 | 3 | import torch 4 | from torch import nn 5 | 6 | from pytorch_symbolic import Input, SymbolicModel, model_tools 7 | 8 | 9 | def AlwaysTheSameConv(in_channels, out_channels): 10 | torch.manual_seed(42) 11 | 12 | return nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1) 13 | 14 | 15 | class VanillaModel(nn.Module): 16 | def __init__(self): 17 | super().__init__() 18 | self.layers = [AlwaysTheSameConv(3, 3) for _ in range(4)] 19 | for idx, layer in enumerate(self.layers): 20 | self.register_module(str(idx), layer) 21 | 22 | def forward(self, x): 23 | for layer in self.layers: 24 | x = layer(x) 25 | return x 26 | 27 | 28 | def test_models_from_slice(): 29 | x = Input(shape=(3, 10, 10)) 30 | 31 | nodes = [] 32 | for _ in range(20): 33 | x = AlwaysTheSameConv(x.C, 3)(x) 34 | unused = nn.Identity()(x) # noqa: F841 35 | nodes.append(x) 36 | 37 | models = [] 38 | for i in range(0, 20, 5): 39 | nodes_slice = nodes[i : i + 5] 40 | 41 | model = SymbolicModel(inputs=nodes_slice[0], outputs=nodes_slice[-1]) 42 | models.append(model) 43 | 44 | data = torch.rand(16, 3, 10, 10) 45 | outs = prev_outs = models[-1](data) 46 | 47 | vanilla_model = VanillaModel() 48 | vanilla_outs = vanilla_model(data) 49 | assert torch.equal(outs, vanilla_outs) 50 | 51 | for model in models: 52 | outs = model(data) 53 | assert torch.equal(outs, prev_outs) 54 | assert model_tools.model_similar(model, vanilla_model) 55 | assert model_tools.models_have_corresponding_parameters(model, vanilla_model) 56 | prev_outs = outs 57 | 58 | 59 | def test_detached_models_from_slice(): 60 | x = Input(shape=(3, 10, 10)) 61 | 62 | nodes = [] 63 | for _ in range(20): 64 | x = AlwaysTheSameConv(x.C, 3)(x) 65 | unused = nn.Identity()(x) # noqa: F841 66 | nodes.append(x) 67 | 68 | models = [] 69 | for i in range(0, 20, 5): 70 | nodes_slice = nodes[i : i + 5] 71 | 72 | model = SymbolicModel(inputs=nodes_slice[0], outputs=nodes_slice[-1]).detach_from_graph() 73 | models.append(model) 74 | 75 | data = torch.rand(16, 3, 10, 10) 76 | outs = prev_outs = models[-1](data) 77 | 78 | vanilla_model = VanillaModel() 79 | vanilla_outs = vanilla_model(data) 80 | assert torch.equal(outs, vanilla_outs) 81 | 82 | for model in models: 83 | outs = model(data) 84 | assert torch.equal(outs, prev_outs) 85 | assert model_tools.model_similar(model, vanilla_model) 86 | assert model_tools.models_have_corresponding_parameters(model, vanilla_model) 87 | prev_outs = outs 88 | -------------------------------------------------------------------------------- /benchmarks/bench2/define_models.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2022 Szymon Mikler 2 | 3 | from torch import nn 4 | 5 | from pytorch_symbolic import Input, SymbolicModel 6 | 7 | 8 | class ToyResNet(nn.Module): 9 | def __init__(self): 10 | super().__init__() 11 | self.relu = nn.ReLU() 12 | self.block1conv1 = nn.Conv2d(3, 32, 3) 13 | self.block1conv2 = nn.Conv2d(32, 64, 3) 14 | self.maxpool = nn.MaxPool2d(3) 15 | 16 | self.block2conv1 = nn.Conv2d(64, 64, 3, padding=1) 17 | self.block2conv2 = nn.Conv2d(64, 64, 3, padding=1) 18 | 19 | self.block3conv1 = nn.Conv2d(64, 64, 3, padding=1) 20 | self.block3conv2 = nn.Conv2d(64, 64, 3, padding=1) 21 | 22 | self.conv1 = nn.Conv2d(64, 64, 3) 23 | 24 | # does not work with varying image size 25 | # kernel_size = 7 # calculated by hand 26 | # self.global_pool = nn.AvgPool2d(kernel_size) 27 | self.adaptive_pool = nn.AdaptiveAvgPool2d((1, 1)) 28 | 29 | self.flatten = nn.Flatten() 30 | self.linear = nn.Linear(64, 256) 31 | self.dropout = nn.Dropout(0.5) 32 | self.classifier = nn.Linear(256, 10) 33 | 34 | def forward(self, x): 35 | x = self.relu(self.block1conv1(x)) 36 | x = self.relu(self.block1conv2(x)) 37 | block_1_output = self.maxpool(x) 38 | 39 | x = self.relu(self.block2conv1(block_1_output)) 40 | x = self.relu(self.block2conv2(x)) 41 | block_2_output = x + block_1_output 42 | 43 | x = self.relu(self.block3conv1(block_2_output)) 44 | x = self.relu(self.block3conv2(x)) 45 | block_3_output = x + block_2_output 46 | 47 | x = self.relu(self.conv1(block_3_output)) 48 | x = self.adaptive_pool(x) 49 | x = self.flatten(x) 50 | x = self.relu(self.linear(x)) 51 | x = self.dropout(x) 52 | return self.classifier(x) 53 | 54 | 55 | def functional_toy_resnet(bs, img_size): 56 | inputs = Input(batch_shape=(bs, 3, img_size, img_size)) 57 | 58 | x = inputs(nn.Conv2d(inputs.channels, 32, 3))(nn.ReLU()) 59 | x = x(nn.Conv2d(x.channels, 64, 3))(nn.ReLU()) 60 | block_1_output = x(nn.MaxPool2d(3)) 61 | 62 | x = block_1_output(nn.Conv2d(block_1_output.channels, 64, 3, padding=1))(nn.ReLU()) 63 | x = x(nn.Conv2d(x.channels, 64, 3, padding=1))(nn.ReLU()) 64 | block_2_output = x + block_1_output 65 | 66 | x = block_2_output(nn.Conv2d(block_2_output.channels, 64, 3, padding=1))(nn.ReLU()) 67 | x = x(nn.Conv2d(x.channels, 64, 3, padding=1))(nn.ReLU()) 68 | block_3_output = x + block_2_output 69 | 70 | x = block_3_output(nn.Conv2d(x.channels, 64, 3))(nn.ReLU()) 71 | x = x(nn.AdaptiveAvgPool2d((1, 1)))(nn.Flatten()) 72 | x = x(nn.Linear(x.features, 256))(nn.ReLU()) 73 | x = x(nn.Dropout(0.5)) 74 | outputs = x(nn.Linear(x.features, 10)) 75 | 76 | return SymbolicModel(inputs, outputs) 77 | 78 | 79 | def create_toy_resnets(bs, img_size): 80 | """ToyResNet example from https://www.tensorflow.org/guide/keras/functional#a_toy_resnet_model.""" 81 | models = [ 82 | (("functional",), functional_toy_resnet(bs, img_size=img_size)), 83 | (("vanilla",), ToyResNet()), 84 | ] 85 | return models 86 | -------------------------------------------------------------------------------- /tests/test_docs_manual.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2022 Szymon Mikler 2 | 3 | from torch import nn 4 | 5 | from pytorch_symbolic import Input, SymbolicModel, model_tools, useful_layers 6 | 7 | 8 | def test_1(): 9 | inputs = Input(shape=(3, 32, 32)) 10 | x = inputs(nn.Conv2d(inputs.channels, 32, 3))(nn.ReLU()) 11 | x = x(nn.Conv2d(x.channels, 64, 3))(nn.ReLU()) 12 | block_1_output = x(nn.MaxPool2d(3)) 13 | 14 | x = block_1_output(nn.Conv2d(block_1_output.channels, 64, 3, padding=1))(nn.ReLU()) 15 | x = x(nn.Conv2d(x.channels, 64, 3, padding=1))(nn.ReLU()) 16 | block_2_output = x + block_1_output 17 | 18 | x = block_2_output(nn.Conv2d(block_2_output.channels, 64, 3, padding=1))(nn.ReLU()) 19 | x = x(nn.Conv2d(x.channels, 64, 3, padding=1))(nn.ReLU()) 20 | block_3_output = x + block_2_output 21 | 22 | x = block_3_output(nn.Conv2d(x.channels, 64, 3))(nn.ReLU()) 23 | x = x(nn.AvgPool2d(kernel_size=(x.H, x.W)))(nn.Flatten()) 24 | x = x(nn.Linear(x.features, 256))(nn.ReLU()) 25 | x = x(nn.Dropout(0.5)) 26 | outputs = x(nn.Linear(x.features, 10)) 27 | 28 | model = SymbolicModel(inputs, outputs) 29 | assert model_tools.get_parameter_count(model) == 223242 30 | 31 | 32 | def test_2(): 33 | inputs = Input((3, 128, 128)) 34 | x = inputs 35 | 36 | x = x(nn.Conv2d(in_channels=x.channels, out_channels=16, kernel_size=3)) 37 | x = x(nn.MaxPool2d(kernel_size=2)) 38 | x = x(nn.ReLU()) 39 | 40 | x = x(nn.Conv2d(in_channels=x.channels, out_channels=32, kernel_size=3)) 41 | x = x(nn.MaxPool2d(kernel_size=2)) 42 | x = x(nn.ReLU()) 43 | 44 | x = x(nn.Conv2d(in_channels=x.channels, out_channels=64, kernel_size=3)) 45 | x = x(nn.MaxPool2d(kernel_size=2)) 46 | x = x(nn.ReLU()) 47 | 48 | x = x(nn.Conv2d(in_channels=x.channels, out_channels=64, kernel_size=3)) 49 | x = x(nn.MaxPool2d(kernel_size=2)) 50 | x = x(nn.ReLU()) 51 | 52 | x = x(nn.Flatten()) 53 | outputs = x(nn.Linear(in_features=x.features, out_features=10)) 54 | model = SymbolicModel(inputs=inputs, outputs=outputs) 55 | assert model.output_shape == (1, 10) 56 | assert model_tools.get_parameter_count(model) == 83562 57 | 58 | 59 | def test_3(): 60 | task1_input = Input(shape=(1, 28, 28)) 61 | task2_input = Input(shape=(3, 32, 32)) 62 | 63 | x = task1_input 64 | x = x(nn.Conv2d(x.channels, 16, 3)) 65 | x = x(nn.MaxPool2d(3))(nn.ReLU()) 66 | x = x(nn.Flatten()) 67 | head1_out = x(nn.Linear(x.features, 200)) 68 | 69 | x = task2_input 70 | x = x(nn.Conv2d(x.channels, 16, 3)) 71 | x = x(nn.MaxPool2d(3))(nn.ReLU()) 72 | x = x(nn.Flatten()) 73 | head2_out = x(nn.Linear(x.features, 200)) 74 | 75 | x = head1_out + head2_out 76 | x = x(nn.Linear(x.features, 400))(nn.ReLU()) 77 | task1_outputs = x(nn.Linear(x.features, 10)) 78 | task2_outputs = x(nn.Linear(x.features, 10)) 79 | 80 | model = SymbolicModel(inputs=(task1_input, task2_input), outputs=(task1_outputs, task2_outputs)) 81 | assert model_tools.get_parameter_count(model) == 614228 82 | 83 | 84 | def test_4(): 85 | x1 = Input(shape=(1, 2, 3)) 86 | x2 = Input(shape=(5, 2, 3)) 87 | x = x1(useful_layers.ConcatLayer(dim=1), x2) 88 | assert x.shape == (1, 6, 2, 3) 89 | -------------------------------------------------------------------------------- /benchmarks/bench2/run_one.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2022 Szymon Mikler 2 | 3 | import argparse 4 | import gc 5 | import logging 6 | import os 7 | import socket 8 | import time 9 | 10 | import dllogger 11 | import torch 12 | from dllogger import DLLLoggerAlreadyInitialized, JSONStreamBackend, StdOutBackend, Verbosity 13 | 14 | from benchmarks.bench2 import define_models 15 | from pytorch_symbolic import optimize_module_calls 16 | 17 | logging.basicConfig(level=logging.ERROR) 18 | 19 | # garbage collector is messing with timing 20 | gc.disable() 21 | 22 | parser = argparse.ArgumentParser() 23 | 24 | parser.add_argument("--n-warmup", type=int, default=50) 25 | parser.add_argument("--n-iter", type=int, default=1000) 26 | parser.add_argument("--img-size", type=int, default=32) 27 | parser.add_argument("--batch-size", type=int, default=4) 28 | parser.add_argument("--model-idx", type=int, required=True) 29 | parser.add_argument("--device", type=str, default="cpu") 30 | parser.add_argument("--run-name", type=str, default="unknown.jsonl") 31 | parser.add_argument("--path", type=str, default="benchmarks/logs") 32 | 33 | args = parser.parse_args() 34 | 35 | BASE_TAGS = ["oct20"] 36 | 37 | os.makedirs(args.path, exist_ok=True) 38 | 39 | N_WARMUP = args.n_warmup 40 | N_ITER = args.n_iter 41 | BATCH_SIZE = args.batch_size 42 | DEVICE = args.device 43 | IMG_SIZE = args.img_size 44 | 45 | try: 46 | dllogger.init( 47 | [ 48 | JSONStreamBackend( 49 | filename=os.path.join(args.path, args.run_name), 50 | verbosity=Verbosity.DEFAULT, 51 | append=True, 52 | ), 53 | StdOutBackend(verbosity=Verbosity.DEFAULT), 54 | ] 55 | ) 56 | except DLLLoggerAlreadyInitialized: 57 | pass 58 | 59 | 60 | def run(model, device): 61 | x = torch.rand(size=(BATCH_SIZE, 3, IMG_SIZE, IMG_SIZE)) 62 | x = x.to(device) 63 | 64 | for _ in range(N_WARMUP): 65 | _ = model(x) 66 | 67 | if DEVICE != "cpu": 68 | torch.cuda.synchronize() 69 | 70 | # TIMING 71 | t0 = time.time() 72 | for _ in range(N_ITER): 73 | _ = model(x) 74 | if DEVICE != "cpu": 75 | torch.cuda.synchronize() 76 | td = time.time() - t0 77 | return td 78 | 79 | 80 | def log(name, time_per_run): 81 | dllogger.log( 82 | step={}, 83 | data={ 84 | "tags": name, 85 | "throughput": 1 / time_per_run, 86 | "IMG_SIZE": IMG_SIZE, 87 | "N_WARMUP": N_WARMUP, 88 | "N_ITER": N_ITER, 89 | "BATCH_SIZE": BATCH_SIZE, 90 | "HOST": host_info, 91 | }, 92 | ) 93 | dllogger.flush() 94 | 95 | 96 | if __name__ == "__main__": 97 | device = torch.device(DEVICE) 98 | 99 | models = define_models.create_toy_resnets(BATCH_SIZE, IMG_SIZE) 100 | tags, model = models[args.model_idx] 101 | model = model.to(device) 102 | optimize_module_calls() 103 | 104 | if not isinstance(tags, tuple): 105 | tags = (tags,) 106 | 107 | host_info = { 108 | "name": socket.gethostname(), 109 | "device": torch.cuda.get_device_name(device) if "cuda" in DEVICE else "cpu", 110 | } 111 | 112 | td = run(model, device) 113 | log([*BASE_TAGS, *tags, "call_optimization"], td / N_ITER) 114 | -------------------------------------------------------------------------------- /benchmarks/bench2/draw_results.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2022 Szymon Mikler 2 | 3 | from __future__ import annotations 4 | 5 | import matplotlib.pyplot as plt 6 | import numpy as np 7 | import pandas as pd 8 | import seaborn as sns 9 | 10 | from benchmarks.tagged_collection import TaggedCollection 11 | 12 | sns.set_style("darkgrid") 13 | 14 | 15 | def plot( 16 | df: pd.DataFrame, 17 | x_axis, 18 | y_axis, 19 | compare, 20 | reference_y=None, 21 | title="UNKNOWN", 22 | xlabel=None, 23 | ylabel=None, 24 | ax: plt.Axes = None, 25 | fig: plt.Figure = None, 26 | linewidth: int | None = None, 27 | ): 28 | if fig is None or ax is None: 29 | fig, ax = plt.subplots(figsize=(12, 8), constrained_layout=True, dpi=300) 30 | 31 | if reference_y: 32 | ys = df.apply(lambda row: reference_y[row[X]], axis=1) 33 | print(ys) 34 | df[Y] = (df[Y] - ys) / ys 35 | 36 | def t(x): 37 | return np.percentile(x, q=20), np.percentile(x, q=80) 38 | 39 | sns.lineplot( 40 | ax=ax, 41 | data=df, 42 | x=x_axis, 43 | y=y_axis, 44 | estimator="median", 45 | err_kws={"alpha": 0.15}, 46 | hue=compare, 47 | # hue=None, 48 | # errorbar=t, 49 | errorbar=("pi", 50), 50 | # errorbar=lambda x: 1, 51 | # errorbar=None, 52 | # markers=True, 53 | ) 54 | 55 | if linewidth is not None: 56 | for line in ax.lines: 57 | line.set_linewidth(2) 58 | 59 | if xlabel: 60 | ax.set_xlabel(xlabel) 61 | if ylabel: 62 | ax.set_ylabel(ylabel) 63 | ax.set_title(title) 64 | 65 | 66 | ######################### 67 | 68 | X = "data.IMG_SIZE" 69 | Y = "data.throughput" 70 | 71 | filters = (lambda row: "call_optimization" in row.tags,) 72 | 73 | tagged = TaggedCollection.from_dllogs("benchmarks/good_logs/bench2.jsonl") 74 | 75 | tagged = tagged.filter(*filters) 76 | assert len(tagged) > 0, "Empty query!" 77 | 78 | references = tagged["vanilla"]["call_optimization"] 79 | reference_y_for_x = {} 80 | 81 | for x, reference_for_x in references.groupby(lambda row: row[X], return_keys=True): 82 | avg = np.median(reference_for_x.get(Y)) 83 | reference_y_for_x[x] = avg 84 | 85 | fig, ax = plt.subplots(nrows=1, ncols=1, figsize=(12, 5), constrained_layout=True, dpi=300) 86 | ax = [ax] 87 | 88 | names = { 89 | "oct20,vanilla,call_optimization": "Inheriting from nn.Module", 90 | "oct20,symbolic,call_optimization": "SymbolicModel", 91 | } 92 | 93 | for tags, group in tagged.groupby("tags", return_keys=True): 94 | assert tags in names 95 | name = names[tags] 96 | for row in group: 97 | row["Definition"] = name 98 | 99 | plot( 100 | tagged.df, 101 | x_axis=X, 102 | y_axis=Y, 103 | compare="Definition", 104 | reference_y=None, 105 | title="Performance comparison between different definitions of Toy ResNet", 106 | xlabel="Height and width of the input images", 107 | ylabel="Batches per second (more is better)", 108 | ax=ax[0], 109 | fig=fig, 110 | ) 111 | # plot( 112 | # tagged.df, 113 | # x_axis=X, 114 | # y_axis=Y, 115 | # compare="tags", 116 | # reference_y=reference_y_for_x, 117 | # title="Close-up", 118 | # xlabel="Number of linear layers", 119 | # ylabel="Throughput difference (%)", 120 | # ax=ax[1], 121 | # fig=fig, 122 | # ) 123 | plt.savefig("latest.png") 124 | plt.show() 125 | -------------------------------------------------------------------------------- /benchmarks/bench1/run_one.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2022 Szymon Mikler 2 | 3 | import argparse 4 | import gc 5 | import logging 6 | import os 7 | import socket 8 | import time 9 | from copy import deepcopy 10 | 11 | import dllogger 12 | import torch 13 | from dllogger import DLLLoggerAlreadyInitialized, JSONStreamBackend, StdOutBackend, Verbosity 14 | 15 | from benchmarks.bench1 import define_models 16 | from pytorch_symbolic import optimize_module_calls 17 | 18 | logging.basicConfig(level=logging.ERROR) 19 | 20 | # garbage collector is messing with timing 21 | gc.disable() 22 | 23 | parser = argparse.ArgumentParser() 24 | 25 | parser.add_argument("--n-layers", type=int, required=True) 26 | parser.add_argument("--n-warmup", type=int, default=50) 27 | parser.add_argument("--n-iter", type=int, default=1000) 28 | parser.add_argument("--n-features", type=int, default=4) 29 | parser.add_argument("--batch-size", type=int, default=4) 30 | parser.add_argument("--model-idx", type=int, required=True) 31 | parser.add_argument("--device", type=str, default="cpu") 32 | parser.add_argument("--run-name", type=str, default="unknown.jsonl") 33 | parser.add_argument("--path", type=str, default="benchmarks/logs") 34 | 35 | args = parser.parse_args() 36 | 37 | BASE_TAGS = ["oct20"] 38 | 39 | os.makedirs(args.path, exist_ok=True) 40 | 41 | N_WARMUP = args.n_warmup 42 | N_ITER = args.n_iter 43 | N_LAYERS = args.n_layers 44 | N_FEATURES = args.n_features 45 | BATCH_SIZE = args.batch_size 46 | DEVICE = args.device 47 | 48 | try: 49 | dllogger.init( 50 | [ 51 | JSONStreamBackend( 52 | filename=os.path.join(args.path, args.run_name), 53 | verbosity=Verbosity.DEFAULT, 54 | append=True, 55 | ), 56 | StdOutBackend(verbosity=Verbosity.DEFAULT), 57 | ] 58 | ) 59 | except DLLLoggerAlreadyInitialized: 60 | pass 61 | 62 | 63 | def run(model, device): 64 | x = torch.rand(size=(BATCH_SIZE, N_FEATURES)) 65 | x = x.to(device) 66 | 67 | for _ in range(N_WARMUP): 68 | _ = model(x) 69 | 70 | if DEVICE != "cpu": 71 | torch.cuda.synchronize() 72 | 73 | # TIMING 74 | t0 = time.time() 75 | for _ in range(N_ITER): 76 | _ = model(x) 77 | if DEVICE != "cpu": 78 | torch.cuda.synchronize() 79 | td = time.time() - t0 80 | return td 81 | 82 | 83 | def log(name, layers, time_per_run): 84 | dllogger.log( 85 | step={}, 86 | data={ 87 | "tags": name, 88 | "layers": layers, 89 | "throughput": 1 / time_per_run, 90 | "N_WARMUP": N_WARMUP, 91 | "N_ITER": N_ITER, 92 | "N_FEATURES": N_FEATURES, 93 | "BATCH_SIZE": BATCH_SIZE, 94 | "HOST": host_info, 95 | }, 96 | ) 97 | dllogger.flush() 98 | 99 | 100 | if __name__ == "__main__": 101 | device = torch.device(DEVICE) 102 | 103 | models = define_models.create_sequential_multi_layer(N_LAYERS, N_FEATURES) 104 | tags, model = models[args.model_idx] 105 | model = deepcopy(model) 106 | optimize_module_calls() 107 | 108 | if not isinstance(tags, tuple): 109 | tags = (tags,) 110 | 111 | host_info = { 112 | "name": socket.gethostname(), 113 | "device": torch.cuda.get_device_name(device) if "cuda" in DEVICE else "cpu", 114 | } 115 | 116 | td = run(model, device) 117 | log([*BASE_TAGS, *tags, "call_optimization"], N_LAYERS, td / N_ITER) 118 | -------------------------------------------------------------------------------- /benchmarks/bench1/draw_results.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2022 Szymon Mikler 2 | 3 | from __future__ import annotations 4 | 5 | import matplotlib.pyplot as plt 6 | import numpy as np 7 | import pandas as pd 8 | import seaborn as sns 9 | 10 | from benchmarks.tagged_collection import TaggedCollection 11 | 12 | sns.set_style("darkgrid") 13 | 14 | 15 | def plot( 16 | df: pd.DataFrame, 17 | x_axis, 18 | y_axis, 19 | compare, 20 | reference_y=None, 21 | title="UNKNOWN", 22 | xlabel=None, 23 | ylabel=None, 24 | ax: plt.Axes = None, 25 | fig: plt.Figure = None, 26 | linewidth: int | None = None, 27 | ): 28 | if fig is None or ax is None: 29 | fig, ax = plt.subplots(figsize=(12, 6), constrained_layout=True, dpi=300) 30 | 31 | if reference_y: 32 | ys = df.apply(lambda row: reference_y[row[X]], axis=1) 33 | print(ys) 34 | df[Y] = (df[Y] - ys) / ys 35 | 36 | def t(x): 37 | return np.percentile(x, q=20), np.percentile(x, q=80) 38 | 39 | sns.lineplot( 40 | ax=ax, 41 | data=df, 42 | x=x_axis, 43 | y=y_axis, 44 | estimator="median", 45 | err_kws={"alpha": 0.15}, 46 | hue=compare, 47 | # hue=None, 48 | # errorbar=t, 49 | errorbar=("pi", 50), 50 | # errorbar=lambda x: 1, 51 | # errorbar=None, 52 | # markers=True, 53 | ) 54 | 55 | if linewidth is not None: 56 | for line in ax.lines: 57 | line.set_linewidth(2) 58 | 59 | if xlabel: 60 | ax.set_xlabel(xlabel) 61 | if ylabel: 62 | ax.set_ylabel(ylabel) 63 | ax.set_title(title) 64 | 65 | 66 | ######################### 67 | 68 | X = "data.layers" 69 | Y = "data.throughput" 70 | 71 | filters = (lambda row: "call_optimization" in row.tags,) 72 | 73 | tagged = TaggedCollection.from_dllogs("benchmarks/good_logs/bench1.jsonl") 74 | 75 | tagged = tagged.filter(*filters) 76 | assert len(tagged) > 0, "Empty query!" 77 | 78 | references = tagged["vanilla"]["call_optimization"] 79 | reference_y_for_x = {} 80 | 81 | for x, reference_for_x in references.groupby(lambda row: row[X], return_keys=True): 82 | avg = np.median(reference_for_x.get(Y)) 83 | reference_y_for_x[x] = avg 84 | 85 | fig, ax = plt.subplots(nrows=2, ncols=1, figsize=(12, 10), constrained_layout=True, dpi=300) 86 | 87 | names = { 88 | "oct20,vanilla,call_optimization": "Inheriting from nn.Module", 89 | "oct20,sequential,call_optimization": "nn.Sequential", 90 | "oct20,symbolic,call_optimization": "SymbolicModel", 91 | } 92 | 93 | for tags, group in tagged.groupby("tags", return_keys=True): 94 | assert tags in names 95 | name = names[tags] 96 | for row in group: 97 | row["Definition"] = name 98 | 99 | plot( 100 | tagged.df, 101 | x_axis=X, 102 | y_axis=Y, 103 | compare="Definition", 104 | reference_y=None, 105 | title="Performance comparison between different definitions of deep linear models", 106 | xlabel="Number of linear layers", 107 | ylabel="Batches per second (more is better)", 108 | ax=ax[0], 109 | fig=fig, 110 | ) 111 | plot( 112 | tagged.df, 113 | x_axis=X, 114 | y_axis=Y, 115 | compare="Definition", 116 | reference_y=reference_y_for_x, 117 | title="Close-up", 118 | xlabel="Number of linear layers", 119 | ylabel="Throughput difference (%)", 120 | ax=ax[1], 121 | fig=fig, 122 | ) 123 | plt.savefig("latest.png") 124 | plt.show() 125 | -------------------------------------------------------------------------------- /tests/test_creating_multi_in_out.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2022 Szymon Mikler 2 | 3 | import torch 4 | from torch import nn 5 | 6 | from pytorch_symbolic import Input, SymbolicModel, model_tools 7 | 8 | NUM_INPUTS = 17 9 | NUM_LAYERS = 11 10 | FEATURES = 7 11 | SEED = 2137 12 | 13 | 14 | class WeightedConcatLayer(nn.Module): 15 | def __init__(self, dim): 16 | """Strange layer with multiple inputs where their correct order is important.""" 17 | super().__init__() 18 | self.dim = dim 19 | 20 | def forward(self, *tensors): 21 | tensors = [tensor * (i + 1) for i, tensor in enumerate(tensors)] 22 | return torch.cat(tensors=tensors, dim=self.dim) 23 | 24 | 25 | class VanillaModel(nn.Module): 26 | def __init__(self): 27 | super().__init__() 28 | 29 | self.layers = [nn.Linear(NUM_INPUTS * FEATURES, FEATURES) for _ in range(NUM_LAYERS)] 30 | for idx, layer in enumerate(self.layers): 31 | self.register_module(str(idx), layer) 32 | 33 | self.wconcat = WeightedConcatLayer(dim=1) 34 | self.final_linear = nn.Linear(NUM_LAYERS * FEATURES, FEATURES) 35 | self.flatten = nn.Flatten() 36 | 37 | def forward(self, *inputs): 38 | concatenated = self.wconcat(*inputs) 39 | transformed = [layer(concatenated) for layer in self.layers] 40 | concatenated2 = self.wconcat(*transformed) 41 | concatenated2 = self.final_linear(concatenated2) 42 | return concatenated2 - sum(transformed) 43 | 44 | 45 | def test_multiple_inputs(): 46 | inputs = [Input(shape=(FEATURES,)) for _ in range(NUM_INPUTS)] 47 | torch.manual_seed(SEED) 48 | 49 | concatenated = WeightedConcatLayer(dim=1)(*inputs) 50 | transformed = [nn.Linear(concatenated.features, FEATURES)(concatenated) for _ in range(NUM_LAYERS)] 51 | concatenated2 = WeightedConcatLayer(dim=1)(*transformed) 52 | concatenated2 = nn.Linear(concatenated2.features, FEATURES)(concatenated2) 53 | result = concatenated2 - sum(transformed) 54 | model = SymbolicModel(inputs=inputs, outputs=result) 55 | 56 | torch.manual_seed(SEED) 57 | vanilla_model = VanillaModel() 58 | 59 | assert model_tools.model_similar(model, vanilla_model) 60 | assert model_tools.models_have_corresponding_parameters(model, vanilla_model) 61 | 62 | for i in range(10): 63 | xs = [torch.rand(i, FEATURES) for _ in range(NUM_INPUTS)] 64 | out = model(*xs) 65 | vanilla_out = vanilla_model(*xs) 66 | assert torch.equal(out, vanilla_out), str((out, vanilla_out)) 67 | 68 | 69 | def test_detached_multiple_inputs(): 70 | inputs = [Input(shape=(FEATURES,)) for _ in range(NUM_INPUTS)] 71 | torch.manual_seed(SEED) 72 | 73 | concatenated = WeightedConcatLayer(dim=1)(*inputs) 74 | transformed = [nn.Linear(concatenated.features, FEATURES)(concatenated) for _ in range(NUM_LAYERS)] 75 | concatenated2 = WeightedConcatLayer(dim=1)(*transformed) 76 | concatenated2 = nn.Linear(concatenated2.features, FEATURES)(concatenated2) 77 | result = concatenated2 - sum(transformed) 78 | model = SymbolicModel(inputs=inputs, outputs=result).detach_from_graph() 79 | 80 | torch.manual_seed(SEED) 81 | vanilla_model = VanillaModel() 82 | 83 | assert model_tools.model_similar(model, vanilla_model) 84 | assert model_tools.models_have_corresponding_parameters(model, vanilla_model) 85 | 86 | for i in range(10): 87 | xs = [torch.rand(i, FEATURES) for _ in range(NUM_INPUTS)] 88 | out = model(*xs) 89 | vanilla_out = vanilla_model(*xs) 90 | assert torch.equal(out, vanilla_out), str((out, vanilla_out)) 91 | -------------------------------------------------------------------------------- /docs/reference/symbolic_data.md: -------------------------------------------------------------------------------- 1 | # symbolic_data 2 | 3 | This is a collection of symbolic data types. 4 | 5 | The main class and a grandfather for all other classes is Symbolic Data. 6 | 7 | Symbolic Tensor has similar API to ``torch.Tensor`` object, 8 | but Symbolic Tensor is used only to define the graph, not to perform actual computations. 9 | You should use it to register new layers in your computation graph and later to 10 | create the model. 11 | You can use methods of `torch.Tensor` when working with Symbolic Tensor. 12 | For example `tensor.t()` or `tensor.T` will return another Symbolic Tensor 13 | with transposition applied to the underlying data. 14 | 15 | Symbolic Data supports slicing too, so you can do: 16 | ```python 17 | from pytorch_symbolic import Input, SymbolicModel 18 | 19 | x = Input(batch_shape=(3, 4, 5)) 20 | 21 | y = x[0] 22 | for row in x[1:]: 23 | y += row 24 | 25 | model = SymbolicModel(x, y) 26 | ``` 27 | 28 | But be careful! Each slice operation creates a new layer, 29 | so if you do a lot of slicing, 30 | it is better enclose it in a custom module. 31 | However, being able to do it directly on Symbolic Data is convenient for prototyping. 32 | 33 | 34 | ::: pytorch_symbolic.Input 35 | options: 36 | show_source: false 37 | heading_level: 2 38 | show_root_heading: true 39 | members_order: source 40 | show_object_full_path: false 41 | docstring_section_style: table 42 | show_signature_annotations: true 43 | separate_signature: true 44 | annotations_path: brief 45 | merge_init_into_class: true 46 | show_root_full_path: true 47 | 48 | ::: pytorch_symbolic.CustomInput 49 | options: 50 | show_source: false 51 | heading_level: 2 52 | show_root_heading: true 53 | members_order: source 54 | show_object_full_path: false 55 | docstring_section_style: table 56 | show_signature_annotations: true 57 | separate_signature: true 58 | annotations_path: brief 59 | merge_init_into_class: true 60 | show_root_full_path: true 61 | 62 | ::: pytorch_symbolic.symbolic_data.SymbolicData 63 | options: 64 | show_source: false 65 | heading_level: 2 66 | show_root_heading: true 67 | members_order: source 68 | show_object_full_path: false 69 | docstring_section_style: table 70 | show_signature_annotations: true 71 | separate_signature: true 72 | annotations_path: brief 73 | merge_init_into_class: true 74 | show_root_full_path: true 75 | 76 | ::: pytorch_symbolic.symbolic_data.SymbolicCallable 77 | options: 78 | show_source: false 79 | heading_level: 2 80 | show_root_heading: true 81 | members_order: source 82 | show_object_full_path: false 83 | docstring_section_style: table 84 | show_signature_annotations: true 85 | separate_signature: true 86 | annotations_path: brief 87 | merge_init_into_class: true 88 | show_root_full_path: true 89 | 90 | ::: pytorch_symbolic.symbolic_data.SymbolicTensor 91 | options: 92 | show_source: false 93 | heading_level: 2 94 | show_root_heading: true 95 | members_order: source 96 | show_object_full_path: false 97 | docstring_section_style: table 98 | show_signature_annotations: true 99 | separate_signature: true 100 | annotations_path: brief 101 | merge_init_into_class: true 102 | show_root_full_path: true 103 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Pytorch Symbolic 2 | 3 | [//]: # (To get badges go to https://shields.io/ and use https://pypi.org/pypi/slicemap/json as data url. Query fields using dot as the separator.) 4 | 5 | [![PyPi version](https://img.shields.io/badge/dynamic/json?label=latest&query=info.version&url=https%3A%2F%2Fpypi.org%2Fpypi%2Fpytorch-symbolic%2Fjson)](https://pypi.org/project/pytorch-symbolic) 6 | [![PyPI license](https://img.shields.io/badge/dynamic/json?label=license&query=info.license&url=https%3A%2F%2Fpypi.org%2Fpypi%2Fpytorch-symbolic%2Fjson)](https://github.com/sjmikler/pytorch-symbolic/blob/main/LICENSE.txt) 7 | [![Notebook](https://github.com/gahaalt/pytorch-symbolic/actions/workflows/run-notebook.yaml/badge.svg)](https://github.com/gahaalt/pytorch-symbolic/actions/workflows/notebook.yaml) 8 | 9 | Pytorch Symbolic is MIT licensed library that adds symbolic API for model creation to PyTorch. 10 | 11 | Pytorch Symbolic makes it easier and faster to define complex models. 12 | It spares you writing boilerplate code. 13 | It aims to be PyTorch equivalent for [Keras Functional API](https://keras.io/guides/functional_api/). 14 | 15 | Features: 16 | 17 | * Small extension of PyTorch 18 | * No dependencies besides PyTorch 19 | * Produces models entirely compatible with PyTorch 20 | * Overhead free as tested in [benchmarks](docs/benchmarks.md) 21 | * Reduces the amount of boilerplate code 22 | * Works well with complex architectures 23 | * Code and documentation is automatically tested 24 | 25 | ## Example 26 | 27 | To create a symbolic model, you need Symbolic Tensors and `torch.nn.Module`. 28 | Register layers and operations in your model by calling ``layer(inputs)`` or 29 | equivalently ``inputs(layer)``. 30 | Layers will be automagically added to your model and 31 | all operations will be replayed on the real data. 32 | That's all! 33 | 34 | Using Pytorch Symbolic, we can define a working classifier in a few lines of code: 35 | 36 | ```python 37 | from torch import nn 38 | from pytorch_symbolic import Input, SymbolicModel 39 | 40 | inputs = Input(shape=(1, 28, 28)) 41 | x = nn.Flatten()(inputs) 42 | x = nn.Linear(x.shape[1], 10)(x)(nn.Softmax(1)) 43 | model = SymbolicModel(inputs=inputs, outputs=x) 44 | model.summary() 45 | ``` 46 | 47 | ```stdout 48 | _______________________________________________________ 49 | Layer Output shape Params Parent 50 | ======================================================= 51 | 1 Input_1 (None, 1, 28, 28) 0 52 | 2 Flatten_1 (None, 784) 0 1 53 | 3 Linear_1 (None, 10) 7850 2 54 | 4* Softmax_1 (None, 10) 0 3 55 | ======================================================= 56 | Total params: 7850 57 | Trainable params: 7850 58 | Non-trainable params: 0 59 | _______________________________________________________ 60 | ``` 61 | 62 | **See more examples 63 | in [Documentation Quick Start](docs/quick_start.md).** 64 | 65 | ## How to start 66 | 67 | See Jupyter Notebook showing the basic usage of Pytorch Symbolic: 68 | 69 | * Learn Pytorch Symbolic in an interactive way. 70 | * Try the package before installing it on your computer. 71 | * See visualizations of graphs that are created under the hood. 72 | 73 | Open in Colab: 74 | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/gahaalt/pytorch-symbolic/blob/develop/introduction.ipynb) 75 | 76 | ## Installation 77 | 78 | Install Pytorch Symbolic easily with pip: 79 | 80 | ```bash 81 | pip install pytorch-symbolic 82 | ``` 83 | 84 | ## Troubleshooting 85 | 86 | Please create an issue if you notice a problem! 87 | 88 | ## Links 89 | 90 | * [See Documentation](https://github.com/sjmikler/pytorch-symbolic/tree/main/docs) 91 | * [See on GitHub](https://github.com/gahaalt/pytorch-symbolic/) 92 | * [See on PyPI](https://pypi.org/project/pytorch-symbolic/) 93 | 94 | -------------------------------------------------------------------------------- /pytorch_symbolic/functions_utility.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2022 Szymon Mikler 2 | 3 | from __future__ import annotations 4 | 5 | from typing import Callable, Dict, Hashable, List, Tuple 6 | 7 | from torch import nn 8 | 9 | from . import useful_layers 10 | from .symbolic_data import SymbolicData 11 | 12 | 13 | def _replace_symbolic_with_value(container, extracted, navigation): 14 | """Search recursively for all occurences of Symbolic and replace them with their value. 15 | 16 | At the same time save navigation to know, how to do indexing to get to them. 17 | 18 | If navigation for container ends up as [..., [1, 2, 0, "TEST", 5], ...] 19 | this means that to get the element you should index container[1][2][0]["TEST"][5]. 20 | """ 21 | if isinstance(container, SymbolicData): 22 | navigation.append(navigation[-1].copy()) 23 | extracted.append(container) 24 | return container.v 25 | 26 | if isinstance(container, List) or isinstance(container, Tuple): 27 | new_list = [] 28 | for i, x in enumerate(container): 29 | navigation[-1].append(i) 30 | new_list.append(_replace_symbolic_with_value(x, extracted, navigation)) 31 | navigation[-1].pop() 32 | return new_list 33 | if isinstance(container, Dict): 34 | new_dict = {} 35 | for k, v in container.items(): 36 | navigation[-1].append(k) 37 | new_dict[k] = _replace_symbolic_with_value(v, extracted, navigation) 38 | navigation[-1].pop() 39 | return new_dict 40 | return container 41 | 42 | 43 | def add_to_graph(func: Callable | nn.Module, *args, custom_name: str | None = None, **kwds): 44 | """Register a custom func or a module in the computation graph. 45 | 46 | This works will arbitrary functions and modules iff at least one Symbolic Data is among ``*args, **kwds``. 47 | 48 | This way of registering is flexible, but might add a small slowdown to the call, 49 | because it adds a wrapper for parsing arguments. 50 | If this is unacceptable, please create an torch.nn.Module that takes only Symbolic Data arguments. 51 | 52 | Here all arguments, including Symbolic Data, should be passed after the ``func`` argument. 53 | The arguments can be mixed and matched, even nested in lists, tuples and dictionaries. 54 | 55 | Convolution func example:: 56 | 57 | inputs = Input(shape=(3, 32, 32)) 58 | kernel = Input(batch_size=(16, 3, 3, 3)) 59 | bias = Input(batch_size=(16,)) 60 | output = add_to_graph(F.conv2d, input=inputs, weight=k, bias=bias, padding=1) 61 | """ 62 | extracted_symbols: List[SymbolicData] = [] 63 | 64 | real_call_args = [] 65 | real_call_kwds = {} 66 | 67 | navigation: List[List[Hashable]] = [[]] 68 | real_call_args = _replace_symbolic_with_value(args, extracted_symbols, navigation) 69 | real_call_kwds = _replace_symbolic_with_value(kwds, extracted_symbols, navigation) 70 | navigation.pop() 71 | 72 | assert len(extracted_symbols) > 0, "No Symbolic Data detected in the input!" 73 | assert all((isinstance(symbol, SymbolicData) for symbol in extracted_symbols)) 74 | assert len(extracted_symbols) == len(navigation) 75 | 76 | def wrapper_function(*args): 77 | assert len(args) == len(navigation), f"Expected {len(navigation)} inputs, not {len(args)}!" 78 | for arg, navi in zip(args, navigation): 79 | obj = real_call_kwds if isinstance(navi[0], str) else real_call_args 80 | 81 | for idx in navi[:-1]: 82 | obj = obj[idx] 83 | 84 | obj[navi[-1]] = arg 85 | 86 | return func(*real_call_args, **real_call_kwds) 87 | 88 | if hasattr(func, "__name__"): 89 | name = func.__name__ 90 | elif hasattr(func, "__class__"): 91 | name = func.__class__.__name__ 92 | else: 93 | name = str(func) 94 | module = useful_layers.NamedLambdaOpLayer(op=wrapper_function, name=f"wrap({name})") 95 | # This might be a Symbolic Callable, so we use `apply_module` instead of `__call__` 96 | return extracted_symbols[0].apply_module(module, *extracted_symbols[1:], custom_name=custom_name) 97 | -------------------------------------------------------------------------------- /pytorch_symbolic/useful_layers.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2022 Szymon Mikler 2 | 3 | from typing import Callable 4 | 5 | import torch 6 | from torch import nn 7 | 8 | 9 | class LambdaOpLayer(nn.Module): 10 | def __init__(self, op: Callable): 11 | super().__init__() 12 | self.forward = op 13 | 14 | 15 | class NamedLambdaOpLayer(nn.Module): 16 | def __init__(self, op: Callable, name: str): 17 | super().__init__() 18 | self.forward = op 19 | self.name = name 20 | 21 | def __repr__(self): 22 | return self.name 23 | 24 | def _get_name(self): 25 | return self.name 26 | 27 | 28 | class CallbackLayer(nn.Module): 29 | def __init__(self, op: Callable): 30 | """Layer that returns its inputs, but executes a callable ``op`` on them before returning. 31 | 32 | This can be used for debugging or logging purposes. Accepts only one argument. 33 | 34 | Example:: 35 | 36 | x = CallbackLayer(print)(x) 37 | 38 | It does not change anything in x, but prints its value. 39 | """ 40 | super().__init__() 41 | self.callback = op 42 | 43 | def forward(self, arg): 44 | self.callback(arg) 45 | return arg 46 | 47 | 48 | class AddOpLayer(nn.Module): 49 | @staticmethod 50 | def forward(a, b): 51 | return a + b 52 | 53 | 54 | class SubOpLayer(nn.Module): 55 | @staticmethod 56 | def forward(a, b): 57 | return a - b 58 | 59 | 60 | class MulOpLayer(nn.Module): 61 | @staticmethod 62 | def forward(a, b): 63 | return a * b 64 | 65 | 66 | class ModOpLayer(nn.Module): 67 | @staticmethod 68 | def forward(a, b): 69 | return torch.remainder(a, b) 70 | 71 | 72 | class MatmulOpLayer(nn.Module): 73 | @staticmethod 74 | def forward(a, b): 75 | return a @ b 76 | 77 | 78 | class ConcatLayer(nn.Module): 79 | def __init__(self, dim): 80 | super().__init__() 81 | self.dim = dim 82 | 83 | def forward(self, *tensors): 84 | return torch.cat(tensors=tensors, dim=self.dim) 85 | 86 | 87 | class StackLayer(nn.Module): 88 | def __init__(self, dim): 89 | super().__init__() 90 | self.dim = dim 91 | 92 | def forward(self, *tensors): 93 | return torch.stack(tensors=tensors, dim=self.dim) 94 | 95 | 96 | class ReshapeLayer(nn.Module): 97 | def __init__(self, shape, batch_size_included=False): 98 | super().__init__() 99 | if not batch_size_included: 100 | self.shape = (-1, *shape) 101 | else: 102 | self.shape = shape 103 | 104 | def forward(self, tensor): 105 | return torch.reshape(input=tensor, shape=self.shape) 106 | 107 | 108 | class ViewCopyLayer(nn.Module): 109 | def __init__(self, shape, batch_size_included=False): 110 | super().__init__() 111 | if not batch_size_included: 112 | self.shape = (-1, *shape) 113 | else: 114 | self.shape = shape 115 | 116 | def forward(self, tensor): 117 | return torch.view_copy(input=tensor, size=self.shape) 118 | 119 | 120 | class AggregateLayer(nn.Module): 121 | def __init__(self, op: Callable, dim=None, keepdim=False): 122 | super().__init__() 123 | self.op = op 124 | 125 | self.dim = dim 126 | self.keepdim = keepdim 127 | if self.dim is None: 128 | self.forward = self.forward_nodim 129 | 130 | def forward(self, tensor): 131 | return self.op(input=tensor, dim=self.dim, keepdim=self.keepdim) 132 | 133 | def forward_nodim(self, tensor): 134 | return self.op(input=tensor) 135 | 136 | 137 | class UnpackLayer(nn.Module): 138 | def forward(self, *args): 139 | return args 140 | 141 | 142 | class SliceLayer(nn.Module): 143 | def __init__(self, idx): 144 | super().__init__() 145 | self.idx = idx 146 | 147 | def forward(self, arg): 148 | return arg[self.idx] 149 | 150 | 151 | class SliceLayerSymbolicIdx(nn.Module): 152 | def forward(self, arg, idx): 153 | return arg[idx] 154 | 155 | 156 | class MethodCall(nn.Module): 157 | def __init__(self, method_name, *args, **kwds): 158 | super().__init__() 159 | self.method_name = method_name 160 | self.args = args 161 | self.kwds = kwds 162 | 163 | def forward(self, tensor): 164 | return getattr(tensor, self.method_name)(*self.args, **self.kwds) 165 | 166 | 167 | class GetAttr(nn.Module): 168 | def __init__(self, name): 169 | super().__init__() 170 | self.name = name 171 | 172 | def forward(self, data): 173 | return getattr(data, self.name) 174 | -------------------------------------------------------------------------------- /tests/test_add_to_graph.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2022 Szymon Mikler 2 | 3 | import torch 4 | import torch.nn.functional as F 5 | 6 | from pytorch_symbolic import Input, SymbolicModel, add_to_graph 7 | 8 | 9 | def test_func_concat(): 10 | in1 = Input((10, 20)) 11 | in2 = Input((10, 20)) 12 | 13 | cat1 = add_to_graph(torch.concat, (in1, in2), dim=2) 14 | cat2 = add_to_graph(torch.concat, tensors=(in1, in2), dim=2) 15 | 16 | model1 = SymbolicModel(inputs=(in1, in2), outputs=cat1) 17 | model2 = SymbolicModel(inputs=(in1, in2), outputs=cat2) 18 | 19 | for _ in range(10): 20 | x1 = torch.rand(1, 10, 20) 21 | x2 = torch.rand(1, 10, 20) 22 | correct_result = torch.concat((x1, x2), dim=2) 23 | assert torch.equal(model1(x1, x2), correct_result) 24 | assert torch.equal(model2(x1, x2), correct_result) 25 | 26 | 27 | def test_detached_func_concat(): 28 | in1 = Input((10, 20)) 29 | in2 = Input((10, 20)) 30 | 31 | cat1 = add_to_graph(torch.concat, (in1, in2), dim=2) 32 | cat2 = add_to_graph(torch.concat, tensors=(in1, in2), dim=2) 33 | 34 | model1 = SymbolicModel(inputs=(in1, in2), outputs=cat1).detach_from_graph() 35 | model2 = SymbolicModel(inputs=(in1, in2), outputs=cat2).detach_from_graph() 36 | 37 | for _ in range(10): 38 | x1 = torch.rand(1, 10, 20) 39 | x2 = torch.rand(1, 10, 20) 40 | correct_result = torch.concat((x1, x2), dim=2) 41 | assert torch.equal(model1(x1, x2), correct_result) 42 | assert torch.equal(model2(x1, x2), correct_result) 43 | 44 | 45 | def sum_list_recursively(multiply_result, container): 46 | container = container.copy() 47 | for idx, el in enumerate(container): 48 | if isinstance(el, list): 49 | container[idx] = sum_list_recursively(1, el) 50 | return sum(container) * multiply_result 51 | 52 | 53 | def test_func_complicated_input(): 54 | sym1 = Input(batch_shape=(10, 20)) 55 | sym2 = Input(batch_shape=(10, 20)) 56 | sym3 = Input(batch_shape=(10, 20)) 57 | 58 | other = torch.rand(10, 20) 59 | 60 | summed1 = add_to_graph( 61 | sum_list_recursively, 62 | 5, 63 | [ 64 | other, 65 | other, 66 | [other, other, sym1], 67 | [other, other, [other, other, [other]]], 68 | [other, sym2, other], 69 | [[[other], [[[[other, other]]], [other, [sym3, [other]]]]]], 70 | ], 71 | ) 72 | summed2 = add_to_graph(sum_list_recursively, multiply_result=5, container=[[other] * 16, [sym1, [sym2, [sym3]]]]) 73 | 74 | model1 = SymbolicModel((sym1, sym2, sym3), outputs=summed1) 75 | model2 = SymbolicModel((sym1, sym2, sym3), outputs=summed2) 76 | 77 | x1 = torch.rand(10, 20) 78 | x2 = torch.rand(10, 20) 79 | x3 = torch.rand(10, 20) 80 | 81 | r1 = model1(x1, x2, x3) 82 | r2 = model2(x1, x2, x3) 83 | assert torch.allclose(r1, r2) 84 | 85 | correct_result = 5 * (other * 16 + x1 + x2 + x3) 86 | assert torch.allclose(r1, correct_result) 87 | 88 | 89 | def test_detached_func_complicated_input(): 90 | sym1 = Input(batch_shape=(10, 20)) 91 | sym2 = Input(batch_shape=(10, 20)) 92 | sym3 = Input(batch_shape=(10, 20)) 93 | 94 | other = torch.rand(10, 20) 95 | 96 | summed1 = add_to_graph( 97 | sum_list_recursively, 98 | 5, 99 | [ 100 | other, 101 | other, 102 | [other, other, sym1], 103 | [other, other, [other, other, [other]]], 104 | [other, sym2, other], 105 | [[[other], [[[[other, other]]], [other, [sym3, [other]]]]]], 106 | ], 107 | ) 108 | summed2 = add_to_graph(sum_list_recursively, multiply_result=5, container=[[other] * 16, [sym1, [sym2, [sym3]]]]) 109 | 110 | model1 = SymbolicModel((sym1, sym2, sym3), outputs=summed1).detach_from_graph() 111 | model2 = SymbolicModel((sym1, sym2, sym3), outputs=summed2).detach_from_graph() 112 | 113 | x1 = torch.rand(10, 20) 114 | x2 = torch.rand(10, 20) 115 | x3 = torch.rand(10, 20) 116 | 117 | r1 = model1(x1, x2, x3) 118 | r2 = model2(x1, x2, x3) 119 | assert torch.allclose(r1, r2) 120 | 121 | correct_result = 5 * (other * 16 + x1 + x2 + x3) 122 | assert torch.allclose(r1, correct_result) 123 | 124 | 125 | def test_conv(): 126 | inputs = Input(shape=(3, 32, 32)) 127 | kernel = Input(batch_shape=(16, 3, 3, 3)) 128 | bias = Input(batch_shape=(16,)) 129 | output = add_to_graph(F.conv2d, input=inputs, weight=kernel, bias=bias, padding=1) 130 | model = SymbolicModel((inputs, kernel, bias), output) 131 | 132 | i = torch.rand(10, 3, 32, 32) 133 | k = torch.rand(16, 3, 3, 3) 134 | b = torch.rand(16) 135 | assert torch.allclose(model(i, k, b), F.conv2d(i, k, b, padding=1)) 136 | -------------------------------------------------------------------------------- /examples/resnet.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2022 Szymon Mikler 2 | 3 | """ 4 | This is a flexible implementation of ResNet architecture. 5 | It allows for creation of standard ResNet v2 or Wide ResNet variants. 6 | """ 7 | 8 | from torch import nn 9 | 10 | from pytorch_symbolic import Input, SymbolicModel 11 | 12 | from .common import classifier 13 | 14 | 15 | def shortcut_func(x, channels, stride): 16 | if x.channels != channels or stride != 1: 17 | return x(nn.Conv2d(x.channels, channels, kernel_size=1, bias=False, stride=stride)) 18 | else: 19 | return x 20 | 21 | 22 | def ToyResNet(input_shape, n_classes): 23 | """A basic example of ResNet with low degree of customizability.""" 24 | 25 | inputs = Input(input_shape) 26 | flow = inputs(nn.Conv2d(inputs.channels, 16, 3, 1, 1)) 27 | 28 | for group_size, width, stride in zip((2, 2, 2), (16, 32, 64), (1, 2, 2)): 29 | for _ in range(group_size): 30 | shortcut = shortcut_func(flow, width, stride) 31 | flow = flow(nn.BatchNorm2d(flow.channels))(nn.ReLU()) 32 | flow = flow(nn.Conv2d(flow.channels, width, 3, stride, 1)) 33 | flow = flow(nn.BatchNorm2d(flow.channels))(nn.ReLU()) 34 | flow = flow(nn.Conv2d(flow.channels, width, 3, 1, 1)) 35 | 36 | flow = flow + shortcut 37 | stride = 1 38 | flow = flow(nn.BatchNorm2d(flow.channels))(nn.ReLU()) 39 | outs = classifier(flow, n_classes, pooling="avgpool") 40 | model = SymbolicModel(inputs=inputs, outputs=outs) 41 | return model 42 | 43 | 44 | relu = nn.ReLU() 45 | 46 | 47 | def ResNet( 48 | input_shape, 49 | n_classes, 50 | version=None, 51 | bootleneck=False, 52 | strides=(1, 2, 2), 53 | group_sizes=(2, 2, 2), 54 | channels=(16, 32, 64), 55 | activation=relu, 56 | final_pooling="avgpool", 57 | dropout=0, 58 | bn_ends_block=False, 59 | **kwargs, 60 | ): 61 | if version: 62 | if version == 20: 63 | group_sizes = (3, 3, 3) 64 | elif version == 32: 65 | group_sizes = (5, 5, 5) 66 | elif version == 44: 67 | group_sizes = (7, 7, 7) 68 | elif version == 56: 69 | group_sizes = (9, 9, 9) 70 | elif version == 110: 71 | group_sizes = (18, 18, 18) 72 | elif version == 164: 73 | bootleneck = True 74 | channels = (64, 128, 256) 75 | group_sizes = (18, 18, 18) 76 | elif isinstance(version, tuple) and version[0] == "WRN": 77 | _, N, K = version 78 | assert (N - 4) % 6 == 0, "N-4 has to be divisible by 6" 79 | lpb = (N - 4) // 6 # layers per block 80 | group_sizes = (lpb, lpb, lpb) 81 | channels = tuple(c * K for c in channels) 82 | else: 83 | raise NotImplementedError(f"Unkown version={version}!") 84 | if kwargs: 85 | print(f"ResNet: unknown parameters: {kwargs.keys()}") 86 | 87 | def simple_block(flow, channels, stride): 88 | if preactivate_block: 89 | flow = flow(nn.BatchNorm2d(flow.channels))(activation) 90 | 91 | flow = flow(nn.Conv2d(flow.channels, channels, 3, stride, 1)) 92 | flow = flow(nn.BatchNorm2d(flow.channels))(activation) 93 | 94 | if dropout: 95 | flow = flow(nn.Dropout(p=dropout)) 96 | flow = flow(nn.Conv2d(flow.channels, channels, 3, 1, 1)) 97 | 98 | if bn_ends_block: 99 | flow = flow(nn.BatchNorm2d(flow.channels))(activation) 100 | return flow 101 | 102 | def bootleneck_block(flow, channels, stride): 103 | if preactivate_block: 104 | flow = flow(nn.BatchNorm2d(flow.channels))(activation) 105 | 106 | flow = flow(nn.Conv2d(flow.channels, channels // 4, 1)) 107 | flow = flow(nn.BatchNorm2d(flow.channels))(activation) 108 | 109 | flow = flow(nn.Conv2d(flow.channels, channels // 4, 3, stride=stride, padding=1)) 110 | flow = flow(nn.BatchNorm2d(flow.channels))(activation) 111 | 112 | flow = flow(nn.Conv2d(flow.channels, channels, 1)) 113 | if bn_ends_block: 114 | flow = flow(nn.BatchNorm2d(flow.channels))(activation) 115 | return flow 116 | 117 | if bootleneck: 118 | block = bootleneck_block 119 | else: 120 | block = simple_block 121 | 122 | inputs = Input(input_shape) 123 | 124 | # BUILDING HEAD OF THE NETWORK 125 | flow = inputs(nn.Conv2d(inputs.channels, 16, 3, 1, 1)) 126 | 127 | # BUILD THE RESIDUAL BLOCKS 128 | for group_size, width, stride in zip(group_sizes, channels, strides): 129 | flow = flow(nn.BatchNorm2d(flow.channels))(activation) 130 | preactivate_block = False 131 | 132 | for _ in range(group_size): 133 | residual = block(flow, width, stride) 134 | shortcut = shortcut_func(flow, width, stride) 135 | flow = residual + shortcut 136 | preactivate_block = True 137 | stride = 1 138 | 139 | # BUILDING THE CLASSIFIER 140 | flow = flow(nn.BatchNorm2d(flow.channels))(activation) 141 | outs = classifier(flow, n_classes, pooling=final_pooling) 142 | model = SymbolicModel(inputs=inputs, outputs=outs) 143 | return model 144 | -------------------------------------------------------------------------------- /pytorch_symbolic/code_generator.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2022 Szymon Mikler 2 | 3 | from __future__ import annotations 4 | 5 | from typing import List, Set, Tuple 6 | 7 | from pytorch_symbolic.symbolic_data import SymbolicData 8 | 9 | 10 | def generate_forward_with_loops( 11 | inputs: List[SymbolicData] | Tuple[SymbolicData, ...], 12 | outputs: List[SymbolicData] | Tuple[SymbolicData, ...], 13 | execution_order: List[SymbolicData] | Tuple[SymbolicData, ...], 14 | nodes_in_subgraph: Set[SymbolicData], 15 | min_loop_length: int | float = float("inf"), 16 | ) -> str: 17 | """Generate code for forward function of SymbolicModel. 18 | 19 | It assumes there is `self._execution_order_layers` available in the class. 20 | 21 | Parameters 22 | ---------- 23 | inputs 24 | Inputs to the model 25 | outputs 26 | Outputs of the model 27 | execution_order 28 | Contains the exact order in which the nodes should be executed. 29 | If there are layers with multiple outputs, this will be a subset of `nodes_in_subgraph`. 30 | In such case, only one output of each layer needs to be in the execution_order. 31 | nodes_in_subgraph 32 | All nodes covered by the subgraph, including all nodes created by multiple-output layers. 33 | min_loop_length 34 | Minimal sequence length to replace sequential layers execution with a loop. 35 | 36 | Returns 37 | ------- 38 | str 39 | Generated code. 40 | """ 41 | assert min_loop_length >= 2, "Loop length cannot be smaller than 2!" 42 | 43 | str_length = len(str(max(len(inputs), len(outputs), len(execution_order)))) 44 | node_to_name = {} 45 | for idx, node in enumerate(inputs): 46 | node_to_name[node] = f"i{str(idx).zfill(str_length)}" 47 | for idx, node in enumerate(execution_order): 48 | node_to_name[node] = f"x{str(idx).zfill(str_length)}" 49 | for idx, node in enumerate(nodes_in_subgraph.difference(execution_order)): 50 | node_to_name[node] = f"y{str(idx).zfill(str_length)}" 51 | for idx, node in enumerate(outputs): 52 | node_to_name[node] = f"o{str(idx).zfill(str_length)}" 53 | 54 | input_names = [node_to_name[node] for node in inputs] 55 | forward_definition = "def forward(self," + ", ".join(input_names) + "):" 56 | code_lines = [forward_definition] 57 | 58 | TAB = " " * 4 59 | code_lines.append(TAB + "l = self._execution_order_layers") 60 | 61 | nodes_looped_over = set() 62 | # All parents must be in the graph. Otherwise, forward is impossible. 63 | parents = {node: node.parents for node in execution_order} 64 | # We only count children in the graph. Thus the intersection. 65 | children = {node: list(nodes_in_subgraph.intersection(node.children)) for node in execution_order} 66 | 67 | siblings = {node: list(nodes_in_subgraph.intersection(node._layer_full_siblings)) for node in execution_order} 68 | 69 | for exec_id, node in enumerate(execution_order): 70 | if node in nodes_looped_over: 71 | continue 72 | 73 | input_names = [node_to_name[node] for node in node.parents] 74 | sequence = [node] 75 | last_node = node 76 | while ( 77 | len(children[last_node]) == 1 78 | # stop iterating when need to unpack something! 79 | and len(last_node._layer_full_siblings) == 1 80 | and len(children[last_node][0]._layer_full_siblings) == 1 # needed, else tests fail 81 | # this should never be false, but just in case we make sure the child is next in execution order 82 | and children[last_node][0] is execution_order[exec_id + len(sequence)] 83 | and len(parents[last_node]) == 1 84 | and len(parents[children[last_node][0]]) == 1 85 | ): 86 | last_node = children[last_node][0] 87 | sequence.append(last_node) 88 | 89 | if len(sequence) >= min_loop_length: 90 | output_name = node_to_name[sequence[-1]] 91 | code_lines.append(TAB + f"{output_name} = {input_names[0]}") 92 | code_lines.append(TAB + f"for layer in l[{exec_id}:{exec_id + len(sequence)}]:") 93 | code_lines.append(TAB + TAB + f"{output_name} = layer({output_name})") 94 | nodes_looped_over.update(sequence) 95 | elif len(node._layer_full_siblings) > 1: # Must unpack all siblings, even if not all are used 96 | output_names = [] 97 | for n in node._layer_full_siblings: 98 | if n in siblings[node]: 99 | output_names.append(node_to_name[n]) 100 | else: 101 | output_names.append("_") # If sibling not used, we don't save it as a variable 102 | 103 | assert len(input_names) == 1, "Layer that has full siblings cannot have more than 1 input!" 104 | code_line = TAB + ", ".join(output_names) + f" = l[{exec_id}](" + "*" + input_names[0] + ")" 105 | code_lines.append(code_line) 106 | else: 107 | code_line = TAB + node_to_name[node] + f" = l[{exec_id}](" + ", ".join(input_names) + ")" 108 | code_lines.append(code_line) 109 | 110 | code_lines.append(TAB + "return " + ", ".join(node_to_name[node] for node in outputs)) 111 | generated_forward = "\n".join(code_lines) + "\n" 112 | return generated_forward 113 | -------------------------------------------------------------------------------- /benchmarks/tagged_collection.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2022 Szymon Mikler 2 | 3 | import json 4 | from typing import Callable 5 | 6 | import numpy as np 7 | 8 | 9 | def flatten_nested_dict(data): 10 | for key, value in list(data.items()): 11 | if isinstance(value, dict): 12 | value = flatten_nested_dict(value) 13 | value = {str(key) + "." + str(k): v for k, v in value.items()} 14 | data.pop(key) 15 | data.update(value) 16 | return data 17 | 18 | 19 | class Tagged(dict): 20 | def __init__(self, tags=None, featured=None, **fields): 21 | flatten_fields = flatten_nested_dict(fields) 22 | self.featured = featured 23 | 24 | if tags is not None: 25 | self._tags = tags 26 | else: 27 | if "tags" in flatten_fields: 28 | self._tags = flatten_fields.pop("tags") 29 | elif "data.tags" in flatten_fields: 30 | self._tags = flatten_fields.pop("data.tags") 31 | else: 32 | raise KeyError("Tags are missing!") 33 | 34 | super().__init__(**fields) 35 | self._tags_str = ",".join(self.tags) 36 | self["tags"] = self._tags_str 37 | 38 | @property 39 | def tags(self): 40 | return self._tags 41 | 42 | def __repr__(self): 43 | if self.featured: 44 | details = "; " + str({k: self[k] for k in self.featured}) 45 | else: 46 | details = "" 47 | 48 | return f"Tagged {self.tags}{details}" 49 | 50 | 51 | class TaggedCollection(list): 52 | def __init__(self, data=None, featured=None): 53 | super().__init__() 54 | self._tags = [] 55 | self._tags_set = set() 56 | self.featured = featured 57 | 58 | if data: 59 | for row in data: 60 | self.convert_and_append(row) 61 | 62 | def convert_and_append(self, row): 63 | if not isinstance(row, Tagged): 64 | row = Tagged(**row, featured=self.featured) 65 | 66 | for tag in row.tags: 67 | if tag not in self._tags_set: 68 | self._tags.append(tag) 69 | self._tags_set.add(tag) 70 | self.append(row) 71 | 72 | def get(self, key): 73 | values = [] 74 | for record in self: 75 | assert key in record 76 | values.append(record[key]) 77 | return np.array(values) 78 | 79 | @property 80 | def tags(self): 81 | return self._tags 82 | 83 | def filter(self, *filters): 84 | new = TaggedCollection() 85 | for record in self: 86 | for flt in filters: 87 | if not flt(record): 88 | break 89 | else: 90 | new.convert_and_append(record) 91 | return new 92 | 93 | def restrict_tags(self, tags): 94 | new = TaggedCollection() 95 | tags_set = set(tags) 96 | 97 | for record in self: 98 | for tag in record.tags: 99 | if tag not in tags_set: 100 | break 101 | else: 102 | new.convert_and_append(record) 103 | return new 104 | 105 | def groupby(self, *filters, return_keys=False): 106 | if not filters: 107 | filters = [lambda x: True] 108 | 109 | groups = {} 110 | for record in self: 111 | results = [] 112 | for flt in filters: 113 | if isinstance(flt, Callable): 114 | result = flt(record) 115 | else: 116 | result = record[flt] 117 | results.append(result) 118 | 119 | if len(results) > 1: 120 | results = tuple(results) 121 | else: 122 | results = results[0] 123 | 124 | if results not in groups: 125 | groups[results] = TaggedCollection() 126 | groups[results].convert_and_append(record) 127 | if return_keys: 128 | return list((k, v) for k, v in groups.items()) 129 | else: 130 | return list(groups.values()) 131 | 132 | @classmethod 133 | def from_dllogs(cls, path, prefix="DLLL", featured=None): 134 | tc = cls(featured=featured) 135 | with open(path, "r") as f: 136 | for line in f.readlines(): 137 | if line.startswith(prefix): 138 | line = line[len(prefix) :].strip() 139 | data = json.loads(line) 140 | tc.convert_and_append(data) 141 | return tc 142 | 143 | def common_keys(self): 144 | keys = None 145 | for _idx, row in enumerate(self): 146 | if keys is None: 147 | keys = set(row.keys()) 148 | keys = keys.intersection(row.keys()) 149 | return keys 150 | 151 | def to_df(self): 152 | import pandas as pd 153 | 154 | return pd.DataFrame.from_records(self) 155 | 156 | @property 157 | def df(self): 158 | return self.to_df() 159 | 160 | def __getitem__(self, item): 161 | if isinstance(item, str): 162 | if item.startswith("!"): 163 | item = item[1:] 164 | return self.filter(lambda x: item not in x.tags) 165 | return self.filter(lambda x: item in x.tags) 166 | if isinstance(item, Callable): 167 | return self.filter(item) 168 | return super().__getitem__(item) 169 | -------------------------------------------------------------------------------- /examples/lstm.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2022 Szymon Mikler 2 | 3 | """ 4 | Original: 5 | https://pytorch.org/tutorials/beginner/nlp/sequence_models_tutorial.html 6 | 7 | CODE FROM THE ABOVE LINK WAS REWRITTEN USING PYTORCH SYMBOLIC 8 | """ 9 | 10 | from typing import Dict 11 | 12 | import torch 13 | import torch.nn as nn 14 | import torch.nn.functional as F 15 | import torch.optim as optim 16 | 17 | from pytorch_symbolic import Input, SymbolicModel, add_to_graph 18 | 19 | 20 | def run() -> None: 21 | """ 22 | Example: An LSTM for Part-of-Speech Tagging 23 | 24 | Prepare data: 25 | """ 26 | 27 | def prepare_sequence(seq, to_ix): 28 | idxs = [to_ix[w] for w in seq] 29 | return torch.tensor(idxs, dtype=torch.long) 30 | 31 | training_data = [ 32 | # Tags are: DET - determiner; NN - noun; V - verb 33 | # For example, the word "The" is a determiner 34 | ("The dog ate the apple".split(), ["DET", "NN", "V", "DET", "NN"]), 35 | ("Everybody read that book".split(), ["NN", "V", "DET", "NN"]), 36 | ] 37 | word_to_ix: Dict[str, int] = {} 38 | # For each words-list (sentence) and tags-list in each tuple of training_data 39 | for sent, _tags in training_data: 40 | for word in sent: 41 | if word not in word_to_ix: # word has not been assigned an index yet 42 | word_to_ix[word] = len(word_to_ix) # Assign each word with a unique index 43 | print(word_to_ix) 44 | tag_to_ix = {"DET": 0, "NN": 1, "V": 2} # Assign each tag with a unique index 45 | 46 | # These will usually be more like 32 or 64 dimensional. 47 | # We will keep them small, so we can see how the weights change as we train. 48 | EMBEDDING_DIM = 6 49 | HIDDEN_DIM = 6 50 | 51 | """ 52 | Create the model: 53 | """ 54 | 55 | class LSTMTagger(nn.Module): 56 | def __init__(self, embedding_dim, hidden_dim, vocab_size, tagset_size): 57 | super(LSTMTagger, self).__init__() 58 | self.hidden_dim = hidden_dim 59 | 60 | self.word_embeddings = nn.Embedding(vocab_size, embedding_dim) 61 | 62 | # The LSTM takes word embeddings as inputs, and outputs hidden states 63 | # with dimensionality hidden_dim. 64 | self.lstm = nn.LSTM(embedding_dim, hidden_dim) 65 | 66 | # The linear layer that maps from hidden state space to tag space 67 | self.hidden2tag = nn.Linear(hidden_dim, tagset_size) 68 | 69 | def forward(self, sentence): 70 | embeds = self.word_embeddings(sentence) 71 | lstm_out, _ = self.lstm(embeds.view(len(sentence), 1, -1)) 72 | tag_space = self.hidden2tag(lstm_out.view(len(sentence), -1)) 73 | tag_scores = F.log_softmax(tag_space, dim=1) 74 | return tag_scores 75 | 76 | def create_lstm_tagger(embedding_dim, hidden_dim, vocab_size, tagset_size): 77 | """Equivalent of LSTMTagger.""" 78 | sentence = Input(batch_shape=(30,), dtype=torch.int) 79 | embeds = nn.Embedding(vocab_size, embedding_dim)(sentence) 80 | lstm_out, _ = nn.LSTM(embedding_dim, hidden_dim)(embeds.unsqueeze(1)) 81 | tag_space = nn.Linear(hidden_dim, tagset_size)(lstm_out.squeeze()) 82 | tag_scores = add_to_graph(F.log_softmax, tag_space, dim=1) 83 | model = SymbolicModel(sentence, tag_scores) 84 | return model 85 | 86 | model = create_lstm_tagger(EMBEDDING_DIM, HIDDEN_DIM, len(word_to_ix), len(tag_to_ix)) 87 | loss_function = nn.NLLLoss() 88 | optimizer = optim.SGD(model.parameters(), lr=0.1) 89 | 90 | # See what the scores are before training 91 | # Note that element i,j of the output is the score for tag j for word i. 92 | # Here we don't need to train, so the code is wrapped in torch.no_grad() 93 | with torch.no_grad(): 94 | inputs = prepare_sequence(training_data[0][0], word_to_ix) 95 | tag_scores = model(inputs) 96 | print(tag_scores.argmax(1).tolist()) 97 | 98 | for _epoch in range(300): # again, normally you would NOT do 300 epochs, it is toy data 99 | for sentence, tags in training_data: 100 | # Step 1. Remember that Pytorch accumulates gradients. 101 | # We need to clear them out before each instance 102 | model.zero_grad() 103 | 104 | # Step 2. Get our inputs ready for the network, that is, turn them into 105 | # Tensors of word indices. 106 | sentence_in = prepare_sequence(sentence, word_to_ix) 107 | targets = prepare_sequence(tags, tag_to_ix) 108 | 109 | # Step 3. Run our forward pass. 110 | tag_scores = model(sentence_in) 111 | 112 | # Step 4. Compute the loss, gradients, and update the parameters by 113 | # calling optimizer.step() 114 | loss = loss_function(tag_scores, targets) 115 | loss.backward() 116 | optimizer.step() 117 | 118 | # See what the scores are after training 119 | with torch.no_grad(): 120 | inputs = prepare_sequence(training_data[0][0], word_to_ix) 121 | tag_scores = model(inputs) 122 | 123 | # The sentence is "the dog ate the apple". i,j corresponds to score for tag j 124 | # for word i. The predicted tag is the maximum scoring tag. 125 | # Here, we can see the predicted sequence below is 0 1 2 0 1 126 | # since 0 is index of the maximum value of row 1, 127 | # 1 is the index of maximum value of row 2, etc. 128 | # Which is DET NOUN VERB DET NOUN, the correct sequence! 129 | print("PREDICT", tag_scores.argmax(1).tolist()) 130 | 131 | print("CORRECT", [0, 1, 2, 0, 1]) 132 | -------------------------------------------------------------------------------- /tests/test_creating_models.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2022 Szymon Mikler 2 | 3 | import torch 4 | from torch import nn 5 | 6 | from pytorch_symbolic import Input, SymbolicModel, model_tools, optimize_module_calls, useful_layers 7 | 8 | 9 | def create_vanilla_pyt(seed): 10 | torch.manual_seed(seed) 11 | 12 | class Model(nn.Module): 13 | def __init__(self): 14 | super().__init__() 15 | 16 | self.layers = [] 17 | for _ in range(5): 18 | self.layers.append(nn.Linear(10, 10)) 19 | 20 | for idx, layer in enumerate(self.layers): 21 | self.add_module(name=str(idx), module=layer) 22 | 23 | def forward(self, x): 24 | for layer in self.layers: 25 | x = layer(x) 26 | return x 27 | 28 | return Model() 29 | 30 | 31 | def create_api_v1(seed): 32 | x = inputs = Input(shape=(10,)) 33 | torch.manual_seed(seed) 34 | 35 | for _ in range(5): 36 | x = x(nn.Linear(10, 10)) 37 | return SymbolicModel(inputs, x) 38 | 39 | 40 | def create_api_v2(seed): 41 | x = inputs = Input(batch_shape=(3, 10)) 42 | torch.manual_seed(seed) 43 | 44 | for _ in range(5): 45 | x = nn.Linear(10, 10)(x) 46 | return SymbolicModel(inputs, x) 47 | 48 | 49 | def test_equal_outputs(): 50 | for seed in range(10): 51 | x = torch.rand(10, 10) 52 | outs0 = create_vanilla_pyt(seed)(x) 53 | outs1 = create_api_v1(seed)(x) 54 | outs2 = create_api_v2(seed)(x) 55 | assert torch.equal(outs1, outs2) 56 | assert torch.equal(outs1, outs0) 57 | 58 | 59 | def test_equal_outputs_with_call_optimization(): 60 | for seed in range(10): 61 | x = torch.rand(10, 10) 62 | model0 = create_vanilla_pyt(seed) 63 | model1 = create_api_v1(seed) 64 | model2 = create_api_v2(seed) 65 | optimize_module_calls() 66 | outs0 = model0(x) 67 | outs1 = model1(x) 68 | outs2 = model2(x) 69 | assert torch.equal(outs1, outs2) 70 | assert torch.equal(outs1, outs0) 71 | 72 | 73 | def test_equal_parameters(): 74 | for seed in range(10): 75 | model0 = create_vanilla_pyt(seed) 76 | model1 = create_api_v1(seed) 77 | model2 = create_api_v2(seed) 78 | assert model_tools.model_similar(model1, model2) 79 | assert model_tools.model_similar(model1, model0) 80 | assert model_tools.models_have_corresponding_parameters(model1, model2) 81 | assert model_tools.models_have_corresponding_parameters(model1, model0) 82 | 83 | 84 | def create_vanilla_pyt_multi_in_out(seed): 85 | torch.manual_seed(seed) 86 | 87 | class Model(nn.Module): 88 | def __init__(self): 89 | super().__init__() 90 | self.l11 = nn.Linear(10, 10) 91 | self.l12 = nn.Linear(10, 10) 92 | self.l21 = nn.Linear(10, 10) 93 | self.l22 = nn.Linear(10, 10) 94 | self.l00 = nn.Linear(20, 10) 95 | self.l13 = nn.Linear(10, 10) 96 | self.l23 = nn.Linear(10, 10) 97 | 98 | def forward(self, x1, x2): 99 | x1 = self.l11(x1) 100 | x1 = self.l12(x1) 101 | 102 | x2 = self.l21(x2) 103 | x2 = self.l22(x2) 104 | x_cat = torch.concat((x1, x2), dim=1) 105 | x_cat = self.l00(x_cat) 106 | 107 | x_out1 = self.l13(x_cat) 108 | x_out2 = self.l23(x_cat) 109 | return x_out1, x_out2 110 | 111 | return Model() 112 | 113 | 114 | def create_api_multi_in_out(seed): 115 | inputs1 = x1 = Input(shape=(10,)) 116 | inputs2 = x2 = Input(shape=(10,)) 117 | torch.manual_seed(seed) 118 | 119 | for _ in range(2): 120 | x1 = nn.Linear(10, 10)(x1) 121 | 122 | for _ in range(2): 123 | x2 = nn.Linear(10, 10)(x2) 124 | 125 | x_cat = useful_layers.ConcatLayer(dim=1)(x1, x2) 126 | x_cat = nn.Linear(20, 10)(x_cat) 127 | 128 | x_out1 = nn.Linear(10, 10)(x_cat) 129 | x_out2 = nn.Linear(10, 10)(x_cat) 130 | return SymbolicModel((inputs1, inputs2), (x_out1, x_out2)) 131 | 132 | 133 | def test_equal_parameters_multi_in_out(): 134 | for seed in range(10): 135 | model1 = create_vanilla_pyt_multi_in_out(seed) 136 | model2 = create_api_multi_in_out(seed) 137 | assert model_tools.model_similar(model1, model2) 138 | assert model_tools.models_have_corresponding_parameters(model1, model2) 139 | 140 | 141 | def test_equal_outputs_multi_in_out(): 142 | for seed in range(10): 143 | x1 = torch.rand(10, 10) 144 | x2 = torch.rand(10, 10) 145 | 146 | model1 = create_vanilla_pyt_multi_in_out(seed) 147 | o11, o12 = model1(x1, x2) 148 | 149 | model2 = create_api_multi_in_out(seed) 150 | o21, o22 = model2(x1, x2) 151 | 152 | assert torch.equal(o11, o21) 153 | assert torch.equal(o12, o22) 154 | 155 | 156 | def test_detached_equal_parameters_multi_in_out(): 157 | for seed in range(10): 158 | model1 = create_vanilla_pyt_multi_in_out(seed) 159 | model2 = create_api_multi_in_out(seed).detach_from_graph() 160 | assert model_tools.model_similar(model1, model2) 161 | assert model_tools.models_have_corresponding_parameters(model1, model2) 162 | 163 | 164 | def test_detached_equal_outputs_multi_in_out(): 165 | for seed in range(10): 166 | x1 = torch.rand(10, 10) 167 | x2 = torch.rand(10, 10) 168 | 169 | model1 = create_vanilla_pyt_multi_in_out(seed) 170 | o11, o12 = model1(x1, x2) 171 | 172 | model2 = create_api_multi_in_out(seed).detach_from_graph() 173 | o21, o22 = model2(x1, x2) 174 | 175 | assert torch.equal(o11, o21) 176 | assert torch.equal(o12, o22) 177 | 178 | 179 | def test_creating_named_layers(): 180 | inputs = Input((10,)) 181 | layer = nn.Linear(10, 10) 182 | 183 | y1 = inputs(layer, custom_name="TEST_LAYER_777") 184 | y2 = layer(inputs, custom_name="TEST_LAYER_888") 185 | y = y1 + y2 186 | 187 | model = SymbolicModel(inputs, y) 188 | assert "TEST_LAYER_777" in str(model) 189 | assert "TEST_LAYER_888" in str(model) 190 | -------------------------------------------------------------------------------- /docs/benchmarks.md: -------------------------------------------------------------------------------- 1 | # Benchmarks 2 | 3 | Symbolic API simplifies and speeds up the prototyping 4 | and the development process. 5 | But does it sacrifice performance of the model itself? 6 | One of the most important principles in building this library was to 7 | avoid this. 8 | It was made with performance in mind. 9 | Standard model definition: a class inheriting form `torch.nn.Module` 10 | is a baseline for us. 11 | Symbolic API aims to create models just as fast in most scenarios. 12 | 13 | > **Note:** some features are designed for the convenience and might slow down 14 | > the models. A good example is `add_to_graph` function. 15 | > To get the very best performance, 16 | > it is crucial to use custom functions and operators 17 | > inside `forward` method of a `torch.nn.Module`. 18 | > Although even if we don't, the slowdowns should not exceed a few milliseconds. 19 | 20 | ## Tweaks 21 | 22 | When using Symbolic Model for performance critical use-case consider 23 | calling `optimize_module_calls` after all models are created. 24 | 25 | ```python 26 | from torch import nn 27 | from pytorch_symbolic import Input, SymbolicModel, optimize_module_calls 28 | 29 | x = inputs = Input(shape=(3, 32, 32)) 30 | x = nn.Identity()(x) 31 | model = SymbolicModel(inputs, x) 32 | 33 | optimize_module_calls() 34 | 35 | # Here goes: 36 | # * Model training 37 | # * Model inference 38 | # * Model benchmarking 39 | ``` 40 | 41 | Not using it might give you a small slowdown in CPU limited workflows, 42 | but should not affect large models on GPU. 43 | We execute it in our benchmarks. 44 | 45 | ## Deep linear model 46 | 47 | Pytorch Symbolic won't change the kernel runtime. 48 | The only place where it _might_ introduce a slowdown, is before the kernel launches. 49 | To see if it does, we will maximize the number of kernel calls. 50 | 51 | We will look at a very thin and deep model with linear layers only. 52 | Each layer will have only 4 features. 53 | If there's _any_ overhead induced by Symbolic Model, it should be visible here. 54 | In larger models, the overhead could be hidden by the kernel computation. 55 | 56 | #### Data 57 | 58 | Data is randomly generated: 59 | 60 | ```python 61 | import torch 62 | 63 | data = torch.rand(size=(4, 4)) 64 | ``` 65 | 66 | #### Model definition 67 | 68 | ```python 69 | from torch import nn 70 | from pytorch_symbolic import Input, SymbolicModel 71 | 72 | n_layers = 250 # other used values: 500, 750, 1000 73 | 74 | x = inputs = Input(shape=(4,)) 75 | for _ in range(n_layers): 76 | x = nn.Linear(4, 4)(x) 77 | model = SymbolicModel(inputs, x) 78 | ``` 79 | 80 | ### Inference (CPU) 81 | 82 | For such small subsequent matrix multiplications, 83 | it can be faster to launch the model on the CPU. 84 | 85 | ![images/many_linear_layers.png](images/many_linear_layers.png) 86 | > Percentile intervals [25, 75] are visible. Sequential model is visibly 87 | > slower than the others. This can be explained by the operations 88 | > introduced by the iterator added in `torch.nn.Sequential`. 89 | > Also, Sequential model seems to be slowing down more as the number of layers increases. 90 | > The other two models seem to be equally fast! 91 | 92 | ## Toy ResNet 93 | 94 | This model is presented in [Advanced Topics](advanced_topics.md), 95 | it was also used as an example in [Keras documentation](https://keras.io/guides/functional_api/). 96 | It is a shallower and thinner version of commonly used ResNet network. 97 | 98 | #### Data 99 | 100 | Data is randomly generated: 101 | 102 | ```python 103 | import torch 104 | 105 | data = torch.rand(size=(4, 3, 16, 16)) # Resolution from 16x16 to 64x64 106 | ``` 107 | 108 | #### Model definition 109 | 110 | Definition can be found in [Advanced Topics](advanced_topics.md). 111 | 112 | ### Inference (GPU) 113 | 114 | ![images/toy_resnet.png](images/toy_resnet.png) 115 | > CUDA Graphs have a huge advantage here due to the small batch size and image size. 116 | > For non CUDA Graphed models GPU is executing kernels much faster than CPU 117 | > is scheduling the work. 118 | > This is why we don't see any slowdown when the image resolution increases. 119 | > Nevertheless, Symbolic Model is slightly faster than the Vanilla model. 120 | > This is due to some implementation details. 121 | > For example, it is quite slow to access a layer by `__getattr__` in forward function. 122 | > In Symbolic Model there is no need to do this. 123 | 124 | ## How is Symbolic Model optimized? 125 | 126 | Symbolic models reside on underlying graph structures. 127 | Each Symbolic Tensor is a node and each layer is an edge that connects two nodes. 128 | Initially, the forward pass was implemented lazily: 129 | by executing `forward` in a layer only when 130 | its output was needed by a child node. 131 | But such back-and-forth between parents and children created an unnecessary overhead. 132 | To avoid this, we precompute the exact order in which the layers needs to be called 133 | and use it during `forward`. 134 | 135 | Even when we know the order of the layers, there's one more trick. 136 | Accessing structures has a significant overhead in Python, so we try to avoid it. 137 | When the Symbolic Model is created we dynamically generate code for its `forward` function. 138 | Thanks to this, Symbolic Model executes exactly the same code it would if you 139 | were to write it as a class. 140 | 141 | You can even see the generated code yourself: 142 | 143 | ```python 144 | config.CODEGEN_MIN_LOOP_LENGTH = 10 145 | ... 146 | print(model._generated_forward_source) 147 | ``` 148 | 149 | ```python 150 | def _generated_forward(self, i00, i01): 151 | l = self._execution_order_layers 152 | h09 = i01 153 | for layer in l[0:10]: 154 | h09 = layer(h09) 155 | h19 = i00 156 | for layer in l[10:20]: 157 | h19 = layer(h19) 158 | h20 = l[20](h19, h09) 159 | h21 = l[21](h20) 160 | h22 = l[22](h21) 161 | h32 = h22 162 | for layer in l[23:33]: 163 | h32 = layer(h32) 164 | h42 = h22 165 | for layer in l[33:43]: 166 | h42 = layer(h42) 167 | o00 = l[43](h42, h32) 168 | return o00 169 | ``` 170 | 171 | ## Hardware 172 | 173 | Unless stated otherwise, experiments were run on a following PC: 174 | 175 | ``` 176 | CPU: Intel i7-12700KF 177 | GPU: NVIDIA RTX 3080 10GB 178 | ``` 179 | -------------------------------------------------------------------------------- /tests/test_basic_ops.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2022 Szymon Mikler 2 | 3 | import torch 4 | from torch import nn 5 | 6 | from pytorch_symbolic import CustomInput, Input, SymbolicModel, add_to_graph, optimize_module_calls 7 | 8 | 9 | def test_all(): 10 | x1 = Input((5, 5)) 11 | x2 = Input((5, 5)) 12 | const = torch.rand(5, 5) 13 | 14 | results1 = [ 15 | abs(x1), 16 | abs(x2), 17 | x1 + x2, 18 | x1 - x2 / 2, 19 | x1 * x2, 20 | x1 / x2, 21 | x1 % x2, 22 | x1 @ x2, 23 | x1**x2, 24 | ] 25 | model1 = SymbolicModel(inputs=(x1, x2), outputs=sum(results1)) 26 | 27 | results2 = [ 28 | x1 + const, 29 | x1 - const / 2, 30 | x1 * const, 31 | x1 / const, 32 | x1 % const, 33 | x1 @ const, 34 | x1**const, 35 | ] 36 | model2 = SymbolicModel(inputs=(x1, x2), outputs=sum(results2)) 37 | 38 | results3 = [ 39 | const + x2, 40 | const - x2 / 2, 41 | const * x2, 42 | const / x2, 43 | const % x2, 44 | const @ x2, 45 | const**x2, 46 | ] 47 | model3 = SymbolicModel(inputs=(x1, x2), outputs=sum(results3)) 48 | 49 | for _ in range(10): 50 | x1 = torch.rand(5, 5) 51 | x2 = torch.rand(5, 5) 52 | out1 = model1(x1, x2) 53 | res1 = abs(x1) + abs(x2) + x1 + x2 + x1 - x2 / 2 + x1 * x2 + x1 / x2 + x1 % x2 + x1 @ x2 + x1**x2 54 | assert torch.allclose(out1, res1) 55 | 56 | out2 = model2(x1, x2) 57 | res2 = x1 + const + x1 - const / 2 + x1 * const + x1 / const + x1 % const + x1 @ const + x1**const 58 | assert torch.allclose(out2, res2) 59 | 60 | out3 = model3(x1, x2) 61 | res3 = const + x2 + const - x2 / 2 + const * x2 + const / x2 + const % x2 + const @ x2 + const**x2 62 | assert torch.allclose(out3, res3) 63 | 64 | 65 | def test_all_detached(): 66 | x1 = Input((5, 5)) 67 | x2 = Input((5, 5)) 68 | const = torch.rand(5, 5) 69 | 70 | results1 = [ 71 | abs(x1), 72 | abs(x2), 73 | x1 + x2, 74 | x1 - x2 / 2, 75 | x1 * x2, 76 | x1 / x2, 77 | x1 % x2, 78 | x1 @ x2, 79 | x1**x2, 80 | ] 81 | model1 = SymbolicModel(inputs=(x1, x2), outputs=sum(results1)).detach_from_graph() 82 | 83 | results2 = [ 84 | x1 + const, 85 | x1 - const / 2, 86 | x1 * const, 87 | x1 / const, 88 | x1 % const, 89 | x1 @ const, 90 | x1**const, 91 | ] 92 | model2 = SymbolicModel(inputs=(x1, x2), outputs=sum(results2)).detach_from_graph() 93 | 94 | results3 = [ 95 | const + x2, 96 | const - x2 / 2, 97 | const * x2, 98 | const / x2, 99 | const % x2, 100 | const @ x2, 101 | const**x2, 102 | ] 103 | model3 = SymbolicModel(inputs=(x1, x2), outputs=sum(results3)).detach_from_graph() 104 | 105 | for _ in range(10): 106 | x1 = torch.rand(5, 5) 107 | x2 = torch.rand(5, 5) 108 | out1 = model1(x1, x2) 109 | res1 = abs(x1) + abs(x2) + x1 + x2 + x1 - x2 / 2 + x1 * x2 + x1 / x2 + x1 % x2 + x1 @ x2 + x1**x2 110 | assert torch.allclose(out1, res1) 111 | 112 | out2 = model2(x1, x2) 113 | res2 = x1 + const + x1 - const / 2 + x1 * const + x1 / const + x1 % const + x1 @ const + x1**const 114 | assert torch.allclose(out2, res2) 115 | 116 | out3 = model3(x1, x2) 117 | res3 = const + x2 + const - x2 / 2 + const * x2 + const / x2 + const % x2 + const @ x2 + const**x2 118 | assert torch.allclose(out3, res3) 119 | 120 | 121 | def test_all_optimized(): 122 | x1 = Input((5, 5)) 123 | x2 = Input((5, 5)) 124 | const = torch.rand(5, 5) 125 | 126 | results1 = [ 127 | abs(x1), 128 | abs(x2), 129 | x1 + x2, 130 | x1 - x2 / 2, 131 | x1 * x2, 132 | x1 / x2, 133 | x1 % x2, 134 | x1 @ x2, 135 | x1**x2, 136 | ] 137 | model1 = SymbolicModel(inputs=(x1, x2), outputs=sum(results1)) 138 | 139 | results2 = [ 140 | x1 + const, 141 | x1 - const / 2, 142 | x1 * const, 143 | x1 / const, 144 | x1 % const, 145 | x1 @ const, 146 | x1**const, 147 | ] 148 | model2 = SymbolicModel(inputs=(x1, x2), outputs=sum(results2)) 149 | 150 | results3 = [ 151 | const + x2, 152 | const - x2 / 2, 153 | const * x2, 154 | const / x2, 155 | const % x2, 156 | const @ x2, 157 | const**x2, 158 | ] 159 | model3 = SymbolicModel(inputs=(x1, x2), outputs=sum(results3)) 160 | 161 | optimize_module_calls() 162 | 163 | for _ in range(10): 164 | x1 = torch.rand(5, 5) 165 | x2 = torch.rand(5, 5) 166 | out1 = model1(x1, x2) 167 | res1 = abs(x1) + abs(x2) + x1 + x2 + x1 - x2 / 2 + x1 * x2 + x1 / x2 + x1 % x2 + x1 @ x2 + x1**x2 168 | assert torch.allclose(out1, res1) 169 | 170 | out2 = model2(x1, x2) 171 | res2 = x1 + const + x1 - const / 2 + x1 * const + x1 / const + x1 % const + x1 @ const + x1**const 172 | assert torch.allclose(out2, res2) 173 | 174 | out3 = model3(x1, x2) 175 | res3 = const + x2 + const - x2 / 2 + const * x2 + const / x2 + const % x2 + const @ x2 + const**x2 176 | assert torch.allclose(out3, res3) 177 | 178 | 179 | def test_mmul_on_list_torch(): 180 | A, B, C = 5, 5, 5 181 | x = CustomInput(data=[[torch.rand(1) for _ in range(B)] for _ in range(A)]) 182 | y = CustomInput(data=[[torch.rand(1) for _ in range(C)] for _ in range(B)]) 183 | 184 | r = None 185 | for i in range(len(x[0])): 186 | for j in range(len(y)): 187 | 188 | class Mmul(nn.Module): 189 | def __init__(self, i, j): 190 | super().__init__() 191 | self.i = i 192 | self.j = j 193 | 194 | def forward(self, x, y): 195 | r = torch.zeros(len(x), len(y[0])) 196 | assert len(x[0]) == len(y), "Wrong shapes!" 197 | r[self.i][self.j] = sum([x[self.i][p] * y[p][self.j] for p in range(len(y))]) 198 | return r 199 | 200 | mmul = Mmul(i, j) 201 | 202 | if r is None: 203 | r = mmul(x, y) 204 | else: 205 | r = r + mmul(x, y) 206 | model = SymbolicModel(inputs=(x, y), outputs=(r,)) 207 | 208 | X = torch.rand(A, B) 209 | Y = torch.rand(B, C) 210 | 211 | R = model(X.tolist(), Y.tolist()) 212 | assert torch.allclose(R, X @ Y) 213 | 214 | 215 | def test_mmul_on_list_numpy(): 216 | import numpy as np 217 | 218 | A, B, C = 5, 5, 5 219 | x = CustomInput(data=[[np.random.rand() for _ in range(B)] for _ in range(A)]) 220 | y = CustomInput(data=[[np.random.rand() for _ in range(C)] for _ in range(B)]) 221 | 222 | rs = [] 223 | for i in range(len(x[0])): 224 | for j in range(len(y)): 225 | 226 | class Mmul(nn.Module): 227 | def __init__(self, i, j): 228 | super().__init__() 229 | self.i = i 230 | self.j = j 231 | 232 | def forward(self, x, y): 233 | r = np.zeros([len(x), len(y[0])]) 234 | assert len(x[0]) == len(y), "Wrong shapes!" 235 | r[self.i][self.j] = sum([x[self.i][p] * y[p][self.j] for p in range(len(y))]) 236 | return r 237 | 238 | mmul = Mmul(i, j) 239 | r = mmul(x, y) 240 | rs.append(r) 241 | 242 | output = add_to_graph(lambda *args: sum(args), *rs) 243 | model = SymbolicModel(inputs=(x, y), outputs=output) 244 | 245 | X = np.random.rand(A, B) 246 | Y = np.random.rand(B, C) 247 | 248 | R = model(X.tolist(), Y.tolist()) 249 | assert np.allclose(R, X @ Y) 250 | 251 | 252 | def test_anypow_layer(): 253 | tensor = Input(shape=(10, 20)) 254 | power = CustomInput(data=1.5) 255 | 256 | class AnyPow(nn.Module): 257 | def forward(self, tensor, power): 258 | return tensor**power 259 | 260 | output = AnyPow()(tensor, power) 261 | 262 | model = SymbolicModel(inputs=(tensor, power), outputs=output) 263 | 264 | for x in range(10): 265 | for y in [0.5, 1.0, 1.5, 2.0]: 266 | assert model(x, y) == x**y 267 | 268 | x = torch.rand(10, 20, 30) 269 | for y in [0.5, 1.0, 1.5, 2.0]: 270 | assert torch.allclose(model(x, y), x**y) 271 | assert torch.allclose(model.detach_from_graph()(x, y), x**y) 272 | 273 | x = torch.rand(10, 20, 30) 274 | for y in [0.5, 1.0, 1.5, 2.0]: 275 | assert torch.allclose(model(y, x), y**x) 276 | assert torch.allclose(model.detach_from_graph()(y, x), y**x) 277 | 278 | 279 | def test_indexing_symbolic_data(): 280 | inputs = CustomInput( 281 | { 282 | "x": torch.rand(10, 20), 283 | "y": torch.rand(10, 20), 284 | } 285 | ) 286 | key = CustomInput(data="x") 287 | value = inputs[key] 288 | outputs = nn.Identity()(value) 289 | 290 | model = SymbolicModel(inputs=(inputs, key), outputs=outputs) 291 | 292 | real_inputs = { 293 | "x": torch.rand(1, 2, 3), 294 | "y": torch.rand(1, 2, 3), 295 | "z": torch.rand(1, 2, 3), 296 | } 297 | 298 | for key in real_inputs: 299 | outs = model(real_inputs, key) 300 | assert torch.equal(outs, real_inputs[key]) 301 | -------------------------------------------------------------------------------- /docs/quick_start.md: -------------------------------------------------------------------------------- 1 | # Quick Start 2 | 3 | ## Features 4 | 5 | * Easy to use: 6 | * Familiar for users of [Keras Functional API](https://keras.io/guides/functional_api/) 7 | * Symbolic Tensors have API similar to `torch.Tensor` 8 | * Lots of flexibility and advanced control flow: 9 | * Reusing layers 10 | * Shared layers 11 | * Multiple inputs and outputs 12 | * Custom, user-defined modules and functions 13 | * No restrictions on module's signature 14 | 15 | ## Introduction to Symbolic Data 16 | 17 | To register a new layer, e.g. ``torch.nn.Linear``, in your model, you have two options: 18 | 19 | * `layer(symbolic_data)` (just like in [Keras Functional API](https://keras.io/guides/functional_api/)) 20 | * `symbolic_data(layer)` (like nowhere else) 21 | 22 | There are no differences between these options. 23 | Models produced with them will be identical. 24 | 25 | To create a linear classifier without hidden layers, you can write the following: 26 | 27 | ```python 28 | from torch import nn 29 | from pytorch_symbolic import Input, SymbolicModel 30 | 31 | inputs = Input(shape=(28 * 28,)) 32 | outputs = nn.Linear(in_features=inputs.features, out_features=10)(inputs) 33 | model = SymbolicModel(inputs, outputs) 34 | ``` 35 | 36 | Or equivalently: 37 | 38 | ```python 39 | from torch import nn 40 | from pytorch_symbolic import Input, SymbolicModel 41 | 42 | inputs = Input(shape=(28 * 28,)) 43 | outputs = inputs(nn.Linear(in_features=inputs.features, out_features=10)) 44 | model = SymbolicModel(inputs, outputs) 45 | ``` 46 | 47 | ### Computation graph 48 | 49 | Under the hood of Pytorch Symbolic, there lives a computation graph. 50 | 51 | Every Symbolic Tensor is a node in it. When interacting with Symbolic Tensor: 52 | 53 | * Think of it as a placeholder for your data 54 | * Use it like `torch.Tensor` (e.g. slicing and iterating over it) 55 | 56 | Let us play with Symbolic Tensor and see what we can do: 57 | 58 | ```python 59 | from pytorch_symbolic import Input, SymbolicModel 60 | 61 | inputs = Input((28 * 28,)) # create a symbolic batch of vectors 62 | print(inputs) 63 | outputs = inputs[0] # access first vector from the batch 64 | print(inputs) 65 | ``` 66 | 67 | ```stdout 68 | 69 | 70 | ``` 71 | 72 | At first, `inputs` was the only node in the graph. It had 0 parents and 0 children. 73 | 74 | But when we created `outputs`, a new node was registered as a child of `inputs`. 75 | 76 | Symbolic Tensors have useful attributes. 77 | Using them we can, for example, instantly obtain shapes of intermediate outputs, 78 | instead of deriving them by hand. Let's see it: 79 | 80 | ```py 81 | print(inputs.shape) 82 | print(inputs.features) 83 | ``` 84 | 85 | ```stdout 86 | torch.Size([1, 784]) 87 | 784 88 | ``` 89 | 90 | In some cases, calculating output shapes in your code adds unnecessary complexity, 91 | for example when padding and convolutions are involved. 92 | By calculating shapes automatically we write less code and write easier code. 93 | 94 | After creating the graph, 95 | you can replay all the defined operations by using a Symbolic Model. 96 | 97 | ```py 98 | model = SymbolicModel(inputs=inputs, outputs=inputs + 1) 99 | ``` 100 | 101 | This is a model that adds 1 to the input tensor. 102 | We defined it using Symbolic Tensors, but it will work on arbitrary tensors now! 103 | Using Pytorch Symbolic, you can perform much more complicated operations. 104 | 105 | ## Examples step by step 106 | 107 | ### Model for RGB images 108 | 109 | 1. Get your symbolic inputs. Specifying batch size is optional. There are a few ways to do it: 110 | * `inputs = Input(shape=(C, H, W))` 111 | * `inputs = Input(shape=(C, H, W), batch_size=B)` 112 | * `inputs = Input(batch_shape=(B, C, H, W))` 113 | 2. Register new modules in the graph. There are a few ways: 114 | * `outputs = inputs(layer)` 115 | * `outputs = layer(inputs)` 116 | * `outputs = add_to_graph(layer, inputs)` 117 | 3. Use standard operations on Symbolic Tensors: `+, -, *, **, /, //, %` and `abs`: 118 | * For example, `x = 2 + inputs` or `x = inputs % y` will work as expected 119 | 4. Similarly, use all the **non in-place** methods of `torch.Tensor`, e.g. `inputs.reshape` 120 | 5. To concatenate (similarly for stacking) Symbolic Tensors: 121 | * use `useful_layers.ConcatOpLayer(dim=1)(x, y)` 122 | * add custom function to the model: `add_to_graph(torch.concat, (x, y), dim=1)` 123 | 6. When working with Symbolic Tensors, use `.shape` property or one of the shortcuts: 124 | * `.C` and `.channels` equals `.shape[1]` for RGB data 125 | * `.H` equals `.shape[2]` for RGB data 126 | * `.W` equals `.shape[3]` for RGB data 127 | * `.HW` is (height, width) tuple for RGB data 128 | 7. Finally, create the model: `model = SymbolicModel(inputs, outputs)` 129 | 8. Use `model` as a normal PyTorch `nn.Module`. It's 100% compatible. 130 | When using the model, 131 | all the operations performed on Symbolic Data will be replayed on the real data. 132 | 133 | ### Sequential topology example 134 | 135 | In PyTorch, there's `torch.nn.Sequential` that allows creating simple sequential models. 136 | 137 | In Pytorch Symbolic, you can create them too: 138 | 139 | ```python 140 | from torch import nn 141 | from pytorch_symbolic import Input, SymbolicModel 142 | 143 | x = inputs = Input((3, 128, 128)) 144 | 145 | x = nn.Conv2d(in_channels=x.channels, out_channels=16, kernel_size=3)(x) 146 | x = nn.MaxPool2d(kernel_size=2)(x)(nn.ReLU()) 147 | 148 | x = nn.Conv2d(in_channels=x.channels, out_channels=32, kernel_size=3)(x) 149 | x = nn.MaxPool2d(kernel_size=2)(x)(nn.ReLU()) 150 | 151 | x = nn.Conv2d(in_channels=x.channels, out_channels=64, kernel_size=3)(x) 152 | x = nn.MaxPool2d(kernel_size=2)(x)(nn.ReLU()) 153 | 154 | x = nn.Conv2d(in_channels=x.channels, out_channels=64, kernel_size=3)(x) 155 | x = nn.MaxPool2d(kernel_size=2)(x)(nn.ReLU())(nn.Flatten()) 156 | 157 | outputs = nn.Linear(in_features=x.features, out_features=10)(x) 158 | model = SymbolicModel(inputs=inputs, outputs=outputs) 159 | model.summary() 160 | ``` 161 | 162 | ```stdout 163 | _____________________________________________________________ 164 | Layer Output shape Params Parent 165 | ============================================================= 166 | 1 Input_1 (None, 3, 128, 128) 0 167 | 2 Conv2d_1 (None, 16, 126, 126) 448 1 168 | 3 MaxPool2d_1 (None, 16, 63, 63) 0 2 169 | 4 ReLU_1 (None, 16, 63, 63) 0 3 170 | 5 Conv2d_2 (None, 32, 61, 61) 4640 4 171 | 6 MaxPool2d_2 (None, 32, 30, 30) 0 5 172 | 7 ReLU_2 (None, 32, 30, 30) 0 6 173 | 8 Conv2d_3 (None, 64, 28, 28) 18496 7 174 | 9 MaxPool2d_3 (None, 64, 14, 14) 0 8 175 | 10 ReLU_3 (None, 64, 14, 14) 0 9 176 | 11 Conv2d_4 (None, 64, 12, 12) 36928 10 177 | 12 MaxPool2d_4 (None, 64, 6, 6) 0 11 178 | 13 ReLU_4 (None, 64, 6, 6) 0 12 179 | 14 Flatten_1 (None, 2304) 0 13 180 | 15* Linear_1 (None, 10) 23050 14 181 | ============================================================= 182 | Total params: 83562 183 | Trainable params: 83562 184 | Non-trainable params: 0 185 | _____________________________________________________________ 186 | ``` 187 | 188 | ### Multiple inputs example 189 | 190 | There's nothing stopping you from using multiple input and output nodes. 191 | 192 | Just create multiple `SymbolicTensor` objects: 193 | 194 | ```python 195 | from torch import nn 196 | from pytorch_symbolic import Input, SymbolicModel 197 | 198 | task1_input = x = Input(shape=(3, 32, 32)) 199 | task2_input = y = Input(shape=(64,)) 200 | ``` 201 | 202 | Define the operations on them: 203 | 204 | ```py 205 | x = x(nn.Conv2d(x.channels, 16, 3)) 206 | x = x(nn.MaxPool2d(3))(nn.ReLU())(nn.Flatten()) 207 | head1_out = x(nn.Linear(x.features, 200)) 208 | 209 | y = y(nn.Linear(y.features, 512))(nn.ReLU()) 210 | y = y(nn.Linear(y.features, 512))(nn.ReLU()) 211 | head2_out = y(nn.Linear(y.features, 200)) 212 | 213 | x = head1_out + head2_out # elementwise sum 214 | x = x(nn.Linear(x.features, 400))(nn.ReLU()) 215 | task1_out = x(nn.Linear(x.features, 10)) 216 | task2_out = x(nn.Linear(x.features, 1)) 217 | ``` 218 | 219 | And create the model, passing tuples or lists as inputs and outputs: 220 | 221 | ```py 222 | model = SymbolicModel([task1_input, task2_input], [task1_out, task2_out]) 223 | model.summary() 224 | ``` 225 | 226 | ```stdout 227 | ____________________________________________________________ 228 | Layer Output shape Params Parent 229 | ============================================================ 230 | 1 Input_1 (None, 3, 32, 32) 0 231 | 2 Input_2 (None, 64) 0 232 | 3 Conv2d_1 (None, 16, 30, 30) 448 1 233 | 4 MaxPool2d_1 (None, 16, 10, 10) 0 3 234 | 5 ReLU_1 (None, 16, 10, 10) 0 4 235 | 6 Flatten_1 (None, 1600) 0 5 236 | 7 Linear_1 (None, 200) 320200 6 237 | 8 Linear_2 (None, 512) 33280 2 238 | 9 ReLU_2 (None, 512) 0 8 239 | 10 Linear_3 (None, 512) 262656 9 240 | 11 ReLU_3 (None, 512) 0 10 241 | 12 Linear_4 (None, 200) 102600 11 242 | 13 AddOpLayer_1 (None, 200) 0 7,12 243 | 14 Linear_5 (None, 400) 80400 13 244 | 15 ReLU_4 (None, 400) 0 14 245 | 16* Linear_6 (None, 10) 4010 15 246 | 17* Linear_7 (None, 1) 401 15 247 | ============================================================ 248 | Total params: 803995 249 | Trainable params: 803995 250 | Non-trainable params: 0 251 | ____________________________________________________________ 252 | ``` 253 | 254 | You can use this model in a following way: 255 | 256 | ```py 257 | import torch 258 | 259 | data1 = torch.rand(16, 3, 32, 32) 260 | data2 = torch.rand(16, 64) 261 | 262 | outs1, outs2 = model(data1, data2) 263 | ``` 264 | 265 | ## Special cases 266 | 267 | ### Use custom functions on Symbolic Tensors 268 | 269 | If you want to use a function from `torch.nn.functional` or basically 270 | any custom function in your model, you have a few options. 271 | The recommended way is to use it in a `torch.nn.Module`. 272 | 273 | As an example, let us say this is the function you want to use: 274 | 275 | ```python 276 | from torch import nn 277 | 278 | 279 | def custom_func(*args): 280 | print("Arguments: ", *args) 281 | return args 282 | ``` 283 | 284 | This function just prints whatever is passed to it and returns it. 285 | 286 | An equivalent `torch.nn.Module` can use this function directly: 287 | 288 | ```py 289 | class CustomModule(nn.Module): 290 | def forward(self, *args): 291 | return custom_func(*args) 292 | ``` 293 | 294 | We already wrapped a few useful functions for you, e.g. `torch.concat` and `torch.stack`. 295 | 296 | They are available in `pytorch_symbolic.useful_layers`. 297 | 298 | #### Alternative for custom functions 299 | 300 | If you really hate classes or are in a hurry, we got you covered. 301 | 302 | You can add almost any function to your Symbolic Model using `add_to_graph`: 303 | 304 | ```python 305 | import torch 306 | from pytorch_symbolic import Input 307 | from pytorch_symbolic.functions_utility import add_to_graph 308 | 309 | x1 = Input(shape=(3, 3)) 310 | x2 = Input(shape=(5, 3)) 311 | x = add_to_graph(torch.concat, (x1, x2), dim=1) 312 | x.shape # (1, 8, 3) 313 | ``` 314 | 315 | Attempting to use a function without `add_to_graph`, e.g. `x = torch.abs(x)` will most likely fail: 316 | 317 | ``` 318 | TypeError: abs(): argument 'input' (position 1) must be Tensor, not Input 319 | ``` 320 | 321 | So use `add_to_graph` instead, like `add_to_graph(torch.abs, x)`. 322 | 323 | > If your `torch.nn.Module` requires named arguments, you can use `add_to_graph` to register it. 324 | 325 | ### Modules with multiple inputs 326 | 327 | The best way to use modules with multiple inputs is to use `layer(*args)` notation. 328 | 329 | Just pass multiple arguments to the layer: 330 | 331 | ```python 332 | from pytorch_symbolic import Input, useful_layers 333 | 334 | x1 = Input(shape=(1, 2, 3)) 335 | x2 = Input(shape=(5, 2, 3)) 336 | x = useful_layers.ConcatLayer(dim=1)(x1, x2) 337 | x.shape # (1, 6, 2, 3) 338 | ``` 339 | 340 | Alternatively, using the other notation, do it like this `arg0(layer, *other_args)`: 341 | 342 | ```py 343 | x = x1(useful_layers.ConcatLayer(dim=1), x2) 344 | x.shape # (1, 6, 2, 3) 345 | ``` 346 | -------------------------------------------------------------------------------- /pytorch_symbolic/graph_algorithms.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2022 Szymon Mikler 2 | 3 | from __future__ import annotations 4 | 5 | from typing import TYPE_CHECKING, Any, Callable, Dict, Iterable, List, Set, Tuple 6 | 7 | from .symbolic_data import SymbolicData, SymbolicTensor, useful_layers 8 | 9 | if TYPE_CHECKING: 10 | from .symbolic_model import SymbolicModel 11 | 12 | from collections import defaultdict 13 | 14 | from torch import nn 15 | 16 | 17 | def check_for_missing_inputs( 18 | used_nodes: Set[SymbolicData], 19 | inputs: Tuple[SymbolicData, ...] | List[SymbolicData], 20 | ): 21 | """Check if there exist nodes which require input from outside of the graph. 22 | 23 | It is forbidden, as it doesn't make sense. 24 | 25 | Example of such violation:: 26 | 27 | x1 = Input(shape=(32,)) 28 | x2 = Input(shape=(32,)) 29 | x3 = x1 + x2 30 | model = SymbolicModel(inputs=x1, outputs=x3) 31 | 32 | Model cannot execute defined operations unless ``x2`` is given, because ``x3`` requires it. 33 | """ 34 | for node in used_nodes: 35 | if node in inputs: 36 | continue 37 | for parent in node.parents: 38 | assert parent in used_nodes, ( 39 | f"Node {node} depends on the output of a foreign node! Perhaps you set the wrong inputs?" 40 | ) 41 | 42 | 43 | def figure_out_nodes_between( 44 | inputs: Tuple[SymbolicData, ...] | List[SymbolicData] | None = None, 45 | outputs: Tuple[SymbolicData, ...] | List[SymbolicData] | None = None, 46 | ) -> Set[SymbolicData]: 47 | """Returns intersection of predecessors tree of outputs and succesors tree of inputs.""" 48 | 49 | all_nodes_above: List[SymbolicData] = [] 50 | if outputs is not None: 51 | for output_leaf in outputs: 52 | nodes_above = output_leaf._get_all_nodes_above() 53 | all_nodes_above.extend(nodes_above) 54 | 55 | all_nodes_below: List[SymbolicData] = [] 56 | if inputs is not None: 57 | for input_leaf in inputs: 58 | nodes_below = input_leaf._get_all_nodes_below() 59 | all_nodes_below.extend(nodes_below) 60 | 61 | # Get the intersection 62 | if inputs is None or outputs is None: 63 | used_nodes = set(all_nodes_above) | set(all_nodes_below) 64 | else: 65 | used_nodes = set(all_nodes_above) & set(all_nodes_below) 66 | if inputs is not None: 67 | check_for_missing_inputs(used_nodes, inputs) 68 | return used_nodes 69 | 70 | 71 | def default_node_text(sym: SymbolicData) -> str: 72 | if isinstance(sym, SymbolicTensor): 73 | return str(tuple(sym.v.shape)) 74 | # mypy hates the next line for some reason 75 | if isinstance(sym.v, Callable): # type: ignore 76 | return "callable" 77 | if hasattr(sym.v, "__len__"): 78 | return f"{type(sym.v).__name__}({len(sym.v)})" 79 | return type(sym.v).__name__ 80 | 81 | 82 | def default_edge_text(layer: nn.Module) -> str: 83 | if isinstance(layer, useful_layers.GetAttr): 84 | return "." + layer.name 85 | else: 86 | return layer._get_name() 87 | 88 | 89 | def _calc_sum_sq_distances(graph, positions_dict): 90 | sum_sq_distances = 0 91 | for node in positions_dict: 92 | related_nodes = list(graph.predecessors(node)) 93 | 94 | for n2 in related_nodes: 95 | x1, y1 = positions_dict[node] 96 | x2, y2 = positions_dict[n2] 97 | sum_sq_distances += (x1 - x2) ** 2 + (y1 - y2) ** 2 98 | return sum_sq_distances 99 | 100 | 101 | def _transpose_positions_dict(pos_dict): 102 | return {k: tuple(reversed(v)) for k, v in pos_dict.items()} 103 | 104 | 105 | def _fix_positions_in_multipartite_layout(graph, orig_positions_dict, align: str = "vertical"): 106 | import networkx as nx 107 | import numpy as np 108 | from scipy.optimize import linear_sum_assignment 109 | 110 | assert isinstance(graph, nx.DiGraph) 111 | 112 | if align == "vertical": 113 | orig_positions_dict = _transpose_positions_dict(orig_positions_dict) 114 | 115 | orig_sum = _calc_sum_sq_distances(graph, orig_positions_dict) 116 | 117 | positions = list(orig_positions_dict.values()) 118 | selected_positions: Dict[int, Tuple[float, float]] = {} 119 | 120 | layers = defaultdict(list) 121 | for pos in positions: 122 | layers[pos[1]].append(pos[0]) 123 | 124 | layer_to_nodes = defaultdict(list) 125 | for node, pos in orig_positions_dict.items(): 126 | layer_to_nodes[pos[1]].append(node) 127 | 128 | # Layers are just unique Y positions on the plot 129 | # Go through the highest layer to the lowest one 130 | for layer, nodes_in_layer in reversed(sorted(layer_to_nodes.items())): 131 | n = len(nodes_in_layer) 132 | distances = np.zeros(shape=(n, n)) 133 | 134 | # Extract available X positions in the layer 135 | avail_xs = layers[layer] 136 | 137 | for node_idx, node in enumerate(nodes_in_layer): 138 | related_nodes = list(graph.predecessors(node)) 139 | related_nodes += list(graph.successors(node)) 140 | 141 | for x_idx, x in enumerate(avail_xs): 142 | sum_distances = 0 143 | for related_node in related_nodes: 144 | if related_node in selected_positions: 145 | r_x, r_y = selected_positions[related_node] 146 | sum_distances += (x - r_x) ** 2 + (layer - r_y) ** 2 147 | else: 148 | # If node doesn't have a position yet 149 | # we use the original position but we weight it down 150 | r_x, r_y = orig_positions_dict[related_node] 151 | sum_distances += ((x - r_x) ** 2 + (layer - r_y) ** 2) / 2 152 | distances[node_idx, x_idx] = sum_distances 153 | 154 | # Minimize the sum of squared distances 155 | # This might work poorly when there are interconnections in the layer 156 | rows, cols = linear_sum_assignment(distances) 157 | for row, col in zip(rows, cols): 158 | selected_positions[nodes_in_layer[row]] = (avail_xs[col], layer) 159 | 160 | new_sum = _calc_sum_sq_distances(graph, selected_positions) 161 | 162 | if new_sum <= orig_sum: 163 | r = selected_positions 164 | else: 165 | r = orig_positions_dict 166 | 167 | if align == "vertical": 168 | r = _transpose_positions_dict(r) 169 | 170 | return r 171 | 172 | 173 | def variable_name_resolver(namespace): 174 | def resolver(x): 175 | names = [name for name in namespace if namespace[name] is x] 176 | return names[0] if names else "?" 177 | 178 | return resolver 179 | 180 | 181 | def draw_graph( 182 | *, 183 | model: SymbolicModel | None = None, 184 | inputs: Iterable[SymbolicData] | SymbolicData | None = None, 185 | outputs: Iterable[SymbolicData] | SymbolicData | None = None, 186 | node_text_func: Callable[[SymbolicData], str] | None = None, 187 | edge_text_func: Callable[[nn.Module], str] | None = None, 188 | node_text_namespace: Dict[str, Any] | None = None, 189 | rotate_graph: bool = False, 190 | rotate_labels: bool = False, 191 | show: bool = False, 192 | figsize=None, 193 | ): 194 | """Plot graph of the computations, nodes being placeholder variables and nn.Modules being edges. 195 | 196 | This is not suitable for large graphs or large neural networks. This is a simple tool that 197 | was designed to demonstrate that Pytorch Symbolic creates sensible graphs that are nice 198 | to visualize. 199 | 200 | Parameters 201 | ---------- 202 | model 203 | A SymbolicModel to be plotted. This or ``inputs`` must be provided. 204 | inputs 205 | Input in the graph of Symbolic computations. This or ``model`` or/and ``inputs`` must be provided. 206 | outputs 207 | Output in the graph of Symbolic computations. This or ``model`` or/and ``outputs`` must be provided. 208 | node_text_func 209 | A function that returns text that will be written on Nodes. 210 | edge_text_func 211 | A function that returns text that will be written on Edges. 212 | node_text_namespace 213 | If used with ``globals()``, it will try to show variable name on each node. 214 | Has an effect only if `node_text_func` is None. 215 | rotate_labels 216 | If True, text on edges will be rotated in the direction of the arrow. 217 | rotate_graph 218 | If True, the consecutive layers will be shown to the right, instead of downwards. 219 | show 220 | Call matplotlib.pyplot.show 221 | """ 222 | try: 223 | import matplotlib.patches 224 | import matplotlib.pyplot as plt 225 | import networkx as nx 226 | except ImportError as e: 227 | print("To plot graphs, you need to install networkx, matplotlib and scipy. Run `pip install networkx`.") 228 | raise e 229 | 230 | from .symbolic_model import SymbolicModel 231 | 232 | if node_text_func is None: 233 | if node_text_namespace is not None: 234 | node_text_func = variable_name_resolver(node_text_namespace) 235 | else: 236 | node_text_func = default_node_text 237 | 238 | if edge_text_func is None: 239 | edge_text_func = default_edge_text 240 | 241 | if model is not None: 242 | assert isinstance(model, SymbolicModel) 243 | inputs = model.inputs 244 | outputs = model.outputs 245 | 246 | elif inputs is not None or outputs is not None: 247 | if inputs is not None: 248 | if isinstance(inputs, SymbolicData): 249 | inputs = (inputs,) 250 | assert all(isinstance(x, SymbolicData) for x in inputs) 251 | inputs = tuple(inputs) 252 | if outputs is not None: 253 | if isinstance(outputs, SymbolicData): 254 | outputs = (outputs,) 255 | assert all(isinstance(x, SymbolicData) for x in outputs) 256 | outputs = tuple(outputs) 257 | 258 | else: 259 | raise KeyError("Provide either `model` or `inputs` or/and `outputs`!") 260 | 261 | used_nodes = figure_out_nodes_between(inputs, outputs) 262 | 263 | graph = nx.DiGraph() 264 | assert isinstance(inputs, Iterable) 265 | to_visit = list(inputs) 266 | visited = set(inputs) 267 | 268 | node_labels = {} 269 | node_colors = {} 270 | edge_labels = {} 271 | 272 | INPUT_COLOR = (1.0, 0.3, 0.3) 273 | OUTPUT_COLOR = (0.3, 1.0, 0.3) 274 | OTHER_COLOR = (0.0, 1.0, 1.0) 275 | 276 | while to_visit: 277 | node = to_visit.pop() 278 | graph.add_node(node, depth=node.depth) 279 | 280 | if inputs and node in inputs: 281 | node_colors[node] = INPUT_COLOR 282 | elif outputs and node in outputs: 283 | node_colors[node] = OUTPUT_COLOR 284 | else: 285 | node_colors[node] = OTHER_COLOR 286 | 287 | node_labels[node] = node_text_func(node) 288 | 289 | for parent in node.parents: 290 | if parent not in used_nodes: 291 | continue 292 | assert node.layer is not None 293 | graph.add_edge(parent, node, layer=node.layer) 294 | edge_labels[parent, node] = edge_text_func(node.layer) 295 | 296 | for child in node.children: 297 | if child in used_nodes and child not in visited: 298 | to_visit.append(child) 299 | visited.add(child) 300 | 301 | scale = 1 if rotate_graph else -1 302 | align = "vertical" if rotate_graph else "horizontal" 303 | pos = nx.multipartite_layout(graph, subset_key="depth", align=align, scale=scale) 304 | pos = _fix_positions_in_multipartite_layout(graph, pos, align=align) 305 | 306 | nx.draw_networkx( 307 | graph, 308 | pos, 309 | with_labels=True, 310 | labels=node_labels, 311 | node_size=1000, 312 | arrowsize=20, 313 | node_shape="o", 314 | node_color=[node_colors[n] for n in graph.nodes], 315 | font_weight="bold", 316 | edge_color="grey", 317 | ) 318 | 319 | nx.draw_networkx_edge_labels( 320 | graph, 321 | pos, 322 | edge_labels, 323 | label_pos=0.5, 324 | rotate=rotate_labels, 325 | clip_on=False, 326 | ) 327 | handles = [ 328 | matplotlib.patches.Patch(color=INPUT_COLOR, label="Input node"), 329 | matplotlib.patches.Patch(color=OUTPUT_COLOR, label="Output node"), 330 | matplotlib.patches.Patch(color=OTHER_COLOR, label="Hidden node"), 331 | ] 332 | plt.legend(handles=handles) 333 | 334 | fig: plt.Figure = plt.gcf() 335 | if figsize: 336 | fig.set_size_inches(*figsize) 337 | 338 | fig.tight_layout() 339 | return fig 340 | # if show: 341 | # fig.show() 342 | 343 | 344 | def sort_graph_and_check_DAG(nodes: Set[SymbolicData]) -> List[SymbolicData]: 345 | """Sort graph topologically. 346 | 347 | Wikipedia: 348 | In graph theory, a topological sort or topological ordering of a directed acyclic graph (DAG) is a 349 | linear ordering of its nodes in which each node comes before all nodes to which it has outbound edges. 350 | Every DAG has one or more topological sorts. 351 | """ 352 | 353 | children = {node: set((c for c in node.children if c in nodes)) for node in nodes} 354 | parents = {node: set((p for p in node.parents if p in nodes)) for node in nodes} 355 | grandparents = [node for node in nodes if len(parents[node]) == 0] # no incoming edges 356 | 357 | n_edges = sum(len(v) for v in children.values()) + sum(len(v) for v in parents.values()) 358 | 359 | topologically_sorted = [] 360 | 361 | while grandparents: 362 | node = grandparents.pop() 363 | topologically_sorted.append(node) 364 | 365 | for neighbor in tuple(children[node]): 366 | children[node].remove(neighbor) 367 | parents[neighbor].remove(node) 368 | 369 | if len(parents[neighbor]) == 0: 370 | grandparents.append(neighbor) 371 | 372 | n_edges = sum(len(v) for v in children.values()) + sum(len(v) for v in parents.values()) 373 | 374 | assert n_edges == 0, "Graph is not a DAG (directed acyclic graph)!" 375 | return topologically_sorted 376 | -------------------------------------------------------------------------------- /pytorch_symbolic/symbolic_model.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2022 Szymon Mikler 2 | 3 | from __future__ import annotations 4 | 5 | import copy 6 | import logging 7 | from types import MethodType 8 | from typing import Any, Dict, List, Set, Tuple 9 | 10 | import torch 11 | from torch import nn 12 | 13 | from . import code_generator, config 14 | from .graph_algorithms import figure_out_nodes_between, sort_graph_and_check_DAG 15 | from .model_tools import get_parameter_count 16 | from .symbolic_data import SymbolicData, SymbolicTensor 17 | 18 | 19 | class DetachedSymbolicModel(nn.Module): 20 | def __init__(self, names: List[str], layers: List[nn.Module], forward_src: str): 21 | """A tiny model detached from the SymbolicModel graph structure. 22 | 23 | It can live, even if the graph structure is removed! 24 | 25 | Parameters 26 | ---------- 27 | names 28 | Names for all the layers. Must be equal in length to ``layers`` 29 | layers 30 | Ordered list of layers to be executed 31 | forward_src 32 | String containing definition of a ``forward`` function. 33 | Must define ``def forward(self, ...)``. 34 | Layers are available under ``self._execution_order_layers`` in the function's body. 35 | """ 36 | assert len(names) == len(layers) 37 | 38 | super().__init__() 39 | self._execution_order_layers = [] 40 | for name, layer in zip(names, layers): 41 | try: 42 | layer = copy.deepcopy(layer) 43 | except Exception as e: 44 | logging.error(f"Deepcopy of {layer} failed!") 45 | raise e 46 | self.add_module(name, layer) 47 | self._execution_order_layers.append(layer) 48 | 49 | self._generated_forward_source = forward_src 50 | 51 | scope = {"self": self} 52 | exec(self._generated_forward_source, {}, scope) 53 | self.forward = MethodType(scope["forward"], self) 54 | 55 | 56 | class SymbolicModel(nn.Module): 57 | def __init__( 58 | self, 59 | inputs: Tuple[SymbolicData, ...] | List[SymbolicData] | SymbolicData, 60 | outputs: Tuple[SymbolicData, ...] | List[SymbolicData] | SymbolicData, 61 | enable_forward_codegen=None, 62 | ): 63 | """A PyTorch model that replays operations defined in the graph. 64 | 65 | All operations that were required to change ``inputs`` into ``outputs`` will be replayed 66 | in the same order, but on the real data provided as input to this model. 67 | 68 | Example:: 69 | 70 | input1 = Input((10,)) 71 | input2 = Input((10,)) 72 | x = input1 + input2 73 | x = nn.Linear(x.features, 1)(x) 74 | model = SymbolicModel((input1, input2), x) 75 | 76 | Parameters 77 | ---------- 78 | inputs 79 | A collection of SymbolicData that represent the input data used by the model. 80 | It is you who provide the specific data when the model is created. 81 | If you have mulitple inputs here, be prepared to pass multiple inputs during training/inference. 82 | outputs 83 | A collection of SymbolicTensors that will end the computations. 84 | These nodes return your final computation result. 85 | So if you have mulitple outputs, SymbolicModel will return a tuple of tensors. 86 | 87 | Attributes 88 | ---------- 89 | inputs : tuple 90 | Non-modifiable tuple of input nodes 91 | outputs : tuple 92 | Non-modifiable tuple of output nodes 93 | """ 94 | super().__init__() 95 | logging.debug("Creating a SymbolicModel...") 96 | 97 | if isinstance(inputs, SymbolicData): 98 | inputs = (inputs,) 99 | assert all(isinstance(x, SymbolicData) for x in inputs), "Only SymbolicData allowed in inputs!" 100 | self.inputs: Tuple[SymbolicData, ...] = tuple(inputs) 101 | 102 | if isinstance(outputs, SymbolicData): 103 | outputs = (outputs,) 104 | assert all(isinstance(x, SymbolicData) for x in outputs), "Only SymbolicData allowed in outputs!" 105 | self.outputs: Tuple[SymbolicData, ...] = tuple(outputs) 106 | 107 | # Initialize helper variables 108 | self._layer_type_counts: Dict[str, int] = {} 109 | self._node_to_layer_name: Dict[SymbolicData, str] = {} 110 | self._execution_order_nodes: List[SymbolicData] = [] 111 | self._execution_order_layers: List[nn.Module] = [] 112 | self._figure_out_execution_order() 113 | 114 | if enable_forward_codegen is None: 115 | enable_forward_codegen = config.CODEGEN_BY_DEFAULT 116 | self._enable_forward_codegen = enable_forward_codegen 117 | if self._enable_forward_codegen: 118 | self._replace_forward_with_codegen() 119 | 120 | self._clear_underlying_values() # clear up after tracing to save memory 121 | 122 | def forward(self, *inputs: torch.Tensor) -> Any: 123 | """This function is executed by __call__. Do not use this directly, use __call__ instead. 124 | 125 | Warning! 126 | 127 | This function will be overwritten by `_replace_forward_with_codegen` if `enable_forward_codegen` 128 | is True. If this happened and you want to see your source, print `self._generated_forward_source`. 129 | """ 130 | assert len(inputs) == len(self.inputs), "Number of inputs doesn't match!" 131 | for input_data, input_node in zip(inputs, self.inputs): 132 | input_node._launch_input(input_data) 133 | 134 | for node in self._execution_order_nodes: 135 | node._launch() 136 | 137 | if len(self.outputs) == 1: 138 | return self.outputs[0]._value 139 | else: 140 | return tuple(output_leaf._value for output_leaf in self.outputs) 141 | 142 | @property 143 | def input_shape(self): 144 | """Return shape of the input or in case of multiple inputs - a tuple of them.""" 145 | shapes = [node.shape if isinstance(node, SymbolicTensor) else None for node in self.inputs] 146 | return tuple(shapes) if len(shapes) > 1 else shapes[0] 147 | 148 | @property 149 | def output_shape(self): 150 | """Return shape of the output or in case of multiple outputs - a tuple of them.""" 151 | shapes = [node.shape if isinstance(node, SymbolicTensor) else None for node in self.outputs] 152 | return tuple(shapes) if len(shapes) > 1 else shapes[0] 153 | 154 | def add_output(self, node: SymbolicData): 155 | assert node not in self.inputs, "Node is an input of this SymbolicModel!" 156 | assert node in self._execution_order_nodes, "Node is out of reach for this SymbolicModel!" 157 | 158 | self.outputs = (*self.outputs, node) 159 | if self._enable_forward_codegen: 160 | self._replace_forward_with_codegen() 161 | 162 | def detach_from_graph(self) -> DetachedSymbolicModel: 163 | names = [self._node_to_layer_name[node] for node in self._execution_order_nodes] 164 | forward_src = code_generator.generate_forward_with_loops( 165 | self.inputs, 166 | self.outputs, 167 | execution_order=self._execution_order_nodes, 168 | nodes_in_subgraph=self._used_nodes(), 169 | min_loop_length=config.CODEGEN_MIN_LOOP_LENGTH, 170 | ) 171 | return DetachedSymbolicModel(names, self._execution_order_layers, forward_src) 172 | 173 | def summary(self): 174 | """Print Keras-like model summary.""" 175 | space_between_cols = 3 176 | 177 | data = [["", "Layer", "Output shape", "Params", "Parent"]] 178 | ncols = len(data[0]) 179 | separators = ["="] 180 | node_to_idx = {} 181 | for node in self.inputs: 182 | node_to_idx[node] = len(data) 183 | if isinstance(node, SymbolicTensor): 184 | shape = list(node.shape) 185 | if not node.batch_size_known: 186 | shape[0] = None 187 | shape = tuple(shape) 188 | else: 189 | shape = node._underlying_type_name 190 | data.append( 191 | [ 192 | f"{len(data)}" + ("*" if node in self.outputs else ""), 193 | f"Input_{len(data)}", 194 | str(shape), 195 | "0", 196 | "", 197 | ] 198 | ) 199 | separators.append(None) 200 | 201 | for node in self._execution_order_nodes: 202 | node_to_idx[node] = len(data) 203 | layer = node.layer 204 | if isinstance(node, SymbolicTensor): 205 | shape = list(node.shape) 206 | if not node.batch_size_known: 207 | shape[0] = None 208 | shape = tuple(shape) 209 | else: 210 | shape = node._underlying_type_name 211 | data.append( 212 | [ 213 | f"{len(data)}" + ("*" if node in self.outputs else ""), 214 | f"{self._node_to_layer_name[node]}", 215 | str(shape), 216 | str(get_parameter_count(layer)), 217 | ",".join(str(node_to_idx[parent]) for parent in node.parents), 218 | ] 219 | ) 220 | separators.append(None) 221 | separators[-1] = "=" 222 | 223 | maxcolwidth = [0 for _ in range(ncols)] 224 | for row in data: 225 | for idx, col in enumerate(row): 226 | if len(col) > maxcolwidth[idx]: 227 | maxcolwidth[idx] = len(col) 228 | 229 | print("_" * (sum(maxcolwidth) + ncols * space_between_cols)) 230 | for sep, row in zip(separators, data): 231 | for idx, col in enumerate(row): 232 | s = col.ljust(maxcolwidth[idx] + space_between_cols, " ") 233 | print(s, end="") 234 | print() 235 | if sep is not None: 236 | print(sep * (sum(maxcolwidth) + ncols * space_between_cols)) 237 | 238 | parameter_count = get_parameter_count(self) 239 | trainable_count = get_parameter_count(self, only_trainable=True) 240 | print(f"Total params: {parameter_count}") 241 | print(f"Trainable params: {trainable_count}") 242 | print(f"Non-trainable params: {parameter_count - trainable_count}") 243 | print("_" * (sum(maxcolwidth) + ncols * space_between_cols)) 244 | 245 | def _clear_underlying_values(self): 246 | """Clear values of the underlying nodes to save memory. 247 | 248 | Does not clear the values of input nodes, as they might be needed. 249 | """ 250 | for node in self._used_nodes().difference(self.inputs): 251 | node._clear_value() 252 | 253 | def _replace_forward_with_codegen(self): 254 | self._generated_forward_source = code_generator.generate_forward_with_loops( 255 | self.inputs, 256 | self.outputs, 257 | execution_order=self._execution_order_nodes, 258 | nodes_in_subgraph=self._used_nodes(), 259 | min_loop_length=config.CODEGEN_MIN_LOOP_LENGTH, 260 | ) 261 | scope = {"self": self} 262 | exec(self._generated_forward_source, {}, scope) 263 | self.forward = MethodType(scope["forward"], self) 264 | 265 | def _used_nodes(self) -> Set[SymbolicData]: 266 | """Return a set of all nodes used in this model.""" 267 | return figure_out_nodes_between(self.inputs, self.outputs) 268 | 269 | def _remove_repeated_execution(self, execution_order_nodes: List[SymbolicData]) -> List[SymbolicData]: 270 | """In case of multiple outputs, we need only one of the output node to launch the layer.""" 271 | nodes_without_repeated_execution = [] 272 | used_nodes = self._used_nodes() 273 | 274 | already_executed: Set[SymbolicData] = set() 275 | for node in execution_order_nodes: 276 | if node in already_executed: 277 | continue 278 | nodes_without_repeated_execution.append(node) 279 | already_executed.update(used_nodes.intersection(node._layer_full_siblings)) 280 | 281 | assert len(already_executed) == len(execution_order_nodes) 282 | return nodes_without_repeated_execution 283 | 284 | def _figure_out_execution_order(self): 285 | used_nodes = self._used_nodes() 286 | sort_graph_and_check_DAG(used_nodes) 287 | execution_order_nodes = sorted(self._used_nodes(), key=lambda node: node._execution_order_idx) 288 | assert len(execution_order_nodes) == len(used_nodes) 289 | 290 | for input_node in used_nodes.intersection(self.inputs): # Not all inputs must be in `used_nodes` 291 | execution_order_nodes.remove(input_node) # Exclude inputs, as we don't execute any layers there 292 | 293 | # To avoid calling layers twice when they have multiple outputs 294 | # We remove all nodes with already executed layer from execution order 295 | execution_order_nodes = self._remove_repeated_execution(execution_order_nodes) 296 | 297 | self._execution_order_nodes = execution_order_nodes 298 | self._execution_order_layers = [node.layer for node in self._execution_order_nodes] 299 | 300 | for _idx, node in enumerate(self._execution_order_nodes): 301 | if node._custom_provided_name is not None: 302 | # Use layer name provided by the user 303 | layer_name = node._custom_provided_name 304 | else: 305 | # Extract the default name of the layer provided by pytorch 306 | layer_name = node.layer._get_name() 307 | 308 | # Check how many layers with this name already exist 309 | self._layer_type_counts.setdefault(layer_name, 0) 310 | self._layer_type_counts[layer_name] += 1 311 | 312 | if node._custom_provided_name is None or self._layer_type_counts[layer_name] > 1: 313 | full_layer_name = f"{layer_name}_{self._layer_type_counts[layer_name]}" 314 | else: 315 | # Skip first index only if it was provided by the user 316 | full_layer_name = layer_name 317 | 318 | self._node_to_layer_name[node] = full_layer_name 319 | self.add_module(name=full_layer_name, module=node.layer) 320 | 321 | def __deepcopy__(self, memo): 322 | """This copies a working Module, but the underlying graph structure is not copied!""" 323 | obj = self.detach_from_graph() 324 | memo[id(self)] = obj 325 | return obj 326 | -------------------------------------------------------------------------------- /pytorch_symbolic/symbolic_data.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2022 Szymon Mikler 2 | 3 | from __future__ import annotations 4 | 5 | import logging 6 | from types import MethodWrapperType 7 | from typing import Any, Callable, List, Set, Tuple 8 | 9 | import torch 10 | from torch import nn 11 | 12 | from . import useful_layers 13 | 14 | _SYMBOLIC_DATA_COUNTER = 0 15 | 16 | 17 | class SymbolicData: 18 | def __init__( 19 | self, 20 | value: Any, 21 | parents: Tuple[SymbolicData, ...] = (), 22 | depth: int = 0, 23 | layer: nn.Module | None = None, 24 | batch_size_known: bool = False, 25 | custom_name: str | None = None, 26 | ): 27 | """Grandfather of all Symbolic datatypes. 28 | 29 | Underlying data is a normal Python object, for example a ``dict``. 30 | You can use methods and operators of the underlying object. 31 | You can also unpack or index it, if only the underlying data allows it. 32 | 33 | If the underlying data is ``torch.Tensor``, it should be created as ``SymbolicTensor`` instead. 34 | 35 | Parameters 36 | ---------- 37 | value 38 | parents 39 | depth 40 | layer 41 | batch_size_known 42 | custom_name 43 | 44 | Attributes 45 | ---------- 46 | v : Any 47 | Underlying data that is used during model tracing 48 | layer : nn.Module 49 | A torch.nn.Module that transforms parents' values into this value. Also it's the incoming edge. 50 | depth : int 51 | Maximum of parents' depths plus one 52 | batch_size_known : bool 53 | In case of Input, whether batch size was provided by the user. 54 | For non-Input nodes, batch size is known iff all parents' batch sizes are known. 55 | custom_name : str 56 | Instead of using the default name for the undelying layer, user can provide his own. 57 | """ 58 | global _SYMBOLIC_DATA_COUNTER 59 | self._execution_order_idx = _SYMBOLIC_DATA_COUNTER 60 | _SYMBOLIC_DATA_COUNTER += 1 61 | 62 | # We use Symbolic Data for inheriting only 63 | assert self.__class__ is not SymbolicData, "Symbolic Data should not be created directly!" 64 | 65 | self._value = value 66 | self._underlying_type_name = type(value).__name__ 67 | self._custom_provided_name = custom_name 68 | 69 | self.layer = layer 70 | self.depth = depth 71 | self.batch_size_known = batch_size_known 72 | 73 | self._children: List[SymbolicData] = [] 74 | self._parents: Tuple[SymbolicData, ...] = parents 75 | self._layer_full_siblings: Tuple[SymbolicData, ...] = (self,) 76 | 77 | self._define_class_operators() 78 | 79 | @property 80 | def v(self): 81 | """Get the underlying value.""" 82 | if self._value is None: 83 | self._recalculate_value() 84 | return self._value 85 | 86 | def _define_class_operators(self): 87 | """Define basic operators, e.g. +, -, *, ... 88 | 89 | This allows using operators for Symbolic Data if and only if the underlying data is compatible. 90 | """ 91 | operators = [ 92 | "__abs__", 93 | "__neg__", 94 | "__add__", 95 | "__radd__", 96 | "__sub__", 97 | "__rsub__", 98 | "__mul__", 99 | "__rmul__", 100 | "__pow__", 101 | "__rpow__", 102 | "__mod__", 103 | "__rmod__", 104 | "__truediv__", 105 | "__rtruediv__", 106 | "__and__", 107 | "__rand__", 108 | "__or__", 109 | "__ror__", 110 | "__xor__", 111 | "__rxor__", 112 | "__matmul__", 113 | "__rmatmul__", 114 | ] 115 | 116 | for operator in operators: 117 | if hasattr(self.v, operator) and ( 118 | not hasattr(self.__class__, operator) 119 | or isinstance(getattr(self.__class__, operator), MethodWrapperType) 120 | ): 121 | logging.debug(f"Adding new operator to {self.__class__.__name__}: {operator}") 122 | 123 | def factory(op): 124 | return lambda self, *args, **kwds: self.__getattr__(op)(*args, **kwds) 125 | 126 | setattr(self.__class__, operator, factory(operator)) 127 | 128 | @property 129 | def parents(self) -> Tuple[SymbolicData, ...]: 130 | """Acces the tuple of parents of this node.""" 131 | return tuple(self._parents) 132 | 133 | @property 134 | def children(self) -> Tuple[SymbolicData, ...]: 135 | """Acces the tuple of children of this node.""" 136 | return tuple(self._children) 137 | 138 | def apply_module( 139 | self, 140 | layer: nn.Module, 141 | *others: SymbolicData, 142 | custom_name: str | None = None, 143 | ) -> SymbolicData | Tuple[SymbolicData, ...]: 144 | """Register a new layer in the graph. Layer must be nn.Module.""" 145 | assert all([isinstance(other, SymbolicData) for other in others]), "Works with SymbolicData only!" 146 | 147 | parents = (self, *others) 148 | new_depth = max(parent.depth for parent in parents) + 1 149 | with torch.no_grad(): 150 | new_value = layer.__call__(self.v, *(o.v for o in others)) 151 | 152 | cls = _figure_out_symbolic_type(new_value) 153 | 154 | new_layer_node = cls( 155 | value=new_value, 156 | parents=parents, 157 | layer=layer, 158 | depth=new_depth, 159 | batch_size_known=self.batch_size_known, 160 | custom_name=custom_name, 161 | ) 162 | for parent in parents: 163 | parent._children.append(new_layer_node) 164 | logging.debug(f"Added {new_layer_node} as child of {parent}") 165 | return new_layer_node 166 | 167 | def _recalculate_value(self): 168 | """Recalulate ._value if it is None. 169 | 170 | Sometimes we remove ._value to reduce the memory footprint. This function restores it. 171 | """ 172 | for parent in self._parents: 173 | if parent._value is None: 174 | parent._recalculate_value() # recursive call to parents 175 | with torch.no_grad(): 176 | outputs = self.layer(*(parent._value for parent in self._parents)) 177 | if len(self._layer_full_siblings) == 1: 178 | outputs = (outputs,) 179 | for full_sibling, output in zip(self._layer_full_siblings, outputs): 180 | full_sibling._value = output 181 | 182 | def _clear_value(self): 183 | """Clear the underlying value to save memory.""" 184 | assert len(self._parents) > 0, "Cannot clear the underlying value of input nodes!" 185 | self._value = None 186 | 187 | def __iter__(self): 188 | """Creates the only layer that has multiple children: UnpackLayer. 189 | 190 | Suitable for unpacking results, even nested ones. 191 | """ 192 | layer = useful_layers.UnpackLayer() 193 | new_outputs = layer.__call__(*self.v) 194 | 195 | new_layer_nodes = [] 196 | for new_value in new_outputs: 197 | cls = _figure_out_symbolic_type(new_value) 198 | 199 | new_layer_nodes.append( 200 | cls( 201 | value=new_value, 202 | parents=(self,), 203 | layer=layer, 204 | depth=self.depth + 1, 205 | batch_size_known=self.batch_size_known, 206 | ) 207 | ) 208 | for new_layer_node in new_layer_nodes: 209 | new_layer_node._layer_full_siblings = tuple(new_layer_nodes) 210 | 211 | self._children.extend(new_layer_nodes) 212 | for new_layer_node in new_layer_nodes: 213 | logging.debug(f"Added {new_layer_node} as child of {self}") 214 | for node in new_layer_nodes: 215 | yield node 216 | 217 | def _get_all_nodes_above(self) -> Set[SymbolicData]: 218 | nodes_seen = {self} 219 | to_expand = [self] 220 | while to_expand: 221 | node = to_expand.pop() 222 | for parent in node._parents: 223 | if parent not in nodes_seen: 224 | to_expand.append(parent) 225 | nodes_seen.add(parent) 226 | return nodes_seen 227 | 228 | def _get_all_nodes_below(self) -> Set[SymbolicData]: 229 | nodes_seen = {self} 230 | to_expand = [self] 231 | while to_expand: 232 | node = to_expand.pop() 233 | for child in node._children: 234 | if child not in nodes_seen: 235 | to_expand.append(child) 236 | nodes_seen.add(child) 237 | return nodes_seen 238 | 239 | def _launch_input(self, x): 240 | self._value = x 241 | 242 | def _launch(self): 243 | if len(self._layer_full_siblings) > 1: 244 | assert len(self._parents) == 1 245 | outputs = self.layer(*self._parents[0]._value) 246 | for node, output in zip(self._layer_full_siblings, outputs): 247 | node._value = output 248 | else: 249 | self._value = self.layer(*(parent._value for parent in self._parents)) 250 | 251 | def __len__(self) -> int: 252 | """Length of the symbolic data.""" 253 | return len(self.v) 254 | 255 | def __getitem__(self, idx): 256 | if isinstance(idx, SymbolicData): 257 | layer = useful_layers.SliceLayerSymbolicIdx() 258 | return layer(self, idx) 259 | else: 260 | layer = useful_layers.SliceLayer(idx) 261 | return layer(self) 262 | 263 | def __call__(self, *args, custom_name: str | None = None): 264 | return self.apply_module(*args, custom_name=custom_name) 265 | 266 | def __repr__(self): 267 | addr = f"{self.__class__.__name__} at {hex(id(self))};" 268 | info = f"{len(self._parents)} parents; {len(self._children)} children" 269 | return "<" + addr + " " + info + ">" 270 | 271 | def __hash__(self): 272 | return id(self) 273 | 274 | def __getattr__(self, item): 275 | if ( 276 | hasattr(self.v, item) 277 | and item != "__torch_function__" # because pytorch is wrapping `+` and other operators 278 | ): 279 | return self(useful_layers.GetAttr(item)) 280 | else: 281 | raise AttributeError 282 | 283 | 284 | class SymbolicCallable(SymbolicData): 285 | def __call__(self, *args, **kwds): 286 | assert isinstance(self.v, Callable) 287 | from . import add_to_graph 288 | 289 | def __func__(obj, *args, **kwds): 290 | return obj(*args, **kwds) 291 | 292 | __func__.__name__ = self.v.__name__ 293 | 294 | returns = add_to_graph(__func__, self, *args, **kwds) 295 | if returns.v is NotImplemented: 296 | dtypes = [type(p.v).__name__ for p in returns.parents[1:]] 297 | raise NotImplementedError(f"Operation on {dtypes} returned NonImplemented object!") 298 | return returns 299 | 300 | 301 | class SymbolicTensor(SymbolicData): 302 | def __init__(self, *args, **kwds): 303 | """Recommended to use Symbolic datatype. It mimics and extends ``torch.Tensor`` API. 304 | 305 | Treat it as a placeholder that will be replaced with real data after the model is created. 306 | For calculation purposes treat it as a normal ``torch.Tensor``: add, subtract, multiply, 307 | take absolute value of, index, slice, etc. 308 | """ 309 | super().__init__(*args, **kwds) 310 | assert isinstance(self.v, torch.Tensor) 311 | self._shape = self.v.shape 312 | 313 | @property 314 | def features(self) -> int: 315 | """Size of the last dimension.""" 316 | return self._shape[-1] 317 | 318 | @property 319 | def C(self) -> int: 320 | """Number of channels in Image data.""" 321 | assert len(self._shape) == 4, "The data is not of [C,H,W] form!" 322 | return self._shape[1] 323 | 324 | @property 325 | def channels(self) -> int: 326 | """Same as ``.C``""" 327 | return self.C 328 | 329 | @property 330 | def H(self) -> int: 331 | """Height in Image data.""" 332 | assert len(self._shape) == 4, "The data is not of [C,H,W] form!" 333 | return self._shape[2] 334 | 335 | @property 336 | def W(self) -> int: 337 | """Width in Image data.""" 338 | assert len(self._shape) == 4, "The data is not of [C,H,W] form!" 339 | return self._shape[3] 340 | 341 | @property 342 | def HW(self) -> Tuple[int, int]: 343 | """Tuple of (height, width) in Image data.""" 344 | return (self.H, self.W) 345 | 346 | @property 347 | def CHW(self) -> Tuple[int, int, int]: 348 | """Tuple of (channels, height, width) in Image data.""" 349 | return (self.C, self.H, self.W) 350 | 351 | @property 352 | def HWC(self) -> Tuple[int, int, int]: 353 | """Tuple of (height, width, channels) in Image data.""" 354 | return (self.H, self.W, self.C) 355 | 356 | @property 357 | def batch_size(self) -> int | None: 358 | """Batch size of the data. Will be default if was not provided.""" 359 | return self._shape[0] 360 | 361 | @property 362 | def shape(self) -> Tuple[int | None, ...]: 363 | """Shape of the underlying Symbolic Tensor, including batch size.""" 364 | return self._shape 365 | 366 | @property 367 | def numel(self) -> int: 368 | """Number of the values in underlying Symbolic Tensor. If batch size is known, it is used too.""" 369 | return self._shape.numel() 370 | 371 | # These methods do not need to be defined, because SymbolicData is redirecting __getattr__. 372 | # However, we define basic methods to ensure they will be used without overhead of __getattr__. 373 | 374 | def reshape(self, *shape) -> SymbolicTensor: 375 | reshape_layer = useful_layers.ReshapeLayer(shape, batch_size_included=True) 376 | return reshape_layer(self) 377 | 378 | def view(self, *shape) -> SymbolicTensor: 379 | view_copy_layer = useful_layers.ViewCopyLayer(shape, batch_size_included=True) 380 | return view_copy_layer(self) 381 | 382 | def t(self) -> SymbolicTensor: 383 | transpose_layer = useful_layers.LambdaOpLayer(op=lambda x: x.t()) 384 | return transpose_layer(self) 385 | 386 | @property 387 | def T(self) -> SymbolicTensor: 388 | transpose_layer = useful_layers.LambdaOpLayer(op=lambda x: x.T) 389 | return transpose_layer(self) 390 | 391 | def mean(self, dim=None, keepdim=False) -> SymbolicTensor: 392 | layer = useful_layers.AggregateLayer(torch.mean, dim=dim, keepdim=keepdim) 393 | return layer(self) 394 | 395 | def sum(self, dim=None, keepdim=False) -> SymbolicTensor: 396 | layer = useful_layers.AggregateLayer(torch.sum, dim=dim, keepdim=keepdim) 397 | return layer(self) 398 | 399 | def median(self, dim=None, keepdim=False) -> SymbolicTensor: 400 | layer = useful_layers.AggregateLayer(torch.median, dim=dim, keepdim=keepdim) 401 | return layer(self) 402 | 403 | def argmax(self, dim=None, keepdim=False) -> SymbolicTensor: 404 | layer = useful_layers.AggregateLayer(torch.argmax, dim=dim, keepdim=keepdim) 405 | return layer(self) 406 | 407 | def argmin(self, dim=None, keepdim=False) -> SymbolicTensor: 408 | layer = useful_layers.AggregateLayer(torch.argmin, dim=dim, keepdim=keepdim) 409 | return layer(self) 410 | 411 | def flatten(self, start_dim=0, end_dim=-1) -> SymbolicTensor: 412 | return nn.Flatten(start_dim, end_dim)(self) 413 | 414 | # These operators do not need to be defined! 415 | # However, we define basic operators to ensure they will be used without overhead of __getattr__. 416 | 417 | def __abs__(self): 418 | return self(useful_layers.LambdaOpLayer(lambda x: abs(x))) 419 | 420 | def __neg__(self): 421 | return self(useful_layers.LambdaOpLayer(op=lambda x: -x)) 422 | 423 | def __add__(self, other): 424 | if isinstance(other, SymbolicTensor): 425 | return self(useful_layers.AddOpLayer(), other) 426 | else: 427 | return self(useful_layers.LambdaOpLayer(op=lambda x: x + other)) 428 | 429 | def __radd__(self, other): 430 | return self.__add__(other) 431 | 432 | def __mul__(self, other): 433 | if isinstance(other, SymbolicTensor): 434 | return self(useful_layers.MulOpLayer(), other) 435 | else: 436 | return self(useful_layers.LambdaOpLayer(op=lambda x: x * other)) 437 | 438 | def __rmul__(self, other): 439 | return self.__mul__(other) 440 | 441 | def __mod__(self, other): 442 | if isinstance(other, SymbolicTensor): 443 | return self(useful_layers.ModOpLayer(), other) 444 | else: 445 | return self(useful_layers.LambdaOpLayer(op=lambda x: x % other)) 446 | 447 | def __rmod__(self, other): 448 | return self(useful_layers.LambdaOpLayer(op=lambda x: other % x)) 449 | 450 | def __pow__(self, other): 451 | if isinstance(other, SymbolicTensor): 452 | return self(useful_layers.LambdaOpLayer(op=lambda x, y: x**y), other) 453 | else: 454 | return self(useful_layers.LambdaOpLayer(op=lambda x: x**other)) 455 | 456 | def __rpow__(self, other): 457 | return self(useful_layers.LambdaOpLayer(op=lambda x: other**x)) 458 | 459 | def __sub__(self, other): 460 | if isinstance(other, SymbolicTensor): 461 | return self(useful_layers.SubOpLayer(), other) 462 | else: 463 | return self(useful_layers.LambdaOpLayer(op=lambda x: x - other)) 464 | 465 | def __rsub__(self, other): 466 | return self(useful_layers.LambdaOpLayer(op=lambda x: other - x)) 467 | 468 | def __truediv__(self, other): 469 | if isinstance(other, SymbolicTensor): 470 | return self(useful_layers.LambdaOpLayer(op=lambda x, y: x / y), other) 471 | else: 472 | return self(useful_layers.LambdaOpLayer(op=lambda x: x / other)) 473 | 474 | def __rtruediv__(self, other): 475 | return self(useful_layers.LambdaOpLayer(op=lambda x: other / x)) 476 | 477 | def __matmul__(self, other): 478 | if isinstance(other, SymbolicTensor): 479 | return self(useful_layers.MatmulOpLayer(), other) 480 | else: 481 | return self(useful_layers.LambdaOpLayer(op=lambda x: x @ other)) 482 | 483 | def __rmatmul__(self, other): 484 | return self(useful_layers.LambdaOpLayer(op=lambda x: other @ x)) 485 | 486 | 487 | _SYMBOLIC_FACTORY_CACHE = {} 488 | 489 | 490 | def SymbolicFactory(dtype): 491 | global _SYMBOLIC_FACTORY_CACHE 492 | dtype_name = dtype.__name__ 493 | 494 | if dtype_name not in _SYMBOLIC_FACTORY_CACHE: 495 | logging.debug(f"New underlying data detected: {dtype_name}!") 496 | cls = type(f"SymbolicData({dtype_name})", (SymbolicData,), {}) 497 | _SYMBOLIC_FACTORY_CACHE[dtype_name] = cls 498 | return _SYMBOLIC_FACTORY_CACHE[dtype_name] 499 | 500 | 501 | def Input( 502 | shape: Tuple | List = (), 503 | batch_size: int = 1, 504 | batch_shape: Tuple | List | None = None, 505 | dtype=torch.float32, 506 | min_value: float = 0.0, 507 | max_value: float = 1.0, 508 | ) -> SymbolicTensor: 509 | """Input to Symbolic Model. Create Symbolic Tensor as a root node in the graph. 510 | 511 | Symbolic Tensor returned by Input has no parents while every other Symbolic Tensor has at least one. 512 | 513 | Parameters 514 | ---------- 515 | shape 516 | Shape of the real data NOT including the batch dimension 517 | batch_size 518 | Optional batch size of the Tensor 519 | batch_shape 520 | Shape of the real data including the batch dimension. 521 | If both ``shape`` and ``batch_shape`` are given, ``batch_shape`` has higher priority. 522 | dtype 523 | Dtype of the real data that will be the input of the network 524 | min_value 525 | In rare cases, if real world data is very specific and some values 526 | cannot work with the model, this should be used to set a 527 | reasonable minimal value that the model can take as an input. 528 | max_value 529 | As above, but the maximal value 530 | 531 | Returns 532 | ------- 533 | SymbolicTensor 534 | Root node in the graph 535 | """ 536 | batch_size_known = True 537 | 538 | if batch_shape is not None: 539 | batch_size = batch_shape[0] 540 | shape = batch_shape[1:] 541 | else: 542 | # By default, we use batch_size of 1 under the hood 543 | batch_size_known = False 544 | 545 | value = torch.rand(batch_size, *shape) * (max_value - min_value) + min_value 546 | value = value.to(dtype) 547 | return SymbolicTensor(value=value, batch_size_known=batch_size_known) 548 | 549 | 550 | def CustomInput(data: Any) -> SymbolicData: 551 | """Input to Symbolic Model. Creates Symbolic Data as a root node in the graph. 552 | 553 | This should be used when Input won't work. 554 | 555 | Parameters 556 | ---------- 557 | data 558 | Speficic data that will be used during the graph tracing. 559 | It can, but doesn't need to be a torch.Tensor. 560 | 561 | Returns 562 | ------- 563 | SymbolicData 564 | Root node in the graph 565 | """ 566 | cls = _figure_out_symbolic_type(data) 567 | return cls(value=data, batch_size_known=True) 568 | 569 | 570 | def _figure_out_symbolic_type(v): 571 | if isinstance(v, torch.Tensor): 572 | cls = SymbolicTensor 573 | elif isinstance(v, Callable): 574 | cls = SymbolicCallable 575 | else: 576 | cls = SymbolicFactory(type(v)) 577 | return cls 578 | -------------------------------------------------------------------------------- /docs/advanced_topics.md: -------------------------------------------------------------------------------- 1 | # Advanced Topics 2 | 3 | ## Underlying graphs 4 | 5 | Deep neural networks can be represented as directed acyclic graphs (DAGs) where: 6 | 7 | * nodes are inputs, outputs or intermediate states 8 | * edges are layers (in general transformations or functions) 9 | 10 | In such graph, 11 | there exists a nonempty set of input nodes and a nonempty set of output nodes. 12 | If your architecture meets the above conditions, it can be created in a symbolic manner. 13 | 14 | Every time you execute `Input`, you create a node in this graph. 15 | 16 | ```python 17 | from pytorch_symbolic import Input 18 | 19 | x = Input(shape=(1,)) 20 | print(x) 21 | ``` 22 | 23 | ```stdout 24 | 25 | ``` 26 | 27 | `Input` is a function for creating Symbolic Data. 28 | Symbolic Tensor is a special case of Symbolic Data. 29 | The difference between them is the underlying object. 30 | Basically, Symbolic Data can have any Python object underneath 31 | and Symbolic Tensor always has `torch.Tensor` underneath. 32 | You should use Symbolic Tensor whenever possible as the library is tested with it in mind. 33 | But if needed, you can define whole graphs of operations with other underlying datatypes. 34 | 35 | Three rules that should hold true: 36 | 37 | 1. `SymbolicData` returned from `Input` has no parents 38 | 2. every other `SymbolicData` has at least one parent 39 | 3. every other `SymbolicData` is a result of an operation on `SymbolicData` 40 | 41 | The provided public API will not let you break these, 42 | but if you are really determined, you can break them by modifying private variables. 43 | This can lead to directed cycles in your graph. If this happens, you probably won't be able to 44 | create the model from your graph, instead you'll be attacked with an assertion error. 45 | 46 | One exception is when your graph has directed cycles, 47 | but you create the model from a subgraph 48 | (all nodes between `inputs` and `outputs` induce the subgraph) 49 | that doesn't contain any directed cycle. Then you'll be fine and the model will be created. 50 | 51 | In general, when you transform any Symbolic Data, a new Symbolic Data is created. 52 | 53 | ```py 54 | y = x + 2 55 | print(y) 56 | ``` 57 | 58 | ```stdout 59 | 60 | ``` 61 | 62 | This _transformation_ is always just an `torch.nn.Module`. 63 | Even when you use `+` or `-` operators, there's an `torch.nn.Module` created underneath. 64 | In case of `+` and `-` it's a module that adds or subtracts two inputs. 65 | These operators are defined using `useful_layers.LambdaOpLayer`. 66 | 67 | ```py 68 | print(y.layer) 69 | ``` 70 | 71 | ```stdout 72 | LambdaOpLayer() 73 | ``` 74 | 75 | When you add new nodes to the graph as a result of some operation, 76 | they will be registered as children of other nodes, 77 | You can check children or parents of a node by accessing: 78 | 79 | * `node.parents: Tuple[SymbolicData, ...]` 80 | * `node.children: Tuple[SymbolicData, ...]` 81 | 82 | It's forbidden to modify parents after Symbolic Data is created, 83 | but new children can always be added by applying new operations. 84 | Every operation on Symbolic Data creates at least one new child. 85 | If your `torch.nn.Module` has more than one output, more children will be created. 86 | Usually modules have just one output: 87 | 88 | ```py 89 | print(y) 90 | a = y - 1 91 | print(y) 92 | b = y / 2 93 | print(y) 94 | ``` 95 | 96 | ```stdout 97 | 98 | 99 | 100 | ``` 101 | 102 | We created a bunch of children for `y`. Let's see them: 103 | 104 | ```py 105 | print(y.children) 106 | ``` 107 | 108 | ```stdout 109 | (, 110 | ) 111 | ``` 112 | 113 | But when working with large graphs, it might not be fun to inspect 114 | them by printing children and parents to stdout. Luckily, there is a nicer alternative: 115 | Pytorch Symbolic provides a basic graph drawing utility. 116 | It will work only if you have optional dependencies installed: 117 | 118 | ``` 119 | networkx 120 | matplotlib 121 | scipy 122 | ``` 123 | 124 | You can install them manually or using `pip install pytorch-symbolic[full]`. 125 | 126 | > If you don't have optional dependencies, you can use the package without drawing utility. 127 | 128 | Let us draw our graph: 129 | 130 | ```py 131 | from pytorch_symbolic import graph_algorithms 132 | 133 | graph_algorithms.draw_graph(inputs=x, node_text_namespace=globals()) 134 | ``` 135 | 136 | ![images/draw_graph1.png](images/draw_graph1.png) 137 | 138 | We use `node_text_namespace=globals()` so that Pytorch Symbolic attempts 139 | to display variable names on nodes. 140 | This is useful for understanding what is going on in the graph. 141 | It can be more difficult to use when graph was defined in a local namespace. 142 | Alternatively, you can use argument `node_text_func` to define your 143 | own custom labels on nodes. 144 | This argument is a `Callable` that takes Symbolic Data as input and returns `str` as output. 145 | By default, it displays underlying tensor shape. 146 | 147 | > **Tip:** consider using `try: ... except: ...` in your `node_text_func`. 148 | > There are many data types a node can have underneath, so the chances are that you will 149 | > get an `AttributeError` otherwise. 150 | 151 | Be careful! Drawing utility is designed to display simple graphs 152 | so it might not work well for large neural networks. 153 | You can tune the figure size and other stuff using `matplotlib` library: 154 | 155 | ```py 156 | import matplotlib.pyplot as plt 157 | 158 | nodes = [a + i for i in range(10)] 159 | out = sum(nodes) 160 | 161 | plt.figure(dpi=300) 162 | 163 | graph_algorithms.draw_graph( 164 | inputs=x, 165 | outputs=out, 166 | edge_text_func=lambda x: "", 167 | node_text_namespace=globals(), 168 | figsize=(10, 10), 169 | ) 170 | ``` 171 | 172 | ![images/draw_graph2.png](images/draw_graph2.png) 173 | 174 | Question mark displayed instead the name of a variable indicates that the variable 175 | was not found in the given namespace. 176 | Similarly to `node_text_func` you can use `edge_text_func` to display custom labels on the edges. 177 | By default, we display the name of the layer, e.g. `Linear` or `Conv2d`. 178 | 179 | ## Creating models 180 | 181 | After creating Symbolic Data and defining the operations, 182 | you want to enclose them in a model. 183 | It might be a neural network or just an arbitrary graph of computations. 184 | After creating Symbolic Model you will be able to use it with your own data. 185 | 186 | ```python 187 | from torch import nn 188 | from pytorch_symbolic import Input, SymbolicModel 189 | 190 | inputs = x = Input((784,)) 191 | 192 | for _ in range(3): 193 | x = nn.Linear(x.features, 100)(x) 194 | 195 | x = nn.Linear(x.features, 10)(x, custom_name="Classifier") 196 | model = SymbolicModel(inputs, x) 197 | model.summary() # Keras-like summary 198 | ``` 199 | 200 | ```stdout 201 | ___________________________________________________ 202 | Layer Output shape Params Parent 203 | =================================================== 204 | 1 Input_1 (None, 784) 0 205 | 2 Linear_1 (None, 100) 78500 1 206 | 3 Linear_2 (None, 100) 10100 2 207 | 4 Linear_3 (None, 100) 10100 3 208 | 5* Classifier (None, 10) 1010 4 209 | =================================================== 210 | Total params: 99710 211 | Trainable params: 99710 212 | Non-trainable params: 0 213 | ___________________________________________________ 214 | ``` 215 | 216 | You can use the plotting utility directly on your model: 217 | 218 | ```py 219 | from pytorch_symbolic import graph_algorithms 220 | 221 | graph_algorithms.draw_graph(model=model) 222 | ``` 223 | 224 | ![images/draw_graph3.png](images/draw_graph3.png) 225 | 226 | As you can see, if the batch size is not specified, Pytorch Symbolic uses batch size of 1. 227 | 228 | ## Multiple inputs and outputs 229 | 230 | If your `torch.nn.Module` has multiple inputs or outputs, that's fine: 231 | 232 | ```python 233 | from torch import nn 234 | from pytorch_symbolic import Input, SymbolicModel, useful_layers, graph_algorithms 235 | 236 | 237 | class AddN(nn.Module): 238 | def forward(self, *tensors): 239 | return sum(tensors) 240 | 241 | 242 | inputs = [Input((5,)) for _ in range(5)] 243 | intermediate = AddN()(*inputs) 244 | outputs = [intermediate / i for i in range(1, 5)] 245 | 246 | model = SymbolicModel(inputs, outputs) 247 | model.summary() 248 | graph_algorithms.draw_graph(model=model, figsize=(9, 6)) 249 | ``` 250 | 251 | ```stdout 252 | ____________________________________________________________ 253 | Layer Output shape Params Parent 254 | ============================================================ 255 | 1 Input_1 (None, 5) 0 256 | 2 Input_2 (None, 5) 0 257 | 3 Input_3 (None, 5) 0 258 | 4 Input_4 (None, 5) 0 259 | 5 Input_5 (None, 5) 0 260 | 6 LambdaOpLayer_1 (None, 5) 0 1,2,3,4,5 261 | 7* LambdaOpLayer_2 (None, 5) 0 6 262 | 8* LambdaOpLayer_3 (None, 5) 0 6 263 | 9* LambdaOpLayer_4 (None, 5) 0 6 264 | 10* LambdaOpLayer_5 (None, 5) 0 6 265 | ============================================================ 266 | Total params: 0 267 | Trainable params: 0 268 | Non-trainable params: 0 269 | ____________________________________________________________ 270 | ``` 271 | 272 | ![images/draw_graph4.png](images/draw_graph4.png) 273 | 274 | Notice that we used custom `AddN` module instead of just `intermediate = sum(inputs)`. 275 | Both versions would be correct. 276 | In fact, they produce equivalent models, but their underlying graphs are different. 277 | 278 | Let us compare two smaller examples: 279 | 280 | ```py 281 | inputs = [Input((5,)) for _ in range(3)] 282 | graph_algorithms.draw_graph(inputs=inputs, outputs=AddN()(*inputs)) 283 | ``` 284 | 285 | ![images/draw_graph5.png](images/draw_graph5.png) 286 | 287 | And: 288 | 289 | ``` 290 | graph_algorithms.draw_graph(inputs=inputs, outputs=sum(inputs)) 291 | ``` 292 | 293 | ![images/draw_graph6.png](images/draw_graph6.png) 294 | 295 | This difference exists because `sum` is executing multiple `__add__` operations under the hood. 296 | 297 | ## Reusing existing layers 298 | 299 | Each node, except those created by `Input`, is associated with some `torch.nn.Module`. 300 | When you are reusing a part of the graph, you are reusing all underlying 301 | `torch.nn.Module` too. 302 | Thanks to this, you can have multiple models sharing the same weights. 303 | Or one model using the same weights multiple times. 304 | 305 | For example, imagine you created a model that classifies RGB images: 306 | 307 | ```python 308 | from torch import nn 309 | from pytorch_symbolic import Input, SymbolicModel 310 | 311 | inputs = x = Input((1, 28, 28)) 312 | 313 | for _ in range(3): 314 | x = nn.Conv2d(x.C, 8, 3, padding=1)(x)(nn.ReLU()) 315 | 316 | features = x 317 | 318 | x = nn.Flatten()(x) 319 | outputs = nn.Linear(x.features, 10)(x) 320 | 321 | classifier = SymbolicModel(inputs, outputs) 322 | classifier.summary() 323 | ``` 324 | 325 | ```stdout 326 | _______________________________________________________ 327 | Layer Output shape Params Parent 328 | ======================================================= 329 | 1 Input_1 (None, 1, 28, 28) 0 330 | 2 Conv2d_1 (None, 8, 28, 28) 80 1 331 | 3 ReLU_1 (None, 8, 28, 28) 0 2 332 | 4 Conv2d_2 (None, 8, 28, 28) 584 3 333 | 5 ReLU_2 (None, 8, 28, 28) 0 4 334 | 6 Conv2d_3 (None, 8, 28, 28) 584 5 335 | 7 ReLU_3 (None, 8, 28, 28) 0 6 336 | 8 Flatten_1 (None, 6272) 0 7 337 | 9* Linear_1 (None, 10) 62730 8 338 | ======================================================= 339 | Total params: 63978 340 | Trainable params: 63978 341 | Non-trainable params: 0 342 | _______________________________________________________ 343 | ``` 344 | 345 | After training `classifier`, you might decide that you want to inspect the intermediate features. 346 | 347 | ```py 348 | feature_extractor = SymbolicModel(inputs, features) 349 | feature_extractor.summary() 350 | ``` 351 | 352 | ```stdout 353 | ______________________________________________________ 354 | Layer Output shape Params Parent 355 | ====================================================== 356 | 1 Input_1 (None, 1, 28, 28) 0 357 | 2 Conv2d_1 (None, 8, 28, 28) 80 1 358 | 3 ReLU_1 (None, 8, 28, 28) 0 2 359 | 4 Conv2d_2 (None, 8, 28, 28) 584 3 360 | 5 ReLU_2 (None, 8, 28, 28) 0 4 361 | 6 Conv2d_3 (None, 8, 28, 28) 584 5 362 | 7* ReLU_3 (None, 8, 28, 28) 0 6 363 | ====================================================== 364 | Total params: 1248 365 | Trainable params: 1248 366 | Non-trainable params: 0 367 | ______________________________________________________ 368 | ``` 369 | 370 | This model, `feature_extractor`, uses the same underlying weights 371 | as already trained `classifier`, 372 | but it outputs the intermediate features you wanted to inspect. 373 | It does so without modifying the original model! 374 | 375 | There's also another way of achieving the end result, this one modifies the original model: 376 | 377 | ```py 378 | classifier.add_output(features) 379 | classifier.summary() 380 | ``` 381 | 382 | ```stdout 383 | _______________________________________________________ 384 | Layer Output shape Params Parent 385 | ======================================================= 386 | 1 Input_1 (None, 1, 28, 28) 0 387 | 2 Conv2d_1 (None, 8, 28, 28) 80 1 388 | 3 ReLU_1 (None, 8, 28, 28) 0 2 389 | 4 Conv2d_2 (None, 8, 28, 28) 584 3 390 | 5 ReLU_2 (None, 8, 28, 28) 0 4 391 | 6 Conv2d_3 (None, 8, 28, 28) 584 5 392 | 7* ReLU_3 (None, 8, 28, 28) 0 6 393 | 8 Flatten_1 (None, 6272) 0 7 394 | 9* Linear_1 (None, 10) 62730 8 395 | ======================================================= 396 | Total params: 63978 397 | Trainable params: 63978 398 | Non-trainable params: 0 399 | _______________________________________________________ 400 | ``` 401 | 402 | Now there's a dot next to the 7th and 9th layers indicating both 7th and 9th layers' outputs 403 | are returned from the model. Previously only 9th layer's output was returned. 404 | 405 | ## Nested models 406 | 407 | Instance of Symbolic Model is just a `torch.nn.Module` and you can use it as such. 408 | This means you can use it anywhere, including another Symbolic Model or vanilla model. 409 | Create new models, using the existing ones. 410 | Here we create a model that calculates how similar are two feature maps generated 411 | by previously defined `feature_extractor`: 412 | 413 | ```py 414 | import torch 415 | from pytorch_symbolic import add_to_graph, graph_algorithms 416 | 417 | inputs1 = Input((1, 32, 32)) 418 | inputs2 = Input((1, 64, 64)) 419 | 420 | noise = Input((1, 64, 64)) 421 | 422 | features1 = feature_extractor(inputs1) 423 | features2 = feature_extractor(inputs2 + noise) 424 | features2 = nn.MaxPool2d(2)(features2) 425 | 426 | diffs = (features1 - features2) ** 2 427 | 428 | outputs = add_to_graph(torch.sum, diffs, dim=(1, 2, 3)) 429 | strange_model = SymbolicModel((inputs1, inputs2, noise), outputs) 430 | strange_model.summary() 431 | ``` 432 | 433 | ```stdout 434 | ______________________________________________________________ 435 | Layer Output shape Params Parent 436 | ============================================================== 437 | 1 Input_1 (None, 1, 32, 32) 0 438 | 2 Input_2 (None, 1, 64, 64) 0 439 | 3 Input_3 (None, 1, 64, 64) 0 440 | 4 SymbolicModel_1 (None, 8, 32, 32) 1248 1 441 | 5 AddOpLayer_1 (None, 1, 64, 64) 0 2,3 442 | 6 SymbolicModel_2 (None, 8, 64, 64) 1248 5 443 | 7 MaxPool2d_1 (None, 8, 32, 32) 0 6 444 | 8 SubOpLayer_1 (None, 8, 32, 32) 0 4,7 445 | 9 LambdaOpLayer_1 (None, 8, 32, 32) 0 8 446 | 10* wrap(sum)_1 (None,) 0 9 447 | ============================================================== 448 | Total params: 1248 449 | Trainable params: 1248 450 | Non-trainable params: 0 451 | ______________________________________________________________ 452 | ``` 453 | 454 | ```py 455 | graph_algorithms.draw_graph(model=strange_model, figsize=(10, 8)) 456 | ``` 457 | 458 | ![images/draw_graph7.png](images/draw_graph7.png) 459 | 460 | ## Recreate Toy ResNet 461 | 462 | We took an example of toy ResNet from 463 | [tensorflow guide](https://www.tensorflow.org/guide/keras/symbolic) and 464 | recreated it in a few different ways. Note that their code for model definition 465 | is **16 lines of code long**, excluding imports and utilities. 466 | 467 | > In [benchmarks](benchmarks.md) you can see toy ResNet benchmarked! 468 | 469 | Using Pytorch Symbolic, you can create toy ResNet using exactly as many lines as using Keras: 470 | 471 | ```python 472 | from torch import nn 473 | from pytorch_symbolic import Input, SymbolicModel 474 | 475 | inputs = Input(shape=(3, 32, 32)) 476 | x = nn.Conv2d(inputs.C, 32, 3)(inputs)(nn.ReLU()) 477 | x = nn.Conv2d(x.C, 64, 3)(x)(nn.ReLU()) 478 | block_1_output = nn.MaxPool2d(3)(x) 479 | 480 | x = nn.Conv2d(block_1_output.C, 64, 3, padding=1)(block_1_output)(nn.ReLU()) 481 | x = nn.Conv2d(x.C, 64, 3, padding=1)(x)(nn.ReLU()) 482 | block_2_output = x + block_1_output 483 | 484 | x = nn.Conv2d(block_2_output.C, 64, 3, padding=1)(block_2_output)(nn.ReLU()) 485 | x = nn.Conv2d(x.C, 64, 3, padding=1)(x)(nn.ReLU()) 486 | block_3_output = x + block_2_output 487 | 488 | x = nn.Conv2d(block_3_output.C, 64, 3)(block_3_output)(nn.ReLU()) 489 | x = nn.AvgPool2d(kernel_size=x.HW)(x)(nn.Flatten()) 490 | x = nn.Linear(x.features, 256)(x)(nn.ReLU()) 491 | x = nn.Dropout(0.5)(x) 492 | outputs = nn.Linear(x.features, 10)(x) 493 | 494 | model = SymbolicModel(inputs, outputs) 495 | ``` 496 | 497 | In fact, every line of code in Pytorch Symbolic and Keras is functionally equivalent. 498 | 499 | For example this line in Keras: 500 | 501 | ```python 502 | ... = layers.Conv2D(64, 3, activation="relu", padding="same")(x) 503 | ``` 504 | 505 | is equivalent to this line in Pytorch Symbolic: 506 | 507 | ```python 508 | ... = nn.Conv2d(x.C, 64, 3, padding=1)(x)(nn.ReLU()) 509 | ``` 510 | 511 | Let's analyze what happens in the above lines: 512 | 513 | * `torch.nn.Conv2d` is PyTorch equivalent of Keras `keras.layers.Conv2d` layer 514 | * Input channels: 515 | * In Keras we don't pass them openly - they'll be calculated automatically from the inputs 516 | * In Pytorch Symbolic we also calculate them automatically using `x.C`, but we pass them explicitly as an argument 517 | * In both frameworks `64, 3` are the number of output channels and the size of the kernel 518 | * Padding: 519 | * We use `padding="same"` in Keras 520 | * We use `padding=1` in PyTorch (for kernel size 3) 521 | * Activation: 522 | * In Keras we simply add an argument `activation='relu'` 523 | * in Pytorch Symbolic we add `torch.nn.ReLU()` as a transformation that happens 524 | after `torch.nn.Conv2d(...)` 525 | 526 | The example below is equivalent, but uses the other way of registering layers in the network: 527 | 528 | ```python 529 | from torch import nn 530 | from pytorch_symbolic import Input, SymbolicModel 531 | 532 | inputs = Input(shape=(3, 32, 32)) 533 | x = inputs(nn.Conv2d(inputs.channels, 32, 3))(nn.ReLU()) 534 | x = x(nn.Conv2d(x.channels, 64, 3))(nn.ReLU()) 535 | block_1_output = x(nn.MaxPool2d(3)) 536 | 537 | x = block_1_output(nn.Conv2d(block_1_output.channels, 64, 3, padding=1))(nn.ReLU()) 538 | x = x(nn.Conv2d(x.channels, 64, 3, padding=1))(nn.ReLU()) 539 | block_2_output = x + block_1_output 540 | 541 | x = block_2_output(nn.Conv2d(block_2_output.channels, 64, 3, padding=1))(nn.ReLU()) 542 | x = x(nn.Conv2d(x.channels, 64, 3, padding=1))(nn.ReLU()) 543 | block_3_output = x + block_2_output 544 | 545 | x = block_3_output(nn.Conv2d(x.channels, 64, 3))(nn.ReLU()) 546 | x = x(nn.AvgPool2d(kernel_size=(x.H, x.W)))(nn.Flatten()) 547 | x = x(nn.Linear(x.features, 256))(nn.ReLU()) 548 | x = x(nn.Dropout(0.5)) 549 | outputs = x(nn.Linear(x.features, 10)) 550 | 551 | model = SymbolicModel(inputs, outputs) 552 | ``` 553 | 554 | This took 16 lines of code as well. 555 | 556 | You can register new layers in whichever way you prefer, you can mix them too. 557 | 558 | ### Vanilla PyTorch 559 | 560 | A usual way to define a model in PyTorch is to create a class that inherits from `torch.nn.Module`. 561 | 562 | PyTorch non-symbolic example of toy ResNet from the previous section: 563 | 564 | ```python 565 | from torch import nn 566 | 567 | 568 | class ToyResNet(nn.Module): 569 | def __init__(self): 570 | super().__init__() 571 | self.relu = nn.ReLU() 572 | self.block1conv1 = nn.Conv2d(3, 32, 3) 573 | self.block1conv2 = nn.Conv2d(32, 64, 3) 574 | self.maxpool = nn.MaxPool2d(3) 575 | 576 | self.block2conv1 = nn.Conv2d(64, 64, 3, padding=1) 577 | self.block2conv2 = nn.Conv2d(64, 64, 3, padding=1) 578 | 579 | self.block3conv1 = nn.Conv2d(64, 64, 3, padding=1) 580 | self.block3conv2 = nn.Conv2d(64, 64, 3, padding=1) 581 | 582 | self.conv1 = nn.Conv2d(64, 64, 3) 583 | 584 | kernel_size = 7 # calculated by hand 585 | self.global_pool = nn.AvgPool2d(kernel_size) 586 | self.flatten = nn.Flatten() 587 | self.linear = nn.Linear(64, 256) 588 | self.dropout = nn.Dropout(0.5) 589 | self.classifier = nn.Linear(256, 10) 590 | 591 | def forward(self, x): 592 | x = self.relu(self.block1conv1(x)) 593 | x = self.relu(self.block1conv2(x)) 594 | block_1_output = self.maxpool(x) 595 | 596 | x = self.relu(self.block2conv1(block_1_output)) 597 | x = self.relu(self.block2conv2(x)) 598 | block_2_output = x + block_1_output 599 | 600 | x = self.relu(self.block3conv1(block_2_output)) 601 | x = self.relu(self.block3conv2(x)) 602 | block_3_output = x + block_2_output 603 | 604 | x = self.relu(self.conv1(block_3_output)) 605 | x = self.global_pool(x) 606 | x = self.flatten(x) 607 | x = self.relu(self.linear(x)) 608 | x = self.dropout(x) 609 | return self.classifier(x) 610 | 611 | 612 | model = ToyResNet() 613 | ``` 614 | 615 | This took over 30 lines of code. 616 | 617 | ## Advanced custom functions 618 | 619 | If you read [Quick Start](quick_start.md), you know that the best way to include a custom 620 | function in your model is to use it inside a `torch.nn.Module`. 621 | Just put the function call in `forward` method. If your function needs additional arguments, 622 | for example for the dimension, you can specify them in `__init__`, like this: 623 | 624 | ```python 625 | import torch 626 | from torch import nn 627 | 628 | 629 | class ConcatLayer(nn.Module): 630 | def __init__(self, dim): 631 | super().__init__() 632 | self.dim = dim 633 | 634 | def forward(self, *tensors): 635 | return torch.cat(tensors=tensors, dim=self.dim) 636 | ``` 637 | 638 | This way provides a robust and fast runtime. However, you might not want to write 639 | so much boilerplate code. Luckily, with Pytorch Symbolic there is a way to avoid it: 640 | 641 | ```py 642 | from pytorch_symbolic import Input, add_to_graph 643 | 644 | x = Input(batch_shape=(2, 10, 20)) 645 | y = Input(batch_shape=(2, 20, 20)) 646 | concatenated = add_to_graph(torch.cat, (x, y), dim=1) 647 | ``` 648 | 649 | Function `add_to_graph` is powerful, but use it responsibly. 650 | It adds some CPU overhead, because it is parsing arguments in search of Symbolic Data. 651 | This won't matter in GPU heavy workloads, 652 | but might contribute to a slowdown in CPU limited scenarios. 653 | 654 | Here is an overcomplicated example 655 | to give you an idea of what `add_to_graph` allows you to do: 656 | 657 | ```python 658 | import torch 659 | from pytorch_symbolic import Input, SymbolicModel, add_to_graph 660 | 661 | x = Input(batch_shape=(10, 20)) 662 | y = Input(batch_shape=(20, 30)) 663 | z = Input(batch_shape=(30, 10)) 664 | 665 | data_dict = { 666 | "tensors": { 667 | "constant": torch.rand(10, 10), 668 | "symbolic": x, 669 | "parameters": [y, z] 670 | }, 671 | "metadata": { 672 | "date": "01.01.2001", 673 | "name": "Batman", 674 | } 675 | } 676 | 677 | 678 | def execute_noisy(data_dict, *, add_noise: bool = False): 679 | params = data_dict["tensors"]["parameters"] 680 | x = data_dict["tensors"]["symbolic"] 681 | 682 | for param in params: 683 | x = x @ param 684 | 685 | if add_noise: 686 | print("Adding noise!") 687 | x = x + data_dict["tensors"]["constant"] 688 | return x, data_dict["tensors"]["constant"] 689 | 690 | 691 | outputs, noise = add_to_graph(execute_noisy, data_dict, add_noise=True) 692 | model = SymbolicModel(inputs=(x, y, z), outputs=outputs) 693 | model.summary() 694 | ``` 695 | 696 | ```stdout 697 | Adding noise! 698 | ______________________________________________________________ 699 | Layer Output shape Params Parent 700 | ============================================================== 701 | 1 Input_1 (10, 20) 0 702 | 2 Input_2 (20, 30) 0 703 | 3 Input_3 (30, 10) 0 704 | 4 wrap(execute_noisy)_1 tuple 0 1,2,3 705 | 5* UnpackLayer_1 (10, 10) 0 4 706 | ============================================================== 707 | Total params: 0 708 | Trainable params: 0 709 | Non-trainable params: 0 710 | ______________________________________________________________ 711 | ``` 712 | 713 | Your function will be called by `add_to_graph` during tracing 714 | and will be registered in the graph as a `torch.nn.Module` that wraps it. 715 | This operation produces a Symbolic Data which you can use to define more operations 716 | or to create a Symbolic Model. 717 | 718 | This model will execute your custom function during the runtime! Use it as always: 719 | 720 | ```py 721 | x = torch.rand((10, 20)) 722 | y = torch.rand((20, 30)) 723 | z = torch.rand((30, 10)) 724 | 725 | outs = model(x, y, z) 726 | ``` 727 | 728 | ```stdout 729 | Adding noise! 730 | ``` 731 | 732 | Everything that is _not_ a Symbolic Data in `add_to_graph` is considered 733 | constant and will stay the same each time you execute the model. 734 | In our example `add_noise` stayed `True` even after creating the model. 735 | Under the hood, Pytorch Symbolic is browsing through 736 | the arguments in search of Symbolic Data. 737 | Only they will be replaced by new, real data during the execution. 738 | Argument parser is able to navigate through nested `list`, `tuple` and `dict`. 739 | 740 | ## Arbitrary Symbolic Data 741 | 742 | So we've been saying Symbolic Data over and over, 743 | but we always used Symbolic Tensor underneath, which makes it 744 | a special case of Symbolic Data. 745 | But the fact is, Symbolic Data can have an arbitrary Python object underneath. 746 | 747 | ```python 748 | from pytorch_symbolic import CustomInput, SymbolicModel 749 | import numpy as np 750 | 751 | input1 = x = CustomInput(np.random.rand(100, 200)) 752 | input2 = y = CustomInput(np.array([1, 2, 3])) 753 | 754 | z = x[:, :, None] + y.reshape(1, 1, -1) 755 | model = SymbolicModel((input1, input2), outputs=z) 756 | model.summary() 757 | 758 | ``` 759 | 760 | ```stdout 761 | ______________________________________________________________________ 762 | Layer Output shape Params Parent 763 | ====================================================================== 764 | 1 Input_1 ndarray 0 765 | 2 Input_2 ndarray 0 766 | 3 SliceLayer_1 ndarray 0 1 767 | 4 GetAttr_1 method-wrapper 0 3 768 | 5 GetAttr_2 builtin_function_or_method 0 2 769 | 6 wrap(reshape)_1 ndarray 0 5 770 | 7* wrap(__add__)_1 ndarray 0 4,6 771 | ====================================================================== 772 | Total params: 0 773 | Trainable params: 0 774 | Non-trainable params: 0 775 | ______________________________________________________________________ 776 | ``` 777 | 778 | Use it as always, but take care to have inputs compatible with the operations you defined. 779 | 780 | ```py 781 | example_outs = model(np.random.rand(100, 200), np.random.rand(3)) 782 | example_outs.shape 783 | ``` 784 | 785 | ```stdout 786 | (100, 200, 3) 787 | ``` 788 | 789 | The model above defines a graph of operations on underlying `numpy` arrays. 790 | 791 | You can create replayable graphs of your favorite Python operations. 792 | Pytorch Symbolic is a generic framework for enclosing arbitrary Python operations 793 | in a `torch.nn.Module` and replaying them. 794 | 795 | But note a few things: 796 | 797 | * You can use most of the methods defined for your data type, 798 | e.g. `symbolic_numpy.reshape(...)` will return another Symbolic Data with 799 | underlying `ndarray` which you can use to define more operations or create a Symbolic Model. 800 | * Pytorch Symbolic won't play nicely with in-place operations. 801 | For example, when working with Symbolic Data with underlying lists it is better 802 | to use `new = symbolic_list + [5]` instead of `symbolic_list.append(5)`. 803 | If you use in-place operation, it will _not_ be replayed during graph re-execution, 804 | unless its output (usually `None`) is included in Symbolic Model's outputs. 805 | It is your responsibility to avoid in-place operations. 806 | --------------------------------------------------------------------------------