├── .gitignore ├── LICENSE ├── README.md ├── best_shape_fit.py ├── configs ├── lens_shape │ ├── conditional_cinn_1.py │ ├── conditional_cinn_2.py │ ├── conditional_cinn_4.py │ ├── conditional_cinn_8.py │ ├── conditional_hint_1_full.py │ ├── conditional_hint_2_full.py │ ├── conditional_hint_4_full.py │ ├── conditional_hint_8_full.py │ ├── unconditional_hint_1_full.py │ ├── unconditional_hint_2_full.py │ ├── unconditional_inn_1.py │ └── unconditional_inn_2.py ├── plus_shape │ ├── conditional_cinn_1.py │ ├── conditional_cinn_2.py │ ├── conditional_cinn_4.py │ ├── conditional_cinn_8.py │ ├── conditional_hint_1_full.py │ ├── conditional_hint_2_full.py │ ├── conditional_hint_4_full.py │ ├── conditional_hint_8_full.py │ ├── conditional_recursive_cinn_4.py │ ├── unconditional_hint_16_0_big.py │ ├── unconditional_hint_16_0_small.py │ ├── unconditional_hint_16_1.py │ ├── unconditional_hint_16_1_big.py │ ├── unconditional_hint_16_1_small.py │ ├── unconditional_hint_1_full.py │ ├── unconditional_hint_2_full.py │ ├── unconditional_hint_32_0_big.py │ ├── unconditional_hint_32_0_small.py │ ├── unconditional_hint_4_0_big.py │ ├── unconditional_hint_4_0_small.py │ ├── unconditional_hint_4_1.py │ ├── unconditional_hint_4_1_big.py │ ├── unconditional_hint_4_1_small.py │ ├── unconditional_hint_4_2.py │ ├── unconditional_hint_4_2_big.py │ ├── unconditional_hint_4_2_small.py │ ├── unconditional_hint_4_3.py │ ├── unconditional_hint_4_3_big.py │ ├── unconditional_hint_4_3_constwidth.py │ ├── unconditional_hint_4_3_lessshrink.py │ ├── unconditional_hint_4_3_reshuffle.py │ ├── unconditional_hint_4_3_small.py │ ├── unconditional_hint_4_full.py │ ├── unconditional_hint_8_0_big.py │ ├── unconditional_hint_8_0_small.py │ ├── unconditional_hint_8_1.py │ ├── unconditional_hint_8_1_big.py │ ├── unconditional_hint_8_1_small.py │ ├── unconditional_hint_8_2.py │ ├── unconditional_hint_8_2_big.py │ ├── unconditional_hint_8_2_small.py │ ├── unconditional_hint_8_full.py │ ├── unconditional_inn_1.py │ ├── unconditional_inn_16.py │ ├── unconditional_inn_2.py │ ├── unconditional_inn_32.py │ ├── unconditional_inn_4.py │ ├── unconditional_inn_4_Q.py │ └── unconditional_inn_8.py └── uci_data │ ├── gas_hint_4.py │ ├── gas_hint_8.py │ ├── gas_inn_4.py │ ├── gas_inn_8.py │ ├── miniboone_hint_4.py │ ├── miniboone_hint_8.py │ ├── miniboone_inn_4.py │ ├── miniboone_inn_8.py │ ├── power_hint_4.py │ ├── power_hint_8.py │ ├── power_inn_4.py │ └── power_inn_8.py ├── data.py ├── data └── frog.json ├── dense_jacobian.png ├── eval_shapes.py ├── fourier_curve.png ├── hint.py ├── monitoring.py ├── plot_data.py ├── recursive_block.png ├── rejection_sampling.py ├── run_experiments.py ├── run_uci_experiments.py ├── train_conditional.py └── train_unconditional.py /.gitignore: -------------------------------------------------------------------------------- 1 | # Byte-compiled / optimized / DLL files 2 | __pycache__/ 3 | *.py[cod] 4 | *$py.class 5 | 6 | # C extensions 7 | *.so 8 | 9 | # Distribution / packaging 10 | .Python 11 | build/ 12 | develop-eggs/ 13 | dist/ 14 | downloads/ 15 | eggs/ 16 | .eggs/ 17 | lib/ 18 | lib64/ 19 | parts/ 20 | sdist/ 21 | var/ 22 | wheels/ 23 | pip-wheel-metadata/ 24 | share/python-wheels/ 25 | *.egg-info/ 26 | .installed.cfg 27 | *.egg 28 | MANIFEST 29 | 30 | # PyInstaller 31 | # Usually these files are written by a python script from a template 32 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 33 | *.manifest 34 | *.spec 35 | 36 | # Installer logs 37 | pip-log.txt 38 | pip-delete-this-directory.txt 39 | 40 | # Unit test / coverage reports 41 | htmlcov/ 42 | .tox/ 43 | .nox/ 44 | .coverage 45 | .coverage.* 46 | .cache 47 | nosetests.xml 48 | coverage.xml 49 | *.cover 50 | *.py,cover 51 | .hypothesis/ 52 | .pytest_cache/ 53 | 54 | # Translations 55 | *.mo 56 | *.pot 57 | 58 | # Django stuff: 59 | *.log 60 | local_settings.py 61 | db.sqlite3 62 | db.sqlite3-journal 63 | 64 | # Flask stuff: 65 | instance/ 66 | .webassets-cache 67 | 68 | # Scrapy stuff: 69 | .scrapy 70 | 71 | # Sphinx documentation 72 | docs/_build/ 73 | 74 | # PyBuilder 75 | target/ 76 | 77 | # Jupyter Notebook 78 | .ipynb_checkpoints 79 | 80 | # IPython 81 | profile_default/ 82 | ipython_config.py 83 | 84 | # pyenv 85 | .python-version 86 | 87 | # pipenv 88 | # According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. 89 | # However, in case of collaboration, if having platform-specific dependencies or dependencies 90 | # having no cross-platform support, pipenv may install dependencies that don't work, or not 91 | # install all needed dependencies. 92 | #Pipfile.lock 93 | 94 | # PEP 582; used by e.g. github.com/David-OConnor/pyflow 95 | __pypackages__/ 96 | 97 | # Celery stuff 98 | celerybeat-schedule 99 | celerybeat.pid 100 | 101 | # SageMath parsed files 102 | *.sage.py 103 | 104 | # Environments 105 | .env 106 | .venv 107 | env/ 108 | venv/ 109 | ENV/ 110 | env.bak/ 111 | venv.bak/ 112 | 113 | # Spyder project settings 114 | .spyderproject 115 | .spyproject 116 | 117 | # Rope project settings 118 | .ropeproject 119 | 120 | # mkdocs documentation 121 | /site 122 | 123 | # mypy 124 | .mypy_cache/ 125 | .dmypy.json 126 | dmypy.json 127 | 128 | # Pyre type checker 129 | .pyre/ 130 | 131 | 132 | # Custom 133 | uci_data/* 134 | output/* 135 | *.npy 136 | *.pkl 137 | data/*.pdf 138 | data/*.png -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2020 Visual Learning Lab Heidelberg 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # HINT 2 | Code for the research paper "HINT: Hierarchical Invertible Neural Transport for Density Estimation and Bayesian Inference". 3 | 4 | ![Dense triangular Jacobian via recursive coupling](dense_jacobian.png) 5 | ![Structure of our recursive coupling block](recursive_block.png) 6 | ![Approximating a 2D shape with Fourier curve parameterization](fourier_curve.png) 7 | 8 | A pre-print is available at [arXiv](https://arxiv.org/abs/1905.10687), for now please cite as: 9 | ``` 10 | @misc{kruse2019hint, 11 | title = {HINT: Hierarchical Invertible Neural Transport for Density Estimation and Bayesian Inference}, 12 | author = {Jakob Kruse and Gianluca Detommaso and Robert Scheichl and Ullrich K\"othe}, 13 | year = {2019}, 14 | eprint = {1905.10687}, 15 | archivePrefix = {arXiv} 16 | } 17 | ``` 18 | 19 | 20 | ### Requirements 21 | 22 | In order to run the code, you will need the following: 23 | 24 | + `PyTorch` (>= v1.0.0) 25 | + `Python` (>= v3.7) 26 | + Packages `numpy`, `scipy`, `pandas`, `shapely` & `visdom` 27 | + [`FrEIA`](https://github.com/VLL-HD/FrEIA/) 28 | 29 | 30 | ### Structure 31 | 32 | There is one script for [training unconditional models](../master/train_unconditional.py) and one for [training conditional models](../master/train_conditional.py). 33 | These scripts expect a running `visdom` server for [visualization](../master/monitoring.py) and the import of a config file specifying the model and hyperparameters. 34 | Config files for all models used in the paper are supplied in the directory [configs](../master/configs). 35 | 36 | The recursive affine coupling block is implemented as a `FrEIA`-compatible module in [hint.py](../master/hint.py) and will be officially added to the framework in the future. 37 | 38 | Data sets for training and evaluation are managed with [an additional script](../master/data.py), but note that we maintain [another repository](https://github.com/VLL-HD/inn_toy_data) with the latest versions of all our toy data sets. The script also provides data loaders for three common UCI density estimation benchmarks, with download instructions. Code for the plots used in our paper is found in [plot_data.py](../master/plot_data.py). 39 | 40 | The [last script file](../master/rejection_sampling.py) deals with the rejection sampling baseline and systematic comparisons between the trained models. The first run of this script will take some time as the baseline is very inefficient, but results are stored and subsequent runs go much faster. 41 | -------------------------------------------------------------------------------- /configs/lens_shape/conditional_cinn_1.py: -------------------------------------------------------------------------------- 1 | from collections import namedtuple 2 | import torch 3 | 4 | from FrEIA.framework import * 5 | from FrEIA.modules import * 6 | 7 | from data import LensShapeModel as model 8 | from data import prepare_data_loaders 9 | n_parameters, n_observations = model.n_parameters, model.n_observations 10 | 11 | 12 | ###################################################### 13 | ### TRAINING CONFIGURATION AND HYPERPARAMETERS ### 14 | ###################################################### 15 | 16 | c = { 17 | # GENERAL STUFF 18 | 'suffix': f'{model.name}_conditional_cinn-1', # identifier for trained models and outputs 19 | 'device': 'cuda', # 'cuda' for GPU, 'cpu' for CPU 20 | 'interactive_visualization': True, # requires visdom package to be installed 21 | 22 | # DATA 23 | 'ndim_x': n_parameters, 24 | 'ndim_y': n_observations, 25 | 'ndim_z': n_parameters, 26 | 'data_model': model(), 27 | 'vis_y_target': (2.0, -1.0), 28 | 29 | # MODEL ARCHITECTURE 30 | 'n_blocks': 1, 31 | 'hidden_layer_sizes': 310 , # 217, # 200k 32 | 'init_scale': 0.005, 33 | 34 | # TRAINING HYPERPARAMETERS 35 | 'n_epochs': 50, # total number of epochs 36 | 'max_batches_per_epoch': 100, # iterations per epoch (or if training data exhausted) 37 | 'batch_size': 10000, 38 | 'n_test': 100000, 39 | 'n_train': 1000000, 40 | 41 | 'lr_init': 0.01, # initial learning rate 42 | 'pre_low_lr': 3, # number of epochs at the start with very low learning rate 43 | 'final_decay': 0.01, # fraction of the learning rate to reach at final epoch 44 | 'l2_weight_reg': 1.86e-05, # strength of the weight regularization 45 | 'adam_betas': (0.9, 0.95), 46 | } 47 | 48 | # DATA LOADERS 49 | train_loader, test_loader = prepare_data_loaders(c['data_model'], c['n_train'], c['n_test'], c['batch_size']) 50 | c['train_loader'] = train_loader 51 | c['test_loader'] = test_loader 52 | 53 | 54 | ############################## 55 | ### MODEL ARCHITECTURE ### 56 | ############################## 57 | 58 | nodes = [ConditionNode(c['ndim_y'], name='y')] 59 | nodes.append(InputNode(c['ndim_x'], name='x')) 60 | 61 | for i in range(c['n_blocks']): 62 | nodes.append(Node(nodes[-1], 63 | HouseholderPerm, 64 | {'fixed': False, 'n_reflections': c['ndim_x']}, 65 | name=f'perm_{i+1}')) 66 | nodes.append(Node(nodes[-1], 67 | AffineCoupling, 68 | {'F_class': F_fully_connected, 69 | 'F_args': {'internal_size': c['hidden_layer_sizes']}}, 70 | conditions=nodes[0], 71 | name=f'ac_{i+1}')) 72 | 73 | nodes.append(OutputNode(nodes[-1], name='z')) 74 | 75 | model = ReversibleGraphNet(nodes, verbose=False) 76 | model.to(c['device']) 77 | model.params_trainable = list(filter(lambda p: p.requires_grad, model.parameters())) 78 | 79 | 80 | def model_inverse(test_y, test_z): 81 | x_test = model([test_z], c=[test_y], rev=True) 82 | return x_test 83 | 84 | 85 | c['model'] = model 86 | c['model_inverse'] = model_inverse 87 | 88 | # create namedtuple from config dictionary 89 | c = namedtuple("Configuration",c.keys())(*c.values()) 90 | -------------------------------------------------------------------------------- /configs/lens_shape/conditional_cinn_2.py: -------------------------------------------------------------------------------- 1 | from collections import namedtuple 2 | import torch 3 | 4 | from FrEIA.framework import * 5 | from FrEIA.modules import * 6 | 7 | from data import LensShapeModel as model 8 | from data import prepare_data_loaders 9 | n_parameters, n_observations = model.n_parameters, model.n_observations 10 | 11 | 12 | ###################################################### 13 | ### TRAINING CONFIGURATION AND HYPERPARAMETERS ### 14 | ###################################################### 15 | 16 | c = { 17 | # GENERAL STUFF 18 | 'suffix': f'{model.name}_conditional_cinn-2', # identifier for trained models and outputs 19 | 'device': 'cuda', # 'cuda' for GPU, 'cpu' for CPU 20 | 'interactive_visualization': True, # requires visdom package to be installed 21 | 22 | # DATA 23 | 'ndim_x': n_parameters, 24 | 'ndim_y': n_observations, 25 | 'ndim_z': n_parameters, 26 | 'data_model': model(), 27 | 'vis_y_target': (2.0, -1.0), 28 | 29 | # MODEL ARCHITECTURE 30 | 'n_blocks': 2, 31 | 'hidden_layer_sizes': 217, # 151, # 200k 32 | 'init_scale': 0.005, 33 | 34 | # TRAINING HYPERPARAMETERS 35 | 'n_epochs': 50, # total number of epochs 36 | 'max_batches_per_epoch': 100, # iterations per epoch (or if training data exhausted) 37 | 'batch_size': 10000, 38 | 'n_test': 100000, 39 | 'n_train': 1000000, 40 | 41 | 'lr_init': 0.01, # initial learning rate 42 | 'pre_low_lr': 3, # number of epochs at the start with very low learning rate 43 | 'final_decay': 0.01, # fraction of the learning rate to reach at final epoch 44 | 'l2_weight_reg': 1.86e-05, # strength of the weight regularization 45 | 'adam_betas': (0.9, 0.95), 46 | } 47 | 48 | # DATA LOADERS 49 | train_loader, test_loader = prepare_data_loaders(c['data_model'], c['n_train'], c['n_test'], c['batch_size']) 50 | c['train_loader'] = train_loader 51 | c['test_loader'] = test_loader 52 | 53 | 54 | ############################## 55 | ### MODEL ARCHITECTURE ### 56 | ############################## 57 | 58 | nodes = [ConditionNode(c['ndim_y'], name='y')] 59 | nodes.append(InputNode(c['ndim_x'], name='x')) 60 | 61 | for i in range(c['n_blocks']): 62 | nodes.append(Node(nodes[-1], 63 | HouseholderPerm, 64 | {'fixed': False, 'n_reflections': c['ndim_x']}, 65 | name=f'perm_{i+1}')) 66 | nodes.append(Node(nodes[-1], 67 | AffineCoupling, 68 | {'F_class': F_fully_connected, 69 | 'F_args': {'internal_size': c['hidden_layer_sizes']}}, 70 | conditions=nodes[0], 71 | name=f'ac_{i+1}')) 72 | 73 | nodes.append(OutputNode(nodes[-1], name='z')) 74 | 75 | model = ReversibleGraphNet(nodes, verbose=False) 76 | model.to(c['device']) 77 | model.params_trainable = list(filter(lambda p: p.requires_grad, model.parameters())) 78 | 79 | 80 | def model_inverse(test_y, test_z): 81 | x_test = model([test_z], c=[test_y], rev=True) 82 | return x_test 83 | 84 | 85 | c['model'] = model 86 | c['model_inverse'] = model_inverse 87 | 88 | # create namedtuple from config dictionary 89 | c = namedtuple("Configuration",c.keys())(*c.values()) 90 | -------------------------------------------------------------------------------- /configs/lens_shape/conditional_cinn_4.py: -------------------------------------------------------------------------------- 1 | from collections import namedtuple 2 | import torch 3 | 4 | from FrEIA.framework import * 5 | from FrEIA.modules import * 6 | 7 | from data import LensShapeModel as model 8 | from data import prepare_data_loaders 9 | n_parameters, n_observations = model.n_parameters, model.n_observations 10 | 11 | 12 | ###################################################### 13 | ### TRAINING CONFIGURATION AND HYPERPARAMETERS ### 14 | ###################################################### 15 | 16 | c = { 17 | # GENERAL STUFF 18 | 'suffix': f'{model.name}_conditional_cinn-4', # identifier for trained models and outputs 19 | 'device': 'cuda', # 'cuda' for GPU, 'cpu' for CPU 20 | 'interactive_visualization': True, # requires visdom package to be installed 21 | 22 | # DATA 23 | 'ndim_x': n_parameters, 24 | 'ndim_y': n_observations, 25 | 'ndim_z': n_parameters, 26 | 'data_model': model(), 27 | 'vis_y_target': (2.0, -1.0), 28 | 29 | # MODEL ARCHITECTURE 30 | 'n_blocks': 4, 31 | 'hidden_layer_sizes': 151, # 400k 32 | 'init_scale': 0.005, 33 | 34 | # TRAINING HYPERPARAMETERS 35 | 'n_epochs': 50, # total number of epochs 36 | 'max_batches_per_epoch': 100, # iterations per epoch (or if training data exhausted) 37 | 'batch_size': 10000, 38 | 'n_test': 100000, 39 | 'n_train': 1000000, 40 | 41 | 'lr_init': 0.01, # initial learning rate 42 | 'pre_low_lr': 3, # number of epochs at the start with very low learning rate 43 | 'final_decay': 0.01, # fraction of the learning rate to reach at final epoch 44 | 'l2_weight_reg': 1.86e-05, # strength of the weight regularization 45 | 'adam_betas': (0.9, 0.95), 46 | } 47 | 48 | # DATA LOADERS 49 | train_loader, test_loader = prepare_data_loaders(c['data_model'], c['n_train'], c['n_test'], c['batch_size']) 50 | c['train_loader'] = train_loader 51 | c['test_loader'] = test_loader 52 | 53 | 54 | ############################## 55 | ### MODEL ARCHITECTURE ### 56 | ############################## 57 | 58 | nodes = [ConditionNode(c['ndim_y'], name='y')] 59 | nodes.append(InputNode(c['ndim_x'], name='x')) 60 | 61 | for i in range(c['n_blocks']): 62 | nodes.append(Node(nodes[-1], 63 | HouseholderPerm, 64 | {'fixed': False, 'n_reflections': c['ndim_x']}, 65 | name=f'perm_{i+1}')) 66 | nodes.append(Node(nodes[-1], 67 | AffineCoupling, 68 | {'F_class': F_fully_connected, 69 | 'F_args': {'internal_size': c['hidden_layer_sizes']}}, 70 | conditions=nodes[0], 71 | name=f'ac_{i+1}')) 72 | 73 | nodes.append(OutputNode(nodes[-1], name='z')) 74 | 75 | model = ReversibleGraphNet(nodes, verbose=False) 76 | model.to(c['device']) 77 | model.params_trainable = list(filter(lambda p: p.requires_grad, model.parameters())) 78 | 79 | 80 | def model_inverse(test_y, test_z): 81 | x_test = model([test_z], c=[test_y], rev=True) 82 | return x_test 83 | 84 | 85 | c['model'] = model 86 | c['model_inverse'] = model_inverse 87 | 88 | # create namedtuple from config dictionary 89 | c = namedtuple("Configuration",c.keys())(*c.values()) 90 | -------------------------------------------------------------------------------- /configs/lens_shape/conditional_cinn_8.py: -------------------------------------------------------------------------------- 1 | from collections import namedtuple 2 | import torch 3 | 4 | from FrEIA.framework import * 5 | from FrEIA.modules import * 6 | 7 | from data import LensShapeModel as model 8 | from data import prepare_data_loaders 9 | n_parameters, n_observations = model.n_parameters, model.n_observations 10 | 11 | 12 | ###################################################### 13 | ### TRAINING CONFIGURATION AND HYPERPARAMETERS ### 14 | ###################################################### 15 | 16 | c = { 17 | # GENERAL STUFF 18 | 'suffix': f'{model.name}_conditional_cinn-8', # identifier for trained models and outputs 19 | 'device': 'cuda', # 'cuda' for GPU, 'cpu' for CPU 20 | 'interactive_visualization': True, # requires visdom package to be installed 21 | 22 | # DATA 23 | 'ndim_x': n_parameters, 24 | 'ndim_y': n_observations, 25 | 'ndim_z': n_parameters, 26 | 'data_model': model(), 27 | 'vis_y_target': (2.0, -1.0), 28 | 29 | # MODEL ARCHITECTURE 30 | 'n_blocks': 8, 31 | 'hidden_layer_sizes': 105, # 400k 32 | 'init_scale': 0.005, 33 | 34 | # TRAINING HYPERPARAMETERS 35 | 'n_epochs': 50, # total number of epochs 36 | 'max_batches_per_epoch': 100, # iterations per epoch (or if training data exhausted) 37 | 'batch_size': 10000, 38 | 'n_test': 100000, 39 | 'n_train': 1000000, 40 | 41 | 'lr_init': 0.01, # initial learning rate 42 | 'pre_low_lr': 3, # number of epochs at the start with very low learning rate 43 | 'final_decay': 0.01, # fraction of the learning rate to reach at final epoch 44 | 'l2_weight_reg': 1.86e-05, # strength of the weight regularization 45 | 'adam_betas': (0.9, 0.95), 46 | } 47 | 48 | # DATA LOADERS 49 | train_loader, test_loader = prepare_data_loaders(c['data_model'], c['n_train'], c['n_test'], c['batch_size']) 50 | c['train_loader'] = train_loader 51 | c['test_loader'] = test_loader 52 | 53 | 54 | ############################## 55 | ### MODEL ARCHITECTURE ### 56 | ############################## 57 | 58 | nodes = [ConditionNode(c['ndim_y'], name='y')] 59 | nodes.append(InputNode(c['ndim_x'], name='x')) 60 | 61 | for i in range(c['n_blocks']): 62 | nodes.append(Node(nodes[-1], 63 | HouseholderPerm, 64 | {'fixed': False, 'n_reflections': c['ndim_x']}, 65 | name=f'perm_{i+1}')) 66 | nodes.append(Node(nodes[-1], 67 | AffineCoupling, 68 | {'F_class': F_fully_connected, 69 | 'F_args': {'internal_size': c['hidden_layer_sizes']}}, 70 | conditions=nodes[0], 71 | name=f'ac_{i+1}')) 72 | 73 | nodes.append(OutputNode(nodes[-1], name='z')) 74 | 75 | model = ReversibleGraphNet(nodes, verbose=False) 76 | model.to(c['device']) 77 | model.params_trainable = list(filter(lambda p: p.requires_grad, model.parameters())) 78 | 79 | 80 | def model_inverse(test_y, test_z): 81 | x_test = model([test_z], c=[test_y], rev=True) 82 | return x_test 83 | 84 | 85 | c['model'] = model 86 | c['model_inverse'] = model_inverse 87 | 88 | # create namedtuple from config dictionary 89 | c = namedtuple("Configuration",c.keys())(*c.values()) 90 | -------------------------------------------------------------------------------- /configs/lens_shape/conditional_hint_1_full.py: -------------------------------------------------------------------------------- 1 | from collections import namedtuple 2 | import torch 3 | 4 | from FrEIA.framework import * 5 | from FrEIA.modules import * 6 | 7 | from data import LensShapeModel as model 8 | from data import prepare_data_loaders 9 | n_parameters, n_observations = model.n_parameters, model.n_observations 10 | 11 | 12 | ###################################################### 13 | ### TRAINING CONFIGURATION AND HYPERPARAMETERS ### 14 | ###################################################### 15 | 16 | c = { 17 | # GENERAL STUFF 18 | 'suffix': f'{model.name}_conditional_hint-1-full', # identifier for trained models and outputs 19 | 'device': 'cuda', # 'cuda' for GPU, 'cpu' for CPU 20 | 'interactive_visualization': True, # requires visdom package to be installed 21 | 22 | # DATA 23 | 'ndim_x': n_parameters, 24 | 'ndim_y': n_observations, 25 | 'ndim_z': n_parameters + n_observations, 26 | 'data_model': model(), 27 | 'vis_y_target': (2.0, -1.0), 28 | 29 | # MODEL ARCHITECTURE 30 | 'n_blocks': 1, 31 | 'hidden_layer_sizes': 205, # 143, # 200k 32 | 'init_scale': 0.005, 33 | 34 | # TRAINING HYPERPARAMETERS 35 | 'n_epochs': 50, # total number of epochs 36 | 'max_batches_per_epoch': 100, # iterations per epoch (or if training data exhausted) 37 | 'batch_size': 10000, 38 | 'n_test': 100000, 39 | 'n_train': 1000000, 40 | 41 | 'lr_init': 0.01, # initial learning rate 42 | 'pre_low_lr': 3, # number of epochs at the start with very low learning rate 43 | 'final_decay': 0.01, # fraction of the learning rate to reach at final epoch 44 | 'l2_weight_reg': 1.86e-05, # strength of the weight regularization 45 | 'adam_betas': (0.9, 0.95), 46 | } 47 | 48 | # DATA LOADERS 49 | train_loader, test_loader = prepare_data_loaders(c['data_model'], c['n_train'], c['n_test'], c['batch_size']) 50 | c['train_loader'] = train_loader 51 | c['test_loader'] = test_loader 52 | 53 | 54 | ############################## 55 | ### MODEL ARCHITECTURE ### 56 | ############################## 57 | 58 | y_lane = [InputNode(c['ndim_y'], name='y')] 59 | x_lane = [InputNode(c['ndim_x'], name='x')] 60 | 61 | for i in range(c['n_blocks']): 62 | if i > 0: 63 | y_lane.append(Node(y_lane[-1], 64 | HouseholderPerm, 65 | {'fixed': False, 'n_reflections': c['ndim_y']}, 66 | name=f'perm_y_{i}')) 67 | x_lane.append(Node(x_lane[-1], 68 | HouseholderPerm, 69 | {'fixed': False, 'n_reflections': c['ndim_x']}, 70 | name=f'perm_x_{i}')) 71 | 72 | x_lane.append(Node(x_lane[-1], 73 | HierarchicalAffineCouplingBlock, 74 | {'c_internal': [c['hidden_layer_sizes'], c['hidden_layer_sizes']//2, c['hidden_layer_sizes']//4]}, 75 | name=f'hac_x_{i+1}')) 76 | 77 | # if i < c.n_blocks-1: 78 | x_lane.append(Node(x_lane[-1], 79 | ExternalAffineCoupling, 80 | {'F_class': F_fully_connected, 81 | 'F_args': {'internal_size': c['hidden_layer_sizes']}}, 82 | conditions=y_lane[-1], 83 | name=f'ac_y_to_x_{i+1}')) 84 | 85 | y_lane.append(Node(y_lane[-1], 86 | AffineCoupling, 87 | {'F_class': F_fully_connected, 88 | 'F_args': {'internal_size': c['hidden_layer_sizes']//4}}, 89 | name=f'ac_y_{i+1}')) 90 | 91 | y_lane.append(OutputNode(y_lane[-1], name='z_y')) 92 | x_lane.append(OutputNode(x_lane[-1], name='z_x')) 93 | 94 | model = ReversibleGraphNet(y_lane + x_lane, verbose=False) 95 | model.to(c['device']) 96 | model.params_trainable = list(filter(lambda p: p.requires_grad, model.parameters())) 97 | 98 | 99 | def model_inverse(test_y, test_z): 100 | z_y, z_x = model([test_y, torch.randn(test_y.shape[0], n_parameters).to(c.device)]) 101 | y_test, x_test = model([z_y, test_z], rev=True) 102 | return x_test 103 | 104 | def sample_joint(n_samples): 105 | return model([torch.randn(n_samples, n_observations).to(c.device), 106 | torch.randn(n_samples, n_parameters).to(c.device)], rev=True) 107 | 108 | def sample_conditional(y, z_x=None): 109 | if z_x is None: 110 | z_x = torch.randn(y.shape[0], n_parameters).to(c.device) 111 | z_y, _ = model([y, z_x]) 112 | y, x = model([z_y, z_x], rev=True) 113 | return x 114 | 115 | 116 | c['model'] = model 117 | c['model_inverse'] = model_inverse 118 | 119 | # create namedtuple from config dictionary 120 | c = namedtuple("Configuration", c.keys())(*c.values()) 121 | -------------------------------------------------------------------------------- /configs/lens_shape/unconditional_hint_1_full.py: -------------------------------------------------------------------------------- 1 | from collections import namedtuple 2 | import torch 3 | 4 | from FrEIA.framework import * 5 | from FrEIA.modules import * 6 | 7 | from data import LensShapeModel as model 8 | from data import prepare_data_loaders 9 | n_parameters = model.n_parameters 10 | 11 | 12 | ###################################################### 13 | ### TRAINING CONFIGURATION AND HYPERPARAMETERS ### 14 | ###################################################### 15 | 16 | c = { 17 | # GENERAL STUFF 18 | 'suffix': f'{model.name}_unconditional_hint-1-full', # identifier for trained models and outputs 19 | 'device': 'cuda', # 'cuda' for GPU, 'cpu' for CPU 20 | 'interactive_visualization': True, # requires visdom package to be installed 21 | 22 | # DATA 23 | 'ndim_x': n_parameters, 24 | 'ndim_y': 0, 25 | 'ndim_z': n_parameters, 26 | 'data_model': model(), 27 | 'vis_y_target': None, 28 | 29 | # MODEL ARCHITECTURE 30 | 'n_blocks': 1, 31 | 'hidden_layer_sizes': 139, # 100k 32 | 'init_scale': 0.005, 33 | 34 | # TRAINING HYPERPARAMETERS 35 | 'n_epochs': 50, # total number of epochs 36 | 'max_batches_per_epoch': 100, # iterations per epoch (or if training data exhausted) 37 | 'batch_size': 10000, 38 | 'n_test': 100000, 39 | 'n_train': 1000000, 40 | 41 | 'lr_init': 0.01, # initial learning rate 42 | 'pre_low_lr': 3, # number of epochs at the start with very low learning rate 43 | 'final_decay': 0.01, # fraction of the learning rate to reach at final epoch 44 | 'l2_weight_reg': 1.86e-05, # strength of the weight regularization 45 | 'adam_betas': (0.9, 0.95), 46 | } 47 | 48 | # DATA LOADERS 49 | train_loader, test_loader = prepare_data_loaders(c['data_model'], c['n_train'], c['n_test'], c['batch_size']) 50 | c['train_loader'] = train_loader 51 | c['test_loader'] = test_loader 52 | 53 | 54 | ############################## 55 | ### MODEL ARCHITECTURE ### 56 | ############################## 57 | 58 | x_lane = [InputNode(c['ndim_x'], name='x')] 59 | 60 | for i in range(c['n_blocks']): 61 | # if i > 0: 62 | x_lane.append(Node(x_lane[-1], 63 | HouseholderPerm, 64 | {'fixed': True, 'n_reflections': c['ndim_x']}, 65 | name=f'perm_{i}')) 66 | 67 | x_lane.append(Node(x_lane[-1], 68 | HierarchicalAffineCouplingBlock, 69 | {'c_internal': [c['hidden_layer_sizes'], c['hidden_layer_sizes']//2, c['hidden_layer_sizes']//4]}, 70 | name=f'hac_{i+1}')) 71 | 72 | x_lane.append(OutputNode(x_lane[-1], name='z')) 73 | 74 | model = ReversibleGraphNet(x_lane, verbose=False) 75 | model.to(c['device']) 76 | model.params_trainable = list(filter(lambda p: p.requires_grad, model.parameters())) 77 | 78 | 79 | def model_inverse(test_z): 80 | return model(test_z, rev=True) 81 | 82 | 83 | c['model'] = model 84 | c['model_inverse'] = model_inverse 85 | 86 | # create namedtuple from config dictionary 87 | c = namedtuple('Configuration', c.keys())(*c.values()) 88 | -------------------------------------------------------------------------------- /configs/lens_shape/unconditional_hint_2_full.py: -------------------------------------------------------------------------------- 1 | from collections import namedtuple 2 | import torch 3 | 4 | from FrEIA.framework import * 5 | from FrEIA.modules import * 6 | 7 | from data import LensShapeModel as model 8 | from data import prepare_data_loaders 9 | n_parameters = model.n_parameters 10 | 11 | 12 | ###################################################### 13 | ### TRAINING CONFIGURATION AND HYPERPARAMETERS ### 14 | ###################################################### 15 | 16 | c = { 17 | # GENERAL STUFF 18 | 'suffix': f'{model.name}_unconditional_hint-2-full', # identifier for trained models and outputs 19 | 'device': 'cuda', # 'cuda' for GPU, 'cpu' for CPU 20 | 'interactive_visualization': True, # requires visdom package to be installed 21 | 22 | # DATA 23 | 'ndim_x': n_parameters, 24 | 'ndim_y': 0, 25 | 'ndim_z': n_parameters, 26 | 'data_model': model(), 27 | 'vis_y_target': None, 28 | 29 | # MODEL ARCHITECTURE 30 | 'n_blocks': 2, 31 | 'hidden_layer_sizes': 95, # 100k 32 | 'init_scale': 0.005, 33 | 34 | # TRAINING HYPERPARAMETERS 35 | 'n_epochs': 50, # total number of epochs 36 | 'max_batches_per_epoch': 100, # iterations per epoch (or if training data exhausted) 37 | 'batch_size': 10000, 38 | 'n_test': 100000, 39 | 'n_train': 1000000, 40 | 41 | 'lr_init': 0.01, # initial learning rate 42 | 'pre_low_lr': 3, # number of epochs at the start with very low learning rate 43 | 'final_decay': 0.01, # fraction of the learning rate to reach at final epoch 44 | 'l2_weight_reg': 1.86e-05, # strength of the weight regularization 45 | 'adam_betas': (0.9, 0.95), 46 | } 47 | 48 | # DATA LOADERS 49 | train_loader, test_loader = prepare_data_loaders(c['data_model'], c['n_train'], c['n_test'], c['batch_size']) 50 | c['train_loader'] = train_loader 51 | c['test_loader'] = test_loader 52 | 53 | 54 | ############################## 55 | ### MODEL ARCHITECTURE ### 56 | ############################## 57 | 58 | x_lane = [InputNode(c['ndim_x'], name='x')] 59 | 60 | for i in range(c['n_blocks']): 61 | if i > 0: 62 | x_lane.append(Node(x_lane[-1], 63 | HouseholderPerm, 64 | {'fixed': False, 'n_reflections': c['ndim_x']}, 65 | name=f'perm_{i}')) 66 | 67 | x_lane.append(Node(x_lane[-1], 68 | HierarchicalAffineCouplingBlock, 69 | {'c_internal': [c['hidden_layer_sizes'], c['hidden_layer_sizes']//2, c['hidden_layer_sizes']//4]}, 70 | name=f'hac_{i+1}')) 71 | 72 | x_lane.append(OutputNode(x_lane[-1], name='z')) 73 | 74 | model = ReversibleGraphNet(x_lane, verbose=False) 75 | model.to(c['device']) 76 | model.params_trainable = list(filter(lambda p: p.requires_grad, model.parameters())) 77 | 78 | 79 | def model_inverse(test_z): 80 | return model(test_z, rev=True) 81 | 82 | 83 | c['model'] = model 84 | c['model_inverse'] = model_inverse 85 | 86 | # create namedtuple from config dictionary 87 | c = namedtuple('Configuration', c.keys())(*c.values()) 88 | -------------------------------------------------------------------------------- /configs/lens_shape/unconditional_inn_1.py: -------------------------------------------------------------------------------- 1 | from collections import namedtuple 2 | import torch 3 | 4 | from FrEIA.framework import * 5 | from FrEIA.modules import * 6 | 7 | from data import LensShapeModel as model 8 | from data import prepare_data_loaders 9 | n_parameters = model.n_parameters 10 | 11 | 12 | ###################################################### 13 | ### TRAINING CONFIGURATION AND HYPERPARAMETERS ### 14 | ###################################################### 15 | 16 | c = { 17 | # GENERAL STUFF 18 | 'suffix': f'{model.name}_unconditional_inn-1', # identifier for trained models and outputs 19 | 'device': 'cuda', # 'cuda' for GPU, 'cpu' for CPU 20 | 'interactive_visualization': True, # requires visdom package to be installed 21 | 22 | # DATA 23 | 'ndim_x': n_parameters, 24 | 'ndim_y': 0, 25 | 'ndim_z': n_parameters, 26 | 'data_model': model(), 27 | 'vis_y_target': None, 28 | 29 | # MODEL ARCHITECTURE 30 | 'n_blocks': 1, 31 | 'hidden_layer_sizes': 152, # 100k 32 | 'init_scale': 0.005, 33 | 34 | # TRAINING HYPERPARAMETERS 35 | 'n_epochs': 50, # total number of epochs 36 | 'max_batches_per_epoch': 100, # iterations per epoch (or if training data exhausted) 37 | 'batch_size': 10000, 38 | 'n_test': 100000, 39 | 'n_train': 1000000, 40 | 41 | 'lr_init': 0.01, # initial learning rate 42 | 'pre_low_lr': 3, # number of epochs at the start with very low learning rate 43 | 'final_decay': 0.01, # fraction of the learning rate to reach at final epoch 44 | 'l2_weight_reg': 1.86e-05, # strength of the weight regularization 45 | 'adam_betas': (0.9, 0.95), 46 | } 47 | 48 | # DATA LOADERS 49 | train_loader, test_loader = prepare_data_loaders(c['data_model'], c['n_train'], c['n_test'], c['batch_size']) 50 | c['train_loader'] = train_loader 51 | c['test_loader'] = test_loader 52 | 53 | 54 | ############################## 55 | ### MODEL ARCHITECTURE ### 56 | ############################## 57 | 58 | nodes = [InputNode(c['ndim_x'], name='x')] 59 | 60 | for i in range(c['n_blocks']): 61 | nodes.append(Node(nodes[-1], 62 | HouseholderPerm, 63 | {'fixed': False, 'n_reflections': c['ndim_x']}, 64 | name=f'perm_{i+1}')) 65 | nodes.append(Node(nodes[-1], 66 | AffineCoupling, 67 | {'F_class': F_fully_connected, 68 | 'F_args': {'internal_size': c['hidden_layer_sizes']}}, 69 | name=f'ac_{i+1}')) 70 | 71 | nodes.append(OutputNode(nodes[-1], name='z')) 72 | 73 | model = ReversibleGraphNet(nodes, verbose=False) 74 | model.to(c['device']) 75 | model.params_trainable = list(filter(lambda p: p.requires_grad, model.parameters())) 76 | 77 | def model_inverse(test_z): 78 | return model([test_z], rev=True) 79 | 80 | 81 | c['model'] = model 82 | c['model_inverse'] = model_inverse 83 | 84 | # create namedtuple from config dictionary 85 | c = namedtuple('Configuration', c.keys())(*c.values()) 86 | -------------------------------------------------------------------------------- /configs/lens_shape/unconditional_inn_2.py: -------------------------------------------------------------------------------- 1 | from collections import namedtuple 2 | import torch 3 | 4 | from FrEIA.framework import * 5 | from FrEIA.modules import * 6 | 7 | from data import LensShapeModel as model 8 | from data import prepare_data_loaders 9 | n_parameters = model.n_parameters 10 | 11 | 12 | ###################################################### 13 | ### TRAINING CONFIGURATION AND HYPERPARAMETERS ### 14 | ###################################################### 15 | 16 | c = { 17 | # GENERAL STUFF 18 | 'suffix': f'{model.name}_unconditional_inn-2', # identifier for trained models and outputs 19 | 'device': 'cuda', # 'cuda' for GPU, 'cpu' for CPU 20 | 'interactive_visualization': True, # requires visdom package to be installed 21 | 22 | # DATA 23 | 'ndim_x': n_parameters, 24 | 'ndim_y': 0, 25 | 'ndim_z': n_parameters, 26 | 'data_model': model(), 27 | 'vis_y_target': None, 28 | 29 | # MODEL ARCHITECTURE 30 | 'n_blocks': 2, 31 | 'hidden_layer_sizes': 106, # 100k 32 | 'init_scale': 0.005, 33 | 34 | # TRAINING HYPERPARAMETERS 35 | 'n_epochs': 50, # total number of epochs 36 | 'max_batches_per_epoch': 100, # iterations per epoch (or if training data exhausted) 37 | 'batch_size': 10000, 38 | 'n_test': 100000, 39 | 'n_train': 1000000, 40 | 41 | 'lr_init': 0.01, # initial learning rate 42 | 'pre_low_lr': 3, # number of epochs at the start with very low learning rate 43 | 'final_decay': 0.01, # fraction of the learning rate to reach at final epoch 44 | 'l2_weight_reg': 1.86e-05, # strength of the weight regularization 45 | 'adam_betas': (0.9, 0.95), 46 | } 47 | 48 | # DATA LOADERS 49 | train_loader, test_loader = prepare_data_loaders(c['data_model'], c['n_train'], c['n_test'], c['batch_size']) 50 | c['train_loader'] = train_loader 51 | c['test_loader'] = test_loader 52 | 53 | 54 | ############################## 55 | ### MODEL ARCHITECTURE ### 56 | ############################## 57 | 58 | nodes = [InputNode(c['ndim_x'], name='x')] 59 | 60 | for i in range(c['n_blocks']): 61 | nodes.append(Node(nodes[-1], 62 | HouseholderPerm, 63 | {'fixed': False, 'n_reflections': c['ndim_x']}, 64 | name=f'perm_{i+1}')) 65 | nodes.append(Node(nodes[-1], 66 | AffineCoupling, 67 | {'F_class': F_fully_connected, 68 | 'F_args': {'internal_size': c['hidden_layer_sizes']}}, 69 | name=f'ac_{i+1}')) 70 | 71 | nodes.append(OutputNode(nodes[-1], name='z')) 72 | 73 | model = ReversibleGraphNet(nodes, verbose=False) 74 | model.to(c['device']) 75 | model.params_trainable = list(filter(lambda p: p.requires_grad, model.parameters())) 76 | 77 | def model_inverse(test_z): 78 | return model([test_z], rev=True) 79 | 80 | 81 | c['model'] = model 82 | c['model_inverse'] = model_inverse 83 | 84 | # create namedtuple from config dictionary 85 | c = namedtuple('Configuration', c.keys())(*c.values()) 86 | -------------------------------------------------------------------------------- /configs/plus_shape/conditional_cinn_1.py: -------------------------------------------------------------------------------- 1 | from collections import namedtuple 2 | import torch 3 | 4 | from FrEIA.framework import * 5 | from FrEIA.modules import * 6 | 7 | from data import FourierCurveModel as model 8 | from data import prepare_data_loaders 9 | n_parameters, n_observations = model.n_parameters, model.n_observations 10 | 11 | 12 | ###################################################### 13 | ### TRAINING CONFIGURATION AND HYPERPARAMETERS ### 14 | ###################################################### 15 | 16 | c = { 17 | # GENERAL STUFF 18 | 'suffix': f'{model.name}_conditional_cinn-1', # identifier for trained models and outputs 19 | 'device': 'cuda', # 'cuda' for GPU, 'cpu' for CPU 20 | 'interactive_visualization': True, # requires visdom package to be installed 21 | 22 | # DATA 23 | 'ndim_x': n_parameters, 24 | 'ndim_y': n_observations, 25 | 'ndim_z': n_parameters, 26 | 'data_model': model(), 27 | 'vis_y_target': (0.52466008, 0.21816375, 2.29708147), 28 | 29 | # MODEL ARCHITECTURE 30 | 'n_blocks': 1, 31 | 'hidden_layer_sizes': 217, # 200k 32 | 'init_scale': 0.005, 33 | 34 | # TRAINING HYPERPARAMETERS 35 | 'n_epochs': 50, # total number of epochs 36 | 'max_batches_per_epoch': 100, # iterations per epoch (or if training data exhausted) 37 | 'batch_size': 10000, 38 | 'n_test': 100000, 39 | 'n_train': 1000000, 40 | 41 | 'lr_init': 0.01, # initial learning rate 42 | 'pre_low_lr': 3, # number of epochs at the start with very low learning rate 43 | 'final_decay': 0.01, # fraction of the learning rate to reach at final epoch 44 | 'l2_weight_reg': 1.86e-05, # strength of the weight regularization 45 | 'adam_betas': (0.9, 0.95), 46 | } 47 | 48 | # DATA LOADERS 49 | train_loader, test_loader = prepare_data_loaders(c['data_model'], c['n_train'], c['n_test'], c['batch_size']) 50 | c['train_loader'] = train_loader 51 | c['test_loader'] = test_loader 52 | 53 | 54 | ############################## 55 | ### MODEL ARCHITECTURE ### 56 | ############################## 57 | 58 | nodes = [ConditionNode(c['ndim_y'], name='y')] 59 | nodes.append(InputNode(c['ndim_x'], name='x')) 60 | 61 | for i in range(c['n_blocks']): 62 | nodes.append(Node(nodes[-1], 63 | HouseholderPerm, 64 | {'fixed': False, 'n_reflections': c['ndim_x']}, 65 | name=f'perm_{i+1}')) 66 | nodes.append(Node(nodes[-1], 67 | AffineCoupling, 68 | {'F_class': F_fully_connected, 69 | 'F_args': {'internal_size': c['hidden_layer_sizes']}}, 70 | conditions=nodes[0], 71 | name=f'ac_{i+1}')) 72 | 73 | nodes.append(OutputNode(nodes[-1], name='z')) 74 | 75 | model = ReversibleGraphNet(nodes, verbose=False) 76 | model.to(c['device']) 77 | model.params_trainable = list(filter(lambda p: p.requires_grad, model.parameters())) 78 | 79 | 80 | def model_inverse(test_y, test_z): 81 | x_test = model([test_z], c=[test_y], rev=True) 82 | return x_test 83 | 84 | 85 | c['model'] = model 86 | c['model_inverse'] = model_inverse 87 | 88 | # create namedtuple from config dictionary 89 | c = namedtuple("Configuration",c.keys())(*c.values()) 90 | -------------------------------------------------------------------------------- /configs/plus_shape/conditional_cinn_2.py: -------------------------------------------------------------------------------- 1 | from collections import namedtuple 2 | import torch 3 | 4 | from FrEIA.framework import * 5 | from FrEIA.modules import * 6 | 7 | from data import FourierCurveModel as model 8 | from data import prepare_data_loaders 9 | n_parameters, n_observations = model.n_parameters, model.n_observations 10 | 11 | 12 | ###################################################### 13 | ### TRAINING CONFIGURATION AND HYPERPARAMETERS ### 14 | ###################################################### 15 | 16 | c = { 17 | # GENERAL STUFF 18 | 'suffix': f'{model.name}_conditional_cinn-2', # identifier for trained models and outputs 19 | 'device': 'cuda', # 'cuda' for GPU, 'cpu' for CPU 20 | 'interactive_visualization': True, # requires visdom package to be installed 21 | 22 | # DATA 23 | 'ndim_x': n_parameters, 24 | 'ndim_y': n_observations, 25 | 'ndim_z': n_parameters, 26 | 'data_model': model(), 27 | 'vis_y_target': (0.52466008, 0.21816375, 2.29708147), 28 | 29 | # MODEL ARCHITECTURE 30 | 'n_blocks': 2, 31 | 'hidden_layer_sizes': 151, # 200k 32 | 'init_scale': 0.005, 33 | 34 | # TRAINING HYPERPARAMETERS 35 | 'n_epochs': 50, # total number of epochs 36 | 'max_batches_per_epoch': 100, # iterations per epoch (or if training data exhausted) 37 | 'batch_size': 10000, 38 | 'n_test': 100000, 39 | 'n_train': 1000000, 40 | 41 | 'lr_init': 0.01, # initial learning rate 42 | 'pre_low_lr': 3, # number of epochs at the start with very low learning rate 43 | 'final_decay': 0.01, # fraction of the learning rate to reach at final epoch 44 | 'l2_weight_reg': 1.86e-05, # strength of the weight regularization 45 | 'adam_betas': (0.9, 0.95), 46 | } 47 | 48 | # DATA LOADERS 49 | train_loader, test_loader = prepare_data_loaders(c['data_model'], c['n_train'], c['n_test'], c['batch_size']) 50 | c['train_loader'] = train_loader 51 | c['test_loader'] = test_loader 52 | 53 | 54 | ############################## 55 | ### MODEL ARCHITECTURE ### 56 | ############################## 57 | 58 | nodes = [ConditionNode(c['ndim_y'], name='y')] 59 | nodes.append(InputNode(c['ndim_x'], name='x')) 60 | 61 | for i in range(c['n_blocks']): 62 | nodes.append(Node(nodes[-1], 63 | HouseholderPerm, 64 | {'fixed': False, 'n_reflections': c['ndim_x']}, 65 | name=f'perm_{i+1}')) 66 | nodes.append(Node(nodes[-1], 67 | AffineCoupling, 68 | {'F_class': F_fully_connected, 69 | 'F_args': {'internal_size': c['hidden_layer_sizes']}}, 70 | conditions=nodes[0], 71 | name=f'ac_{i+1}')) 72 | 73 | nodes.append(OutputNode(nodes[-1], name='z')) 74 | 75 | model = ReversibleGraphNet(nodes, verbose=False) 76 | model.to(c['device']) 77 | model.params_trainable = list(filter(lambda p: p.requires_grad, model.parameters())) 78 | 79 | 80 | def model_inverse(test_y, test_z): 81 | x_test = model([test_z], c=[test_y], rev=True) 82 | return x_test 83 | 84 | 85 | c['model'] = model 86 | c['model_inverse'] = model_inverse 87 | 88 | # create namedtuple from config dictionary 89 | c = namedtuple("Configuration",c.keys())(*c.values()) 90 | -------------------------------------------------------------------------------- /configs/plus_shape/conditional_cinn_4.py: -------------------------------------------------------------------------------- 1 | from collections import namedtuple 2 | import torch 3 | 4 | from FrEIA.framework import * 5 | from FrEIA.modules import * 6 | 7 | from data import PlusShapeModel as model 8 | from data import prepare_data_loaders 9 | n_parameters, n_observations = model.n_parameters, model.n_observations 10 | 11 | 12 | ###################################################### 13 | ### TRAINING CONFIGURATION AND HYPERPARAMETERS ### 14 | ###################################################### 15 | 16 | c = { 17 | # GENERAL STUFF 18 | 'suffix': f'{model.name}_conditional_cinn-4', # identifier for trained models and outputs 19 | 'device': 'cuda', # 'cuda' for GPU, 'cpu' for CPU 20 | 'interactive_visualization': True, # requires visdom package to be installed 21 | 22 | # DATA 23 | 'ndim_x': n_parameters, 24 | 'ndim_y': n_observations, 25 | 'ndim_z': n_parameters, 26 | 'data_model': model(), 27 | 'vis_y_target': (0.75, 0.0, 1.0, 3.0), 28 | 29 | # MODEL ARCHITECTURE 30 | 'n_blocks': 4, 31 | 'hidden_layer_sizes': 472, # 4M 32 | 'init_scale': 0.005, 33 | 34 | # TRAINING HYPERPARAMETERS 35 | 'n_epochs': 50, # total number of epochs 36 | 'max_batches_per_epoch': 100, # iterations per epoch (or if training data exhausted) 37 | 'batch_size': 10000, 38 | 'n_test': 100000, 39 | 'n_train': 1000000, 40 | 41 | 'lr_init': 0.01, # initial learning rate 42 | 'pre_low_lr': 3, # number of epochs at the start with very low learning rate 43 | 'final_decay': 0.01, # fraction of the learning rate to reach at final epoch 44 | 'l2_weight_reg': 1.86e-05, # strength of the weight regularization 45 | 'adam_betas': (0.9, 0.95), 46 | } 47 | 48 | # DATA LOADERS 49 | train_loader, test_loader = prepare_data_loaders(c['data_model'], c['n_train'], c['n_test'], c['batch_size']) 50 | c['train_loader'] = train_loader 51 | c['test_loader'] = test_loader 52 | 53 | 54 | ############################## 55 | ### MODEL ARCHITECTURE ### 56 | ############################## 57 | 58 | nodes = [ConditionNode(c['ndim_y'], name='y')] 59 | nodes.append(InputNode(c['ndim_x'], name='x')) 60 | 61 | for i in range(c['n_blocks']): 62 | nodes.append(Node(nodes[-1], 63 | HouseholderPerm, 64 | {'fixed': False, 'n_reflections': c['ndim_x']}, 65 | name=f'perm_{i+1}')) 66 | nodes.append(Node(nodes[-1], 67 | AffineCoupling, 68 | {'F_class': F_fully_connected, 69 | 'F_args': {'internal_size': c['hidden_layer_sizes']}}, 70 | conditions=nodes[0], 71 | name=f'ac_{i+1}')) 72 | 73 | nodes.append(OutputNode(nodes[-1], name='z')) 74 | 75 | model = ReversibleGraphNet(nodes, verbose=False) 76 | model.to(c['device']) 77 | model.params_trainable = list(filter(lambda p: p.requires_grad, model.parameters())) 78 | 79 | 80 | def model_inverse(test_y, test_z): 81 | x_test = model([test_z], c=[test_y], rev=True) 82 | return x_test 83 | 84 | 85 | c['model'] = model 86 | c['model_inverse'] = model_inverse 87 | 88 | # create namedtuple from config dictionary 89 | c = namedtuple("Configuration",c.keys())(*c.values()) 90 | -------------------------------------------------------------------------------- /configs/plus_shape/conditional_cinn_8.py: -------------------------------------------------------------------------------- 1 | from collections import namedtuple 2 | import torch 3 | 4 | from FrEIA.framework import * 5 | from FrEIA.modules import * 6 | 7 | from data import PlusShapeModel as model 8 | from data import prepare_data_loaders 9 | n_parameters, n_observations = model.n_parameters, model.n_observations 10 | 11 | 12 | ###################################################### 13 | ### TRAINING CONFIGURATION AND HYPERPARAMETERS ### 14 | ###################################################### 15 | 16 | c = { 17 | # GENERAL STUFF 18 | 'suffix': f'{model.name}_conditional_cinn-8', # identifier for trained models and outputs 19 | 'device': 'cuda', # 'cuda' for GPU, 'cpu' for CPU 20 | 'interactive_visualization': True, # requires visdom package to be installed 21 | 22 | # DATA 23 | 'ndim_x': n_parameters, 24 | 'ndim_y': n_observations, 25 | 'ndim_z': n_parameters, 26 | 'data_model': model(), 27 | 'vis_y_target': (0.75, 0.0, 1.0, 3.0), 28 | 29 | # MODEL ARCHITECTURE 30 | 'n_blocks': 8, 31 | 'hidden_layer_sizes': 328, # 4M 32 | 'init_scale': 0.005, 33 | 34 | # TRAINING HYPERPARAMETERS 35 | 'n_epochs': 50, # total number of epochs 36 | 'max_batches_per_epoch': 100, # iterations per epoch (or if training data exhausted) 37 | 'batch_size': 10000, 38 | 'n_test': 100000, 39 | 'n_train': 1000000, 40 | 41 | 'lr_init': 0.01, # initial learning rate 42 | 'pre_low_lr': 3, # number of epochs at the start with very low learning rate 43 | 'final_decay': 0.01, # fraction of the learning rate to reach at final epoch 44 | 'l2_weight_reg': 1.86e-05, # strength of the weight regularization 45 | 'adam_betas': (0.9, 0.95), 46 | } 47 | 48 | # DATA LOADERS 49 | train_loader, test_loader = prepare_data_loaders(c['data_model'], c['n_train'], c['n_test'], c['batch_size']) 50 | c['train_loader'] = train_loader 51 | c['test_loader'] = test_loader 52 | 53 | 54 | ############################## 55 | ### MODEL ARCHITECTURE ### 56 | ############################## 57 | 58 | nodes = [ConditionNode(c['ndim_y'], name='y')] 59 | nodes.append(InputNode(c['ndim_x'], name='x')) 60 | 61 | for i in range(c['n_blocks']): 62 | nodes.append(Node(nodes[-1], 63 | HouseholderPerm, 64 | {'fixed': False, 'n_reflections': c['ndim_x']}, 65 | name=f'perm_{i+1}')) 66 | nodes.append(Node(nodes[-1], 67 | AffineCoupling, 68 | {'F_class': F_fully_connected, 69 | 'F_args': {'internal_size': c['hidden_layer_sizes']}}, 70 | conditions=nodes[0], 71 | name=f'ac_{i+1}')) 72 | 73 | nodes.append(OutputNode(nodes[-1], name='z')) 74 | 75 | model = ReversibleGraphNet(nodes, verbose=False) 76 | model.to(c['device']) 77 | model.params_trainable = list(filter(lambda p: p.requires_grad, model.parameters())) 78 | 79 | 80 | def model_inverse(test_y, test_z): 81 | x_test = model([test_z], c=[test_y], rev=True) 82 | return x_test 83 | 84 | 85 | c['model'] = model 86 | c['model_inverse'] = model_inverse 87 | 88 | # create namedtuple from config dictionary 89 | c = namedtuple("Configuration",c.keys())(*c.values()) 90 | -------------------------------------------------------------------------------- /configs/plus_shape/conditional_hint_4_full.py: -------------------------------------------------------------------------------- 1 | from collections import namedtuple 2 | import torch 3 | 4 | from FrEIA.framework import * 5 | from FrEIA.modules import * 6 | 7 | from data import PlusShapeModel as model 8 | from data import prepare_data_loaders 9 | n_parameters, n_observations = model.n_parameters, model.n_observations 10 | 11 | 12 | ###################################################### 13 | ### TRAINING CONFIGURATION AND HYPERPARAMETERS ### 14 | ###################################################### 15 | 16 | c = { 17 | # GENERAL STUFF 18 | 'suffix': f'{model.name}_conditional_hint-4-full', # identifier for trained models and outputs 19 | 'device': 'cuda', # 'cuda' for GPU, 'cpu' for CPU 20 | 'interactive_visualization': True, # requires visdom package to be installed 21 | 22 | # DATA 23 | 'ndim_x': n_parameters, 24 | 'ndim_y': n_observations, 25 | 'ndim_z': n_parameters + n_observations, 26 | 'data_model': model(), 27 | 'vis_y_target': (0.75, 0.0, 1.0, 3.0), 28 | 29 | # MODEL ARCHITECTURE 30 | 'n_blocks': 4, 31 | 'hidden_layer_sizes': 224, # 4M 32 | 'init_scale': 0.005, 33 | 34 | # TRAINING HYPERPARAMETERS 35 | 'n_epochs': 50, # total number of epochs 36 | 'max_batches_per_epoch': 100, # iterations per epoch (or if training data exhausted) 37 | 'batch_size': 10000, 38 | 'n_test': 100000, 39 | 'n_train': 1000000, 40 | 41 | 'lr_init': 0.01, # initial learning rate 42 | 'pre_low_lr': 3, # number of epochs at the start with very low learning rate 43 | 'final_decay': 0.01, # fraction of the learning rate to reach at final epoch 44 | 'l2_weight_reg': 1.86e-05, # strength of the weight regularization 45 | 'adam_betas': (0.9, 0.95), 46 | } 47 | 48 | # DATA LOADERS 49 | train_loader, test_loader = prepare_data_loaders(c['data_model'], c['n_train'], c['n_test'], c['batch_size']) 50 | c['train_loader'] = train_loader 51 | c['test_loader'] = test_loader 52 | 53 | 54 | ############################## 55 | ### MODEL ARCHITECTURE ### 56 | ############################## 57 | 58 | y_lane = [InputNode(c['ndim_y'], name='y')] 59 | x_lane = [InputNode(c['ndim_x'], name='x')] 60 | 61 | for i in range(c['n_blocks']): 62 | if i > 0: 63 | y_lane.append(Node(y_lane[-1], 64 | HouseholderPerm, 65 | {'fixed': True, 'n_reflections': c['ndim_y']}, 66 | name=f'perm_y_{i}')) 67 | x_lane.append(Node(x_lane[-1], 68 | HouseholderPerm, 69 | {'fixed': True, 'n_reflections': c['ndim_x']}, 70 | name=f'perm_x_{i}')) 71 | 72 | x_lane.append(Node(x_lane[-1], 73 | HierarchicalAffineCouplingBlock, 74 | {'c_internal': [c['hidden_layer_sizes'], c['hidden_layer_sizes']//2, c['hidden_layer_sizes']//4]}, 75 | name=f'hac_x_{i+1}')) 76 | 77 | # if i < c['n_blocks']-1: 78 | x_lane.append(Node(x_lane[-1], 79 | ExternalAffineCoupling, 80 | {'F_class': F_fully_connected, 81 | 'F_args': {'internal_size': c['hidden_layer_sizes']}}, 82 | conditions=y_lane[-1], 83 | name=f'ac_y_to_x_{i+1}')) 84 | 85 | y_lane.append(Node(y_lane[-1], 86 | AffineCoupling, 87 | {'F_class': F_fully_connected, 88 | 'F_args': {'internal_size': c['hidden_layer_sizes']//2}}, 89 | name=f'ac_y_{i+1}')) 90 | 91 | y_lane.append(OutputNode(y_lane[-1], name='z_y')) 92 | x_lane.append(OutputNode(x_lane[-1], name='z_x')) 93 | 94 | model = ReversibleGraphNet(y_lane + x_lane, verbose=False) 95 | model.to(c['device']) 96 | model.params_trainable = list(filter(lambda p: p.requires_grad, model.parameters())) 97 | 98 | 99 | def model_inverse(test_y, test_z): 100 | z_y, z_x = model([test_y, torch.randn(test_y.shape[0], n_parameters).to(c.device)]) 101 | y_test, x_test = model([z_y, test_z], rev=True) 102 | return x_test 103 | 104 | def sample_joint(n_samples): 105 | return model([torch.randn(n_samples, n_observations).to(device), 106 | torch.randn(n_samples, n_parameters).to(device)], rev=True) 107 | 108 | def sample_conditional(y, z_x=None): 109 | if z_x is None: 110 | z_x = torch.randn(y.shape[0], n_parameters).to(device) 111 | z_y, _ = model([y, z_x]) 112 | y, x = model([z_y, z_x], rev=True) 113 | return x 114 | 115 | 116 | c['model'] = model 117 | c['model_inverse'] = model_inverse 118 | 119 | # create namedtuple from config dictionary 120 | c = namedtuple("Configuration",c.keys())(*c.values()) 121 | -------------------------------------------------------------------------------- /configs/plus_shape/conditional_hint_8_full.py: -------------------------------------------------------------------------------- 1 | from collections import namedtuple 2 | import torch 3 | 4 | from FrEIA.framework import * 5 | from FrEIA.modules import * 6 | 7 | from data import PlusShapeModel as model 8 | from data import prepare_data_loaders 9 | n_parameters, n_observations = model.n_parameters, model.n_observations 10 | 11 | 12 | ###################################################### 13 | ### TRAINING CONFIGURATION AND HYPERPARAMETERS ### 14 | ###################################################### 15 | 16 | c = { 17 | # GENERAL STUFF 18 | 'suffix': f'{model.name}_conditional_hint-8-full', # identifier for trained models and outputs 19 | 'device': 'cuda', # 'cuda' for GPU, 'cpu' for CPU 20 | 'interactive_visualization': True, # requires visdom package to be installed 21 | 22 | # DATA 23 | 'ndim_x': n_parameters, 24 | 'ndim_y': n_observations, 25 | 'ndim_z': n_parameters + n_observations, 26 | 'data_model': model(), 27 | 'vis_y_target': (0.75, 0.0, 1.0, 3.0), 28 | 29 | # MODEL ARCHITECTURE 30 | 'n_blocks': 8, 31 | 'hidden_layer_sizes': 152, # 4M 32 | 'init_scale': 0.005, 33 | 34 | # TRAINING HYPERPARAMETERS 35 | 'n_epochs': 50, # total number of epochs 36 | 'max_batches_per_epoch': 100, # iterations per epoch (or if training data exhausted) 37 | 'batch_size': 10000, 38 | 'n_test': 100000, 39 | 'n_train': 1000000, 40 | 41 | 'lr_init': 0.01, # initial learning rate 42 | 'pre_low_lr': 3, # number of epochs at the start with very low learning rate 43 | 'final_decay': 0.01, # fraction of the learning rate to reach at final epoch 44 | 'l2_weight_reg': 1.86e-05, # strength of the weight regularization 45 | 'adam_betas': (0.9, 0.95), 46 | } 47 | 48 | # DATA LOADERS 49 | train_loader, test_loader = prepare_data_loaders(c['data_model'], c['n_train'], c['n_test'], c['batch_size']) 50 | c['train_loader'] = train_loader 51 | c['test_loader'] = test_loader 52 | 53 | 54 | ############################## 55 | ### MODEL ARCHITECTURE ### 56 | ############################## 57 | 58 | y_lane = [InputNode(c['ndim_y'], name='y')] 59 | x_lane = [InputNode(c['ndim_x'], name='x')] 60 | 61 | for i in range(c['n_blocks']): 62 | if i > 0: 63 | y_lane.append(Node(y_lane[-1], 64 | HouseholderPerm, 65 | {'fixed': True, 'n_reflections': c['ndim_y']}, 66 | name=f'perm_y_{i}')) 67 | x_lane.append(Node(x_lane[-1], 68 | HouseholderPerm, 69 | {'fixed': True, 'n_reflections': c['ndim_x']}, 70 | name=f'perm_x_{i}')) 71 | 72 | x_lane.append(Node(x_lane[-1], 73 | HierarchicalAffineCouplingBlock, 74 | {'c_internal': [c['hidden_layer_sizes'], c['hidden_layer_sizes']//2, c['hidden_layer_sizes']//4]}, 75 | name=f'hac_x_{i+1}')) 76 | 77 | # if i < c['n_blocks']-1: 78 | x_lane.append(Node(x_lane[-1], 79 | ExternalAffineCoupling, 80 | {'F_class': F_fully_connected, 81 | 'F_args': {'internal_size': c['hidden_layer_sizes']}}, 82 | conditions=y_lane[-1], 83 | name=f'ac_y_to_x_{i+1}')) 84 | 85 | y_lane.append(Node(y_lane[-1], 86 | AffineCoupling, 87 | {'F_class': F_fully_connected, 88 | 'F_args': {'internal_size': c['hidden_layer_sizes']//2}}, 89 | name=f'ac_y_{i+1}')) 90 | 91 | y_lane.append(OutputNode(y_lane[-1], name='z_y')) 92 | x_lane.append(OutputNode(x_lane[-1], name='z_x')) 93 | 94 | model = ReversibleGraphNet(y_lane + x_lane, verbose=False) 95 | model.to(c['device']) 96 | model.params_trainable = list(filter(lambda p: p.requires_grad, model.parameters())) 97 | 98 | 99 | def model_inverse(test_y, test_z): 100 | z_y, z_x = model([test_y, torch.randn(test_y.shape[0], n_parameters).to(c.device)]) 101 | y_test, x_test = model([z_y, test_z], rev=True) 102 | return x_test 103 | 104 | def sample_joint(n_samples): 105 | return model([torch.randn(n_samples, n_observations).to(device), 106 | torch.randn(n_samples, n_parameters).to(device)], rev=True) 107 | 108 | def sample_conditional(y, z_x=None): 109 | if z_x is None: 110 | z_x = torch.randn(y.shape[0], n_parameters).to(device) 111 | z_y, _ = model([y, z_x]) 112 | y, x = model([z_y, z_x], rev=True) 113 | return x 114 | 115 | 116 | c['model'] = model 117 | c['model_inverse'] = model_inverse 118 | 119 | # create namedtuple from config dictionary 120 | c = namedtuple("Configuration",c.keys())(*c.values()) 121 | -------------------------------------------------------------------------------- /configs/plus_shape/conditional_recursive_cinn_4.py: -------------------------------------------------------------------------------- 1 | from collections import namedtuple 2 | import torch 3 | 4 | from FrEIA.framework import * 5 | from FrEIA.modules import * 6 | 7 | from data import PlusShapeModel as model 8 | from data import prepare_data_loaders 9 | n_parameters, n_observations = model.n_parameters, model.n_observations 10 | 11 | 12 | ###################################################### 13 | ### TRAINING CONFIGURATION AND HYPERPARAMETERS ### 14 | ###################################################### 15 | 16 | c = { 17 | # GENERAL STUFF 18 | 'suffix': f'{model.name}_conditional_recursive_cinn-4', # identifier for trained models and outputs 19 | 'device': 'cuda', # 'cuda' for GPU, 'cpu' for CPU 20 | 'interactive_visualization': True, # requires visdom package to be installed 21 | 22 | # DATA 23 | 'ndim_x': n_parameters, 24 | 'ndim_y': n_observations, 25 | 'ndim_z': n_parameters, 26 | 'data_model': model(), 27 | 'vis_y_target': (0.75, 0.0, 1.0, 3.0), 28 | 29 | # MODEL ARCHITECTURE 30 | 'n_blocks': 4, 31 | 'hidden_layer_sizes': 267, # 4M 32 | 'init_scale': 0.005, 33 | 34 | # TRAINING HYPERPARAMETERS 35 | 'n_epochs': 50, # total number of epochs 36 | 'max_batches_per_epoch': 100, # iterations per epoch (or if training data exhausted) 37 | 'batch_size': 10000, 38 | 'n_test': 100000, 39 | 'n_train': 1000000, 40 | 41 | 'lr_init': 0.01, # initial learning rate 42 | 'pre_low_lr': 3, # number of epochs at the start with very low learning rate 43 | 'final_decay': 0.01, # fraction of the learning rate to reach at final epoch 44 | 'l2_weight_reg': 1.86e-05, # strength of the weight regularization 45 | 'adam_betas': (0.9, 0.95), 46 | } 47 | 48 | # DATA LOADERS 49 | train_loader, test_loader = prepare_data_loaders(c['data_model'], c['n_train'], c['n_test'], c['batch_size']) 50 | c['train_loader'] = train_loader 51 | c['test_loader'] = test_loader 52 | 53 | 54 | ############################## 55 | ### MODEL ARCHITECTURE ### 56 | ############################## 57 | 58 | nodes = [ConditionNode(c['ndim_y'], name='y')] 59 | nodes.append(InputNode(c['ndim_x'], name='x')) 60 | 61 | for i in range(c['n_blocks']): 62 | nodes.append(Node(nodes[-1], 63 | HouseholderPerm, 64 | {'fixed': True, 'n_reflections': c['ndim_x']}, 65 | name=f'perm_{i+1}')) 66 | nodes.append(Node(nodes[-1], 67 | HierarchicalAffineCouplingBlock, 68 | {'c_internal': [c['hidden_layer_sizes'], c['hidden_layer_sizes']//2, c['hidden_layer_sizes']//4]}, 69 | conditions=nodes[0], 70 | name=f'ac_{i+1}')) 71 | 72 | nodes.append(OutputNode(nodes[-1], name='z')) 73 | 74 | model = ReversibleGraphNet(nodes, verbose=False) 75 | model.to(c['device']) 76 | model.params_trainable = list(filter(lambda p: p.requires_grad, model.parameters())) 77 | 78 | 79 | def model_inverse(test_y, test_z): 80 | x_test = model([test_z], c=[test_y], rev=True) 81 | return x_test 82 | 83 | 84 | c['model'] = model 85 | c['model_inverse'] = model_inverse 86 | 87 | # create namedtuple from config dictionary 88 | c = namedtuple("Configuration",c.keys())(*c.values()) 89 | -------------------------------------------------------------------------------- /configs/plus_shape/unconditional_hint_16_0_big.py: -------------------------------------------------------------------------------- 1 | from collections import namedtuple 2 | import torch 3 | 4 | from FrEIA.framework import * 5 | from FrEIA.modules import * 6 | 7 | from data import PlusShapeModel as model 8 | from data import prepare_data_loaders 9 | n_parameters = model.n_parameters 10 | 11 | 12 | ###################################################### 13 | ### TRAINING CONFIGURATION AND HYPERPARAMETERS ### 14 | ###################################################### 15 | 16 | c = { 17 | # GENERAL STUFF 18 | 'suffix': f'{model.name}_unconditional_hint-16-0-big', # identifier for trained models and outputs 19 | 'device': 'cuda', # 'cuda' for GPU, 'cpu' for CPU 20 | 'interactive_visualization': True, # requires visdom package to be installed 21 | 22 | # DATA 23 | 'ndim_x': n_parameters, 24 | 'ndim_y': 0, 25 | 'ndim_z': n_parameters, 26 | 'data_model': model(), 27 | 'vis_y_target': None, 28 | 29 | # MODEL ARCHITECTURE 30 | 'n_blocks': 16, 31 | 'hidden_layer_sizes': 512, 32 | 'init_scale': 0.005, 33 | 'recursion_depth': 0, 34 | 35 | # TRAINING HYPERPARAMETERS 36 | 'n_epochs': 50, # total number of epochs 37 | 'max_batches_per_epoch': 100, # iterations per epoch (or if training data exhausted) 38 | 'batch_size': 10000, 39 | 'n_test': 100000, 40 | 'n_train': 1000000, 41 | 42 | 'lr_init': 0.01, # initial learning rate 43 | 'pre_low_lr': 3, # number of epochs at the start with very low learning rate 44 | 'final_decay': 0.01, # fraction of the learning rate to reach at final epoch 45 | 'l2_weight_reg': 1.86e-05, # strength of the weight regularization 46 | 'adam_betas': (0.9, 0.95), 47 | } 48 | 49 | # DATA LOADERS 50 | train_loader, test_loader = prepare_data_loaders(c['data_model'], c['n_train'], c['n_test'], c['batch_size']) 51 | c['train_loader'] = train_loader 52 | c['test_loader'] = test_loader 53 | 54 | 55 | ############################## 56 | ### MODEL ARCHITECTURE ### 57 | ############################## 58 | 59 | x_lane = [InputNode(c['ndim_x'], name='x')] 60 | 61 | for i in range(c['n_blocks']): 62 | if i > 0: 63 | x_lane.append(Node(x_lane[-1], 64 | HouseholderPerm, 65 | {'fixed': False, 'n_reflections': c['ndim_x']}, 66 | name=f'perm_{i}')) 67 | 68 | x_lane.append(Node(x_lane[-1], 69 | HierarchicalAffineCouplingBlock, 70 | {'c_internal': [c['hidden_layer_sizes'], c['hidden_layer_sizes']//2, c['hidden_layer_sizes']//4, c['hidden_layer_sizes']//8], 71 | 'max_splits': c['recursion_depth']}, 72 | name=f'hac_{i+1}')) 73 | 74 | x_lane.append(OutputNode(x_lane[-1], name='z')) 75 | 76 | model = ReversibleGraphNet(x_lane, verbose=False) 77 | model.to(c['device']) 78 | model.params_trainable = list(filter(lambda p: p.requires_grad, model.parameters())) 79 | 80 | 81 | def model_inverse(test_z): 82 | return model(test_z, rev=True) 83 | 84 | 85 | c['model'] = model 86 | c['model_inverse'] = model_inverse 87 | 88 | # create namedtuple from config dictionary 89 | c = namedtuple("Configuration",c.keys())(*c.values()) 90 | -------------------------------------------------------------------------------- /configs/plus_shape/unconditional_hint_16_0_small.py: -------------------------------------------------------------------------------- 1 | from collections import namedtuple 2 | import torch 3 | 4 | from FrEIA.framework import * 5 | from FrEIA.modules import * 6 | 7 | from data import PlusShapeModel as model 8 | from data import prepare_data_loaders 9 | n_parameters = model.n_parameters 10 | 11 | 12 | ###################################################### 13 | ### TRAINING CONFIGURATION AND HYPERPARAMETERS ### 14 | ###################################################### 15 | 16 | c = { 17 | # GENERAL STUFF 18 | 'suffix': f'{model.name}_unconditional_hint-16-0-small', # identifier for trained models and outputs 19 | 'device': 'cuda', # 'cuda' for GPU, 'cpu' for CPU 20 | 'interactive_visualization': True, # requires visdom package to be installed 21 | 22 | # DATA 23 | 'ndim_x': n_parameters, 24 | 'ndim_y': 0, 25 | 'ndim_z': n_parameters, 26 | 'data_model': model(), 27 | 'vis_y_target': None, 28 | 29 | # MODEL ARCHITECTURE 30 | 'n_blocks': 16, 31 | 'hidden_layer_sizes': 128, 32 | 'init_scale': 0.005, 33 | 'recursion_depth': 0, 34 | 35 | # TRAINING HYPERPARAMETERS 36 | 'n_epochs': 50, # total number of epochs 37 | 'max_batches_per_epoch': 100, # iterations per epoch (or if training data exhausted) 38 | 'batch_size': 10000, 39 | 'n_test': 100000, 40 | 'n_train': 1000000, 41 | 42 | 'lr_init': 0.01, # initial learning rate 43 | 'pre_low_lr': 3, # number of epochs at the start with very low learning rate 44 | 'final_decay': 0.01, # fraction of the learning rate to reach at final epoch 45 | 'l2_weight_reg': 1.86e-05, # strength of the weight regularization 46 | 'adam_betas': (0.9, 0.95), 47 | } 48 | 49 | # DATA LOADERS 50 | train_loader, test_loader = prepare_data_loaders(c['data_model'], c['n_train'], c['n_test'], c['batch_size']) 51 | c['train_loader'] = train_loader 52 | c['test_loader'] = test_loader 53 | 54 | 55 | ############################## 56 | ### MODEL ARCHITECTURE ### 57 | ############################## 58 | 59 | x_lane = [InputNode(c['ndim_x'], name='x')] 60 | 61 | for i in range(c['n_blocks']): 62 | if i > 0: 63 | x_lane.append(Node(x_lane[-1], 64 | HouseholderPerm, 65 | {'fixed': False, 'n_reflections': c['ndim_x']}, 66 | name=f'perm_{i}')) 67 | 68 | x_lane.append(Node(x_lane[-1], 69 | HierarchicalAffineCouplingBlock, 70 | {'c_internal': [c['hidden_layer_sizes'], c['hidden_layer_sizes']//2, c['hidden_layer_sizes']//4, c['hidden_layer_sizes']//8], 71 | 'max_splits': c['recursion_depth']}, 72 | name=f'hac_{i+1}')) 73 | 74 | x_lane.append(OutputNode(x_lane[-1], name='z')) 75 | 76 | model = ReversibleGraphNet(x_lane, verbose=False) 77 | model.to(c['device']) 78 | model.params_trainable = list(filter(lambda p: p.requires_grad, model.parameters())) 79 | 80 | 81 | def model_inverse(test_z): 82 | return model(test_z, rev=True) 83 | 84 | 85 | c['model'] = model 86 | c['model_inverse'] = model_inverse 87 | 88 | # create namedtuple from config dictionary 89 | c = namedtuple("Configuration",c.keys())(*c.values()) 90 | -------------------------------------------------------------------------------- /configs/plus_shape/unconditional_hint_16_1.py: -------------------------------------------------------------------------------- 1 | from collections import namedtuple 2 | import torch 3 | 4 | from FrEIA.framework import * 5 | from FrEIA.modules import * 6 | 7 | from data import PlusShapeModel as model 8 | from data import prepare_data_loaders 9 | n_parameters = model.n_parameters 10 | 11 | 12 | ###################################################### 13 | ### TRAINING CONFIGURATION AND HYPERPARAMETERS ### 14 | ###################################################### 15 | 16 | c = { 17 | # GENERAL STUFF 18 | 'suffix': f'{model.name}_unconditional_hint-16-1', # identifier for trained models and outputs 19 | 'device': 'cuda', # 'cuda' for GPU, 'cpu' for CPU 20 | 'interactive_visualization': True, # requires visdom package to be installed 21 | 22 | # DATA 23 | 'ndim_x': n_parameters, 24 | 'ndim_y': 0, 25 | 'ndim_z': n_parameters, 26 | 'data_model': model(), 27 | 'vis_y_target': None, 28 | 29 | # MODEL ARCHITECTURE 30 | 'n_blocks': 16, 31 | 'hidden_layer_sizes': 159, # 2M 32 | 'init_scale': 0.005, 33 | 34 | # TRAINING HYPERPARAMETERS 35 | 'n_epochs': 50, # total number of epochs 36 | 'max_batches_per_epoch': 100, # iterations per epoch (or if training data exhausted) 37 | 'batch_size': 10000, 38 | 'n_test': 100000, 39 | 'n_train': 1000000, 40 | 41 | 'lr_init': 0.01, # initial learning rate 42 | 'pre_low_lr': 3, # number of epochs at the start with very low learning rate 43 | 'final_decay': 0.01, # fraction of the learning rate to reach at final epoch 44 | 'l2_weight_reg': 1.86e-05, # strength of the weight regularization 45 | 'adam_betas': (0.9, 0.95), 46 | } 47 | 48 | # DATA LOADERS 49 | train_loader, test_loader = prepare_data_loaders(c['data_model'], c['n_train'], c['n_test'], c['batch_size']) 50 | c['train_loader'] = train_loader 51 | c['test_loader'] = test_loader 52 | 53 | 54 | ############################## 55 | ### MODEL ARCHITECTURE ### 56 | ############################## 57 | 58 | x_lane = [InputNode(c['ndim_x'], name='x')] 59 | 60 | for i in range(c['n_blocks']): 61 | if i > 0: 62 | x_lane.append(Node(x_lane[-1], 63 | HouseholderPerm, 64 | {'fixed': True, 'n_reflections': c['ndim_x']}, 65 | name=f'perm_{i}')) 66 | 67 | x_lane.append(Node(x_lane[-1], 68 | HierarchicalAffineCouplingBlock, 69 | {'c_internal': [c['hidden_layer_sizes'], c['hidden_layer_sizes']//2], 70 | 'max_splits': 1}, 71 | name=f'hac_{i+1}')) 72 | 73 | x_lane.append(OutputNode(x_lane[-1], name='z')) 74 | 75 | model = ReversibleGraphNet(x_lane, verbose=False) 76 | model.to(c['device']) 77 | model.params_trainable = list(filter(lambda p: p.requires_grad, model.parameters())) 78 | 79 | 80 | def model_inverse(test_z): 81 | return model(test_z, rev=True) 82 | 83 | 84 | c['model'] = model 85 | c['model_inverse'] = model_inverse 86 | 87 | # create namedtuple from config dictionary 88 | c = namedtuple("Configuration",c.keys())(*c.values()) 89 | -------------------------------------------------------------------------------- /configs/plus_shape/unconditional_hint_16_1_big.py: -------------------------------------------------------------------------------- 1 | from collections import namedtuple 2 | import torch 3 | 4 | from FrEIA.framework import * 5 | from FrEIA.modules import * 6 | 7 | from data import PlusShapeModel as model 8 | from data import prepare_data_loaders 9 | n_parameters = model.n_parameters 10 | 11 | 12 | ###################################################### 13 | ### TRAINING CONFIGURATION AND HYPERPARAMETERS ### 14 | ###################################################### 15 | 16 | c = { 17 | # GENERAL STUFF 18 | 'suffix': f'{model.name}_unconditional_hint-16-1-big', # identifier for trained models and outputs 19 | 'device': 'cuda', # 'cuda' for GPU, 'cpu' for CPU 20 | 'interactive_visualization': True, # requires visdom package to be installed 21 | 22 | # DATA 23 | 'ndim_x': n_parameters, 24 | 'ndim_y': 0, 25 | 'ndim_z': n_parameters, 26 | 'data_model': model(), 27 | 'vis_y_target': None, 28 | 29 | # MODEL ARCHITECTURE 30 | 'n_blocks': 16, 31 | 'hidden_layer_sizes': 512, 32 | 'init_scale': 0.005, 33 | 'recursion_depth': 1, 34 | 35 | # TRAINING HYPERPARAMETERS 36 | 'n_epochs': 50, # total number of epochs 37 | 'max_batches_per_epoch': 100, # iterations per epoch (or if training data exhausted) 38 | 'batch_size': 10000, 39 | 'n_test': 100000, 40 | 'n_train': 1000000, 41 | 42 | 'lr_init': 0.001, # initial learning rate 43 | 'pre_low_lr': 3, # number of epochs at the start with very low learning rate 44 | 'final_decay': 0.01, # fraction of the learning rate to reach at final epoch 45 | 'l2_weight_reg': 1.86e-05, # strength of the weight regularization 46 | 'adam_betas': (0.9, 0.95), 47 | } 48 | 49 | # DATA LOADERS 50 | train_loader, test_loader = prepare_data_loaders(c['data_model'], c['n_train'], c['n_test'], c['batch_size']) 51 | c['train_loader'] = train_loader 52 | c['test_loader'] = test_loader 53 | 54 | 55 | ############################## 56 | ### MODEL ARCHITECTURE ### 57 | ############################## 58 | 59 | x_lane = [InputNode(c['ndim_x'], name='x')] 60 | 61 | for i in range(c['n_blocks']): 62 | if i > 0: 63 | x_lane.append(Node(x_lane[-1], 64 | HouseholderPerm, 65 | {'fixed': False, 'n_reflections': c['ndim_x']}, 66 | name=f'perm_{i}')) 67 | 68 | x_lane.append(Node(x_lane[-1], 69 | HierarchicalAffineCouplingBlock, 70 | {'c_internal': [c['hidden_layer_sizes'], c['hidden_layer_sizes']//2, c['hidden_layer_sizes']//4, c['hidden_layer_sizes']//8], 71 | 'max_splits': c['recursion_depth']}, 72 | name=f'hac_{i+1}')) 73 | 74 | x_lane.append(OutputNode(x_lane[-1], name='z')) 75 | 76 | model = ReversibleGraphNet(x_lane, verbose=False) 77 | model.to(c['device']) 78 | model.params_trainable = list(filter(lambda p: p.requires_grad, model.parameters())) 79 | 80 | 81 | def model_inverse(test_z): 82 | return model(test_z, rev=True) 83 | 84 | 85 | c['model'] = model 86 | c['model_inverse'] = model_inverse 87 | 88 | # create namedtuple from config dictionary 89 | c = namedtuple("Configuration",c.keys())(*c.values()) 90 | -------------------------------------------------------------------------------- /configs/plus_shape/unconditional_hint_16_1_small.py: -------------------------------------------------------------------------------- 1 | from collections import namedtuple 2 | import torch 3 | 4 | from FrEIA.framework import * 5 | from FrEIA.modules import * 6 | 7 | from data import PlusShapeModel as model 8 | from data import prepare_data_loaders 9 | n_parameters = model.n_parameters 10 | 11 | 12 | ###################################################### 13 | ### TRAINING CONFIGURATION AND HYPERPARAMETERS ### 14 | ###################################################### 15 | 16 | c = { 17 | # GENERAL STUFF 18 | 'suffix': f'{model.name}_unconditional_hint-16-1-small', # identifier for trained models and outputs 19 | 'device': 'cuda', # 'cuda' for GPU, 'cpu' for CPU 20 | 'interactive_visualization': True, # requires visdom package to be installed 21 | 22 | # DATA 23 | 'ndim_x': n_parameters, 24 | 'ndim_y': 0, 25 | 'ndim_z': n_parameters, 26 | 'data_model': model(), 27 | 'vis_y_target': None, 28 | 29 | # MODEL ARCHITECTURE 30 | 'n_blocks': 16, 31 | 'hidden_layer_sizes': 128, 32 | 'init_scale': 0.005, 33 | 'recursion_depth': 1, 34 | 35 | # TRAINING HYPERPARAMETERS 36 | 'n_epochs': 50, # total number of epochs 37 | 'max_batches_per_epoch': 100, # iterations per epoch (or if training data exhausted) 38 | 'batch_size': 10000, 39 | 'n_test': 100000, 40 | 'n_train': 1000000, 41 | 42 | 'lr_init': 0.001, # initial learning rate 43 | 'pre_low_lr': 3, # number of epochs at the start with very low learning rate 44 | 'final_decay': 0.01, # fraction of the learning rate to reach at final epoch 45 | 'l2_weight_reg': 1.86e-05, # strength of the weight regularization 46 | 'adam_betas': (0.9, 0.95), 47 | } 48 | 49 | # DATA LOADERS 50 | train_loader, test_loader = prepare_data_loaders(c['data_model'], c['n_train'], c['n_test'], c['batch_size']) 51 | c['train_loader'] = train_loader 52 | c['test_loader'] = test_loader 53 | 54 | 55 | ############################## 56 | ### MODEL ARCHITECTURE ### 57 | ############################## 58 | 59 | x_lane = [InputNode(c['ndim_x'], name='x')] 60 | 61 | for i in range(c['n_blocks']): 62 | if i > 0: 63 | x_lane.append(Node(x_lane[-1], 64 | HouseholderPerm, 65 | {'fixed': False, 'n_reflections': c['ndim_x']}, 66 | name=f'perm_{i}')) 67 | 68 | x_lane.append(Node(x_lane[-1], 69 | HierarchicalAffineCouplingBlock, 70 | {'c_internal': [c['hidden_layer_sizes'], c['hidden_layer_sizes']//2, c['hidden_layer_sizes']//4, c['hidden_layer_sizes']//8], 71 | 'max_splits': c['recursion_depth']}, 72 | name=f'hac_{i+1}')) 73 | 74 | x_lane.append(OutputNode(x_lane[-1], name='z')) 75 | 76 | model = ReversibleGraphNet(x_lane, verbose=False) 77 | model.to(c['device']) 78 | model.params_trainable = list(filter(lambda p: p.requires_grad, model.parameters())) 79 | 80 | 81 | def model_inverse(test_z): 82 | return model(test_z, rev=True) 83 | 84 | 85 | c['model'] = model 86 | c['model_inverse'] = model_inverse 87 | 88 | # create namedtuple from config dictionary 89 | c = namedtuple("Configuration",c.keys())(*c.values()) 90 | -------------------------------------------------------------------------------- /configs/plus_shape/unconditional_hint_1_full.py: -------------------------------------------------------------------------------- 1 | from collections import namedtuple 2 | import torch 3 | 4 | from FrEIA.framework import * 5 | from FrEIA.modules import * 6 | 7 | from data import PlusShapeModel as model 8 | from data import prepare_data_loaders 9 | n_parameters = model.n_parameters 10 | 11 | 12 | ###################################################### 13 | ### TRAINING CONFIGURATION AND HYPERPARAMETERS ### 14 | ###################################################### 15 | 16 | c = { 17 | # GENERAL STUFF 18 | 'suffix': f'{model.name}_unconditional_hint-1-full', # identifier for trained models and outputs 19 | 'device': 'cuda', # 'cuda' for GPU, 'cpu' for CPU 20 | 'interactive_visualization': True, # requires visdom package to be installed 21 | 22 | # DATA 23 | 'ndim_x': n_parameters, 24 | 'ndim_y': 0, 25 | 'ndim_z': n_parameters, 26 | 'data_model': model(), 27 | 'vis_y_target': None, 28 | 29 | # MODEL ARCHITECTURE 30 | 'n_blocks': 1, 31 | 'hidden_layer_sizes': 110, # 200k 32 | 'init_scale': 0.005, 33 | 34 | # TRAINING HYPERPARAMETERS 35 | 'n_epochs': 50, # total number of epochs 36 | 'max_batches_per_epoch': 100, # iterations per epoch (or if training data exhausted) 37 | 'batch_size': 10000, 38 | 'n_test': 100000, 39 | 'n_train': 1000000, 40 | 41 | 'lr_init': 0.01, # initial learning rate 42 | 'pre_low_lr': 3, # number of epochs at the start with very low learning rate 43 | 'final_decay': 0.01, # fraction of the learning rate to reach at final epoch 44 | 'l2_weight_reg': 1.86e-05, # strength of the weight regularization 45 | 'adam_betas': (0.9, 0.95), 46 | } 47 | 48 | # DATA LOADERS 49 | train_loader, test_loader = prepare_data_loaders(c['data_model'], c['n_train'], c['n_test'], c['batch_size']) 50 | c['train_loader'] = train_loader 51 | c['test_loader'] = test_loader 52 | 53 | 54 | ############################## 55 | ### MODEL ARCHITECTURE ### 56 | ############################## 57 | 58 | x_lane = [InputNode(c['ndim_x'], name='x')] 59 | 60 | for i in range(c['n_blocks']): 61 | if i > 0: 62 | x_lane.append(Node(x_lane[-1], 63 | HouseholderPerm, 64 | {'fixed': False, 'n_reflections': c['ndim_x']}, 65 | name=f'perm_{i}')) 66 | 67 | x_lane.append(Node(x_lane[-1], 68 | HierarchicalAffineCouplingBlock, 69 | {'c_internal': [c['hidden_layer_sizes'], c['hidden_layer_sizes']//2, c['hidden_layer_sizes']//4]}, 70 | name=f'hac_{i+1}')) 71 | 72 | x_lane.append(OutputNode(x_lane[-1], name='z')) 73 | 74 | model = ReversibleGraphNet(x_lane, verbose=False) 75 | model.to(c['device']) 76 | model.params_trainable = list(filter(lambda p: p.requires_grad, model.parameters())) 77 | 78 | 79 | def model_inverse(test_z): 80 | return model(test_z, rev=True) 81 | 82 | 83 | c['model'] = model 84 | c['model_inverse'] = model_inverse 85 | 86 | # create namedtuple from config dictionary 87 | c = namedtuple("Configuration",c.keys())(*c.values()) 88 | -------------------------------------------------------------------------------- /configs/plus_shape/unconditional_hint_2_full.py: -------------------------------------------------------------------------------- 1 | from collections import namedtuple 2 | import torch 3 | 4 | from FrEIA.framework import * 5 | from FrEIA.modules import * 6 | 7 | from data import PlusShapeModel as model 8 | from data import prepare_data_loaders 9 | n_parameters = model.n_parameters 10 | 11 | 12 | ###################################################### 13 | ### TRAINING CONFIGURATION AND HYPERPARAMETERS ### 14 | ###################################################### 15 | 16 | c = { 17 | # GENERAL STUFF 18 | 'suffix': f'{model.name}_unconditional_hint-2-full', # identifier for trained models and outputs 19 | 'device': 'cuda', # 'cuda' for GPU, 'cpu' for CPU 20 | 'interactive_visualization': True, # requires visdom package to be installed 21 | 22 | # DATA 23 | 'ndim_x': n_parameters, 24 | 'ndim_y': 0, 25 | 'ndim_z': n_parameters, 26 | 'data_model': model(), 27 | 'vis_y_target': None, 28 | 29 | # MODEL ARCHITECTURE 30 | 'n_blocks': 2, 31 | 'hidden_layer_sizes': 175, # 500k 32 | 'init_scale': 0.005, 33 | 34 | # TRAINING HYPERPARAMETERS 35 | 'n_epochs': 50, # total number of epochs 36 | 'max_batches_per_epoch': 100, # iterations per epoch (or if training data exhausted) 37 | 'batch_size': 10000, 38 | 'n_test': 100000, 39 | 'n_train': 1000000, 40 | 41 | 'lr_init': 0.01, # initial learning rate 42 | 'pre_low_lr': 3, # number of epochs at the start with very low learning rate 43 | 'final_decay': 0.01, # fraction of the learning rate to reach at final epoch 44 | 'l2_weight_reg': 1.86e-05, # strength of the weight regularization 45 | 'adam_betas': (0.9, 0.95), 46 | } 47 | 48 | # DATA LOADERS 49 | train_loader, test_loader = prepare_data_loaders(c['data_model'], c['n_train'], c['n_test'], c['batch_size']) 50 | c['train_loader'] = train_loader 51 | c['test_loader'] = test_loader 52 | 53 | 54 | ############################## 55 | ### MODEL ARCHITECTURE ### 56 | ############################## 57 | 58 | x_lane = [InputNode(c['ndim_x'], name='x')] 59 | 60 | for i in range(c['n_blocks']): 61 | if i > 0: 62 | x_lane.append(Node(x_lane[-1], 63 | HouseholderPerm, 64 | {'fixed': False, 'n_reflections': c['ndim_x']}, 65 | name=f'perm_{i}')) 66 | 67 | x_lane.append(Node(x_lane[-1], 68 | HierarchicalAffineCouplingBlock, 69 | {'c_internal': [c['hidden_layer_sizes'], c['hidden_layer_sizes']//2, c['hidden_layer_sizes']//4, c['hidden_layer_sizes']//8]}, 70 | name=f'hac_{i+1}')) 71 | 72 | x_lane.append(OutputNode(x_lane[-1], name='z')) 73 | 74 | model = ReversibleGraphNet(x_lane, verbose=False) 75 | model.to(c['device']) 76 | model.params_trainable = list(filter(lambda p: p.requires_grad, model.parameters())) 77 | 78 | 79 | def model_inverse(test_z): 80 | return model(test_z, rev=True) 81 | 82 | 83 | c['model'] = model 84 | c['model_inverse'] = model_inverse 85 | 86 | # create namedtuple from config dictionary 87 | c = namedtuple("Configuration",c.keys())(*c.values()) 88 | -------------------------------------------------------------------------------- /configs/plus_shape/unconditional_hint_32_0_big.py: -------------------------------------------------------------------------------- 1 | from collections import namedtuple 2 | import torch 3 | 4 | from FrEIA.framework import * 5 | from FrEIA.modules import * 6 | 7 | from data import PlusShapeModel as model 8 | from data import prepare_data_loaders 9 | n_parameters = model.n_parameters 10 | 11 | 12 | ###################################################### 13 | ### TRAINING CONFIGURATION AND HYPERPARAMETERS ### 14 | ###################################################### 15 | 16 | c = { 17 | # GENERAL STUFF 18 | 'suffix': f'{model.name}_unconditional_hint-32-0-big', # identifier for trained models and outputs 19 | 'device': 'cuda', # 'cuda' for GPU, 'cpu' for CPU 20 | 'interactive_visualization': True, # requires visdom package to be installed 21 | 22 | # DATA 23 | 'ndim_x': n_parameters, 24 | 'ndim_y': 0, 25 | 'ndim_z': n_parameters, 26 | 'data_model': model(), 27 | 'vis_y_target': None, 28 | 29 | # MODEL ARCHITECTURE 30 | 'n_blocks': 32, 31 | 'hidden_layer_sizes': 512, 32 | 'init_scale': 0.005, 33 | 'recursion_depth': 0, 34 | 35 | # TRAINING HYPERPARAMETERS 36 | 'n_epochs': 50, # total number of epochs 37 | 'max_batches_per_epoch': 100, # iterations per epoch (or if training data exhausted) 38 | 'batch_size': 10000, 39 | 'n_test': 100000, 40 | 'n_train': 1000000, 41 | 42 | 'lr_init': 0.001, # initial learning rate 43 | 'pre_low_lr': 1, # number of epochs at the start with very low learning rate 44 | 'final_decay': 0.01, # fraction of the learning rate to reach at final epoch 45 | 'l2_weight_reg': 1.86e-05, # strength of the weight regularization 46 | 'adam_betas': (0.9, 0.95), 47 | } 48 | 49 | # DATA LOADERS 50 | train_loader, test_loader = prepare_data_loaders(c['data_model'], c['n_train'], c['n_test'], c['batch_size']) 51 | c['train_loader'] = train_loader 52 | c['test_loader'] = test_loader 53 | 54 | 55 | ############################## 56 | ### MODEL ARCHITECTURE ### 57 | ############################## 58 | 59 | x_lane = [InputNode(c['ndim_x'], name='x')] 60 | 61 | for i in range(c['n_blocks']): 62 | if i > 0: 63 | x_lane.append(Node(x_lane[-1], 64 | HouseholderPerm, 65 | {'fixed': False, 'n_reflections': c['ndim_x']}, 66 | name=f'perm_{i}')) 67 | 68 | x_lane.append(Node(x_lane[-1], 69 | HierarchicalAffineCouplingBlock, 70 | {'c_internal': [c['hidden_layer_sizes'], c['hidden_layer_sizes']//2, c['hidden_layer_sizes']//4, c['hidden_layer_sizes']//8], 71 | 'max_splits': c['recursion_depth']}, 72 | name=f'hac_{i+1}')) 73 | 74 | x_lane.append(OutputNode(x_lane[-1], name='z')) 75 | 76 | model = ReversibleGraphNet(x_lane, verbose=False) 77 | model.to(c['device']) 78 | model.params_trainable = list(filter(lambda p: p.requires_grad, model.parameters())) 79 | 80 | 81 | def model_inverse(test_z): 82 | return model(test_z, rev=True) 83 | 84 | 85 | c['model'] = model 86 | c['model_inverse'] = model_inverse 87 | 88 | # create namedtuple from config dictionary 89 | c = namedtuple("Configuration",c.keys())(*c.values()) 90 | -------------------------------------------------------------------------------- /configs/plus_shape/unconditional_hint_32_0_small.py: -------------------------------------------------------------------------------- 1 | from collections import namedtuple 2 | import torch 3 | 4 | from FrEIA.framework import * 5 | from FrEIA.modules import * 6 | 7 | from data import PlusShapeModel as model 8 | from data import prepare_data_loaders 9 | n_parameters = model.n_parameters 10 | 11 | 12 | ###################################################### 13 | ### TRAINING CONFIGURATION AND HYPERPARAMETERS ### 14 | ###################################################### 15 | 16 | c = { 17 | # GENERAL STUFF 18 | 'suffix': f'{model.name}_unconditional_hint-32-0-small', # identifier for trained models and outputs 19 | 'device': 'cuda', # 'cuda' for GPU, 'cpu' for CPU 20 | 'interactive_visualization': True, # requires visdom package to be installed 21 | 22 | # DATA 23 | 'ndim_x': n_parameters, 24 | 'ndim_y': 0, 25 | 'ndim_z': n_parameters, 26 | 'data_model': model(), 27 | 'vis_y_target': None, 28 | 29 | # MODEL ARCHITECTURE 30 | 'n_blocks': 32, 31 | 'hidden_layer_sizes': 128, 32 | 'init_scale': 0.005, 33 | 'recursion_depth': 0, 34 | 35 | # TRAINING HYPERPARAMETERS 36 | 'n_epochs': 50, # total number of epochs 37 | 'max_batches_per_epoch': 100, # iterations per epoch (or if training data exhausted) 38 | 'batch_size': 10000, 39 | 'n_test': 100000, 40 | 'n_train': 1000000, 41 | 42 | 'lr_init': 0.001, # initial learning rate 43 | 'pre_low_lr': 1, # number of epochs at the start with very low learning rate 44 | 'final_decay': 0.01, # fraction of the learning rate to reach at final epoch 45 | 'l2_weight_reg': 1.86e-05, # strength of the weight regularization 46 | 'adam_betas': (0.9, 0.95), 47 | } 48 | 49 | # DATA LOADERS 50 | train_loader, test_loader = prepare_data_loaders(c['data_model'], c['n_train'], c['n_test'], c['batch_size']) 51 | c['train_loader'] = train_loader 52 | c['test_loader'] = test_loader 53 | 54 | 55 | ############################## 56 | ### MODEL ARCHITECTURE ### 57 | ############################## 58 | 59 | x_lane = [InputNode(c['ndim_x'], name='x')] 60 | 61 | for i in range(c['n_blocks']): 62 | if i > 0: 63 | x_lane.append(Node(x_lane[-1], 64 | HouseholderPerm, 65 | {'fixed': False, 'n_reflections': c['ndim_x']}, 66 | name=f'perm_{i}')) 67 | 68 | x_lane.append(Node(x_lane[-1], 69 | HierarchicalAffineCouplingBlock, 70 | {'c_internal': [c['hidden_layer_sizes'], c['hidden_layer_sizes']//2, c['hidden_layer_sizes']//4, c['hidden_layer_sizes']//8], 71 | 'max_splits': c['recursion_depth']}, 72 | name=f'hac_{i+1}')) 73 | 74 | x_lane.append(OutputNode(x_lane[-1], name='z')) 75 | 76 | model = ReversibleGraphNet(x_lane, verbose=False) 77 | model.to(c['device']) 78 | model.params_trainable = list(filter(lambda p: p.requires_grad, model.parameters())) 79 | 80 | 81 | def model_inverse(test_z): 82 | return model(test_z, rev=True) 83 | 84 | 85 | c['model'] = model 86 | c['model_inverse'] = model_inverse 87 | 88 | # create namedtuple from config dictionary 89 | c = namedtuple("Configuration",c.keys())(*c.values()) 90 | -------------------------------------------------------------------------------- /configs/plus_shape/unconditional_hint_4_0_big.py: -------------------------------------------------------------------------------- 1 | from collections import namedtuple 2 | import torch 3 | 4 | from FrEIA.framework import * 5 | from FrEIA.modules import * 6 | 7 | from data import PlusShapeModel as model 8 | from data import prepare_data_loaders 9 | n_parameters = model.n_parameters 10 | 11 | 12 | ###################################################### 13 | ### TRAINING CONFIGURATION AND HYPERPARAMETERS ### 14 | ###################################################### 15 | 16 | c = { 17 | # GENERAL STUFF 18 | 'suffix': f'{model.name}_unconditional_hint-4-0-big', # identifier for trained models and outputs 19 | 'device': 'cuda', # 'cuda' for GPU, 'cpu' for CPU 20 | 'interactive_visualization': True, # requires visdom package to be installed 21 | 22 | # DATA 23 | 'ndim_x': n_parameters, 24 | 'ndim_y': 0, 25 | 'ndim_z': n_parameters, 26 | 'data_model': model(), 27 | 'vis_y_target': None, 28 | 29 | # MODEL ARCHITECTURE 30 | 'n_blocks': 4, 31 | 'hidden_layer_sizes': 512, 32 | 'init_scale': 0.005, 33 | 'recursion_depth': 0, 34 | 35 | # TRAINING HYPERPARAMETERS 36 | 'n_epochs': 50, # total number of epochs 37 | 'max_batches_per_epoch': 100, # iterations per epoch (or if training data exhausted) 38 | 'batch_size': 10000, 39 | 'n_test': 100000, 40 | 'n_train': 1000000, 41 | 42 | 'lr_init': 0.01, # initial learning rate 43 | 'pre_low_lr': 3, # number of epochs at the start with very low learning rate 44 | 'final_decay': 0.01, # fraction of the learning rate to reach at final epoch 45 | 'l2_weight_reg': 1.86e-05, # strength of the weight regularization 46 | 'adam_betas': (0.9, 0.95), 47 | } 48 | 49 | # DATA LOADERS 50 | train_loader, test_loader = prepare_data_loaders(c['data_model'], c['n_train'], c['n_test'], c['batch_size']) 51 | c['train_loader'] = train_loader 52 | c['test_loader'] = test_loader 53 | 54 | 55 | ############################## 56 | ### MODEL ARCHITECTURE ### 57 | ############################## 58 | 59 | x_lane = [InputNode(c['ndim_x'], name='x')] 60 | 61 | for i in range(c['n_blocks']): 62 | if i > 0: 63 | x_lane.append(Node(x_lane[-1], 64 | HouseholderPerm, 65 | {'fixed': False, 'n_reflections': c['ndim_x']}, 66 | name=f'perm_{i}')) 67 | 68 | x_lane.append(Node(x_lane[-1], 69 | HierarchicalAffineCouplingBlock, 70 | {'c_internal': [c['hidden_layer_sizes'], c['hidden_layer_sizes']//2, c['hidden_layer_sizes']//4, c['hidden_layer_sizes']//8], 71 | 'max_splits': c['recursion_depth']}, 72 | name=f'hac_{i+1}')) 73 | 74 | x_lane.append(OutputNode(x_lane[-1], name='z')) 75 | 76 | model = ReversibleGraphNet(x_lane, verbose=False) 77 | model.to(c['device']) 78 | model.params_trainable = list(filter(lambda p: p.requires_grad, model.parameters())) 79 | 80 | 81 | def model_inverse(test_z): 82 | return model(test_z, rev=True) 83 | 84 | 85 | c['model'] = model 86 | c['model_inverse'] = model_inverse 87 | 88 | # create namedtuple from config dictionary 89 | c = namedtuple("Configuration",c.keys())(*c.values()) 90 | -------------------------------------------------------------------------------- /configs/plus_shape/unconditional_hint_4_0_small.py: -------------------------------------------------------------------------------- 1 | from collections import namedtuple 2 | import torch 3 | 4 | from FrEIA.framework import * 5 | from FrEIA.modules import * 6 | 7 | from data import PlusShapeModel as model 8 | from data import prepare_data_loaders 9 | n_parameters = model.n_parameters 10 | 11 | 12 | ###################################################### 13 | ### TRAINING CONFIGURATION AND HYPERPARAMETERS ### 14 | ###################################################### 15 | 16 | c = { 17 | # GENERAL STUFF 18 | 'suffix': f'{model.name}_unconditional_hint-4-0-small', # identifier for trained models and outputs 19 | 'device': 'cuda', # 'cuda' for GPU, 'cpu' for CPU 20 | 'interactive_visualization': True, # requires visdom package to be installed 21 | 22 | # DATA 23 | 'ndim_x': n_parameters, 24 | 'ndim_y': 0, 25 | 'ndim_z': n_parameters, 26 | 'data_model': model(), 27 | 'vis_y_target': None, 28 | 29 | # MODEL ARCHITECTURE 30 | 'n_blocks': 4, 31 | 'hidden_layer_sizes': 128, 32 | 'init_scale': 0.005, 33 | 'recursion_depth': 0, 34 | 35 | # TRAINING HYPERPARAMETERS 36 | 'n_epochs': 50, # total number of epochs 37 | 'max_batches_per_epoch': 100, # iterations per epoch (or if training data exhausted) 38 | 'batch_size': 10000, 39 | 'n_test': 100000, 40 | 'n_train': 1000000, 41 | 42 | 'lr_init': 0.01, # initial learning rate 43 | 'pre_low_lr': 3, # number of epochs at the start with very low learning rate 44 | 'final_decay': 0.01, # fraction of the learning rate to reach at final epoch 45 | 'l2_weight_reg': 1.86e-05, # strength of the weight regularization 46 | 'adam_betas': (0.9, 0.95), 47 | } 48 | 49 | # DATA LOADERS 50 | train_loader, test_loader = prepare_data_loaders(c['data_model'], c['n_train'], c['n_test'], c['batch_size']) 51 | c['train_loader'] = train_loader 52 | c['test_loader'] = test_loader 53 | 54 | 55 | ############################## 56 | ### MODEL ARCHITECTURE ### 57 | ############################## 58 | 59 | x_lane = [InputNode(c['ndim_x'], name='x')] 60 | 61 | for i in range(c['n_blocks']): 62 | if i > 0: 63 | x_lane.append(Node(x_lane[-1], 64 | HouseholderPerm, 65 | {'fixed': False, 'n_reflections': c['ndim_x']}, 66 | name=f'perm_{i}')) 67 | 68 | x_lane.append(Node(x_lane[-1], 69 | HierarchicalAffineCouplingBlock, 70 | {'c_internal': [c['hidden_layer_sizes'], c['hidden_layer_sizes']//2, c['hidden_layer_sizes']//4, c['hidden_layer_sizes']//8], 71 | 'max_splits': c['recursion_depth']}, 72 | name=f'hac_{i+1}')) 73 | 74 | x_lane.append(OutputNode(x_lane[-1], name='z')) 75 | 76 | model = ReversibleGraphNet(x_lane, verbose=False) 77 | model.to(c['device']) 78 | model.params_trainable = list(filter(lambda p: p.requires_grad, model.parameters())) 79 | 80 | 81 | def model_inverse(test_z): 82 | return model(test_z, rev=True) 83 | 84 | 85 | c['model'] = model 86 | c['model_inverse'] = model_inverse 87 | 88 | # create namedtuple from config dictionary 89 | c = namedtuple("Configuration",c.keys())(*c.values()) 90 | -------------------------------------------------------------------------------- /configs/plus_shape/unconditional_hint_4_1.py: -------------------------------------------------------------------------------- 1 | from collections import namedtuple 2 | import torch 3 | 4 | from FrEIA.framework import * 5 | from FrEIA.modules import * 6 | 7 | from data import PlusShapeModel as model 8 | from data import prepare_data_loaders 9 | n_parameters = model.n_parameters 10 | 11 | 12 | ###################################################### 13 | ### TRAINING CONFIGURATION AND HYPERPARAMETERS ### 14 | ###################################################### 15 | 16 | c = { 17 | # GENERAL STUFF 18 | 'suffix': f'{model.name}_unconditional_hint-4-1', # identifier for trained models and outputs 19 | 'device': 'cuda', # 'cuda' for GPU, 'cpu' for CPU 20 | 'interactive_visualization': True, # requires visdom package to be installed 21 | 22 | # DATA 23 | 'ndim_x': n_parameters, 24 | 'ndim_y': 0, 25 | 'ndim_z': n_parameters, 26 | 'data_model': model(), 27 | 'vis_y_target': None, 28 | 29 | # MODEL ARCHITECTURE 30 | 'n_blocks': 4, 31 | 'hidden_layer_sizes': 358, # 2M 32 | 'init_scale': 0.005, 33 | 34 | # TRAINING HYPERPARAMETERS 35 | 'n_epochs': 50, # total number of epochs 36 | 'max_batches_per_epoch': 100, # iterations per epoch (or if training data exhausted) 37 | 'batch_size': 10000, 38 | 'n_test': 100000, 39 | 'n_train': 1000000, 40 | 41 | 'lr_init': 0.01, # initial learning rate 42 | 'pre_low_lr': 3, # number of epochs at the start with very low learning rate 43 | 'final_decay': 0.01, # fraction of the learning rate to reach at final epoch 44 | 'l2_weight_reg': 1.86e-05, # strength of the weight regularization 45 | 'adam_betas': (0.9, 0.95), 46 | } 47 | 48 | # DATA LOADERS 49 | train_loader, test_loader = prepare_data_loaders(c['data_model'], c['n_train'], c['n_test'], c['batch_size']) 50 | c['train_loader'] = train_loader 51 | c['test_loader'] = test_loader 52 | 53 | 54 | ############################## 55 | ### MODEL ARCHITECTURE ### 56 | ############################## 57 | 58 | x_lane = [InputNode(c['ndim_x'], name='x')] 59 | 60 | for i in range(c['n_blocks']): 61 | if i > 0: 62 | x_lane.append(Node(x_lane[-1], 63 | HouseholderPerm, 64 | {'fixed': False, 'n_reflections': c['ndim_x']}, 65 | name=f'perm_{i}')) 66 | 67 | x_lane.append(Node(x_lane[-1], 68 | HierarchicalAffineCouplingBlock, 69 | {'c_internal': [c['hidden_layer_sizes'], c['hidden_layer_sizes']//2], 70 | 'max_splits': 1}, 71 | name=f'hac_{i+1}')) 72 | 73 | x_lane.append(OutputNode(x_lane[-1], name='z')) 74 | 75 | model = ReversibleGraphNet(x_lane, verbose=False) 76 | model.to(c['device']) 77 | model.params_trainable = list(filter(lambda p: p.requires_grad, model.parameters())) 78 | 79 | 80 | def model_inverse(test_z): 81 | return model(test_z, rev=True) 82 | 83 | 84 | c['model'] = model 85 | c['model_inverse'] = model_inverse 86 | 87 | # create namedtuple from config dictionary 88 | c = namedtuple("Configuration",c.keys())(*c.values()) 89 | -------------------------------------------------------------------------------- /configs/plus_shape/unconditional_hint_4_1_big.py: -------------------------------------------------------------------------------- 1 | from collections import namedtuple 2 | import torch 3 | 4 | from FrEIA.framework import * 5 | from FrEIA.modules import * 6 | 7 | from data import PlusShapeModel as model 8 | from data import prepare_data_loaders 9 | n_parameters = model.n_parameters 10 | 11 | 12 | ###################################################### 13 | ### TRAINING CONFIGURATION AND HYPERPARAMETERS ### 14 | ###################################################### 15 | 16 | c = { 17 | # GENERAL STUFF 18 | 'suffix': f'{model.name}_unconditional_hint-4-1-big', # identifier for trained models and outputs 19 | 'device': 'cuda', # 'cuda' for GPU, 'cpu' for CPU 20 | 'interactive_visualization': True, # requires visdom package to be installed 21 | 22 | # DATA 23 | 'ndim_x': n_parameters, 24 | 'ndim_y': 0, 25 | 'ndim_z': n_parameters, 26 | 'data_model': model(), 27 | 'vis_y_target': None, 28 | 29 | # MODEL ARCHITECTURE 30 | 'n_blocks': 4, 31 | 'hidden_layer_sizes': 512, 32 | 'init_scale': 0.005, 33 | 'recursion_depth': 1, 34 | 35 | # TRAINING HYPERPARAMETERS 36 | 'n_epochs': 50, # total number of epochs 37 | 'max_batches_per_epoch': 100, # iterations per epoch (or if training data exhausted) 38 | 'batch_size': 10000, 39 | 'n_test': 100000, 40 | 'n_train': 1000000, 41 | 42 | 'lr_init': 0.01, # initial learning rate 43 | 'pre_low_lr': 3, # number of epochs at the start with very low learning rate 44 | 'final_decay': 0.01, # fraction of the learning rate to reach at final epoch 45 | 'l2_weight_reg': 1.86e-05, # strength of the weight regularization 46 | 'adam_betas': (0.9, 0.95), 47 | } 48 | 49 | # DATA LOADERS 50 | train_loader, test_loader = prepare_data_loaders(c['data_model'], c['n_train'], c['n_test'], c['batch_size']) 51 | c['train_loader'] = train_loader 52 | c['test_loader'] = test_loader 53 | 54 | 55 | ############################## 56 | ### MODEL ARCHITECTURE ### 57 | ############################## 58 | 59 | x_lane = [InputNode(c['ndim_x'], name='x')] 60 | 61 | for i in range(c['n_blocks']): 62 | if i > 0: 63 | x_lane.append(Node(x_lane[-1], 64 | HouseholderPerm, 65 | {'fixed': False, 'n_reflections': c['ndim_x']}, 66 | name=f'perm_{i}')) 67 | 68 | x_lane.append(Node(x_lane[-1], 69 | HierarchicalAffineCouplingBlock, 70 | {'c_internal': [c['hidden_layer_sizes'], c['hidden_layer_sizes']//2, c['hidden_layer_sizes']//4, c['hidden_layer_sizes']//8], 71 | 'max_splits': c['recursion_depth']}, 72 | name=f'hac_{i+1}')) 73 | 74 | x_lane.append(OutputNode(x_lane[-1], name='z')) 75 | 76 | model = ReversibleGraphNet(x_lane, verbose=False) 77 | model.to(c['device']) 78 | model.params_trainable = list(filter(lambda p: p.requires_grad, model.parameters())) 79 | 80 | 81 | def model_inverse(test_z): 82 | return model(test_z, rev=True) 83 | 84 | 85 | c['model'] = model 86 | c['model_inverse'] = model_inverse 87 | 88 | # create namedtuple from config dictionary 89 | c = namedtuple("Configuration",c.keys())(*c.values()) 90 | -------------------------------------------------------------------------------- /configs/plus_shape/unconditional_hint_4_1_small.py: -------------------------------------------------------------------------------- 1 | from collections import namedtuple 2 | import torch 3 | 4 | from FrEIA.framework import * 5 | from FrEIA.modules import * 6 | 7 | from data import PlusShapeModel as model 8 | from data import prepare_data_loaders 9 | n_parameters = model.n_parameters 10 | 11 | 12 | ###################################################### 13 | ### TRAINING CONFIGURATION AND HYPERPARAMETERS ### 14 | ###################################################### 15 | 16 | c = { 17 | # GENERAL STUFF 18 | 'suffix': f'{model.name}_unconditional_hint-4-1-small', # identifier for trained models and outputs 19 | 'device': 'cuda', # 'cuda' for GPU, 'cpu' for CPU 20 | 'interactive_visualization': True, # requires visdom package to be installed 21 | 22 | # DATA 23 | 'ndim_x': n_parameters, 24 | 'ndim_y': 0, 25 | 'ndim_z': n_parameters, 26 | 'data_model': model(), 27 | 'vis_y_target': None, 28 | 29 | # MODEL ARCHITECTURE 30 | 'n_blocks': 4, 31 | 'hidden_layer_sizes': 128, 32 | 'init_scale': 0.005, 33 | 'recursion_depth': 1, 34 | 35 | # TRAINING HYPERPARAMETERS 36 | 'n_epochs': 50, # total number of epochs 37 | 'max_batches_per_epoch': 100, # iterations per epoch (or if training data exhausted) 38 | 'batch_size': 10000, 39 | 'n_test': 100000, 40 | 'n_train': 1000000, 41 | 42 | 'lr_init': 0.01, # initial learning rate 43 | 'pre_low_lr': 3, # number of epochs at the start with very low learning rate 44 | 'final_decay': 0.01, # fraction of the learning rate to reach at final epoch 45 | 'l2_weight_reg': 1.86e-05, # strength of the weight regularization 46 | 'adam_betas': (0.9, 0.95), 47 | } 48 | 49 | # DATA LOADERS 50 | train_loader, test_loader = prepare_data_loaders(c['data_model'], c['n_train'], c['n_test'], c['batch_size']) 51 | c['train_loader'] = train_loader 52 | c['test_loader'] = test_loader 53 | 54 | 55 | ############################## 56 | ### MODEL ARCHITECTURE ### 57 | ############################## 58 | 59 | x_lane = [InputNode(c['ndim_x'], name='x')] 60 | 61 | for i in range(c['n_blocks']): 62 | if i > 0: 63 | x_lane.append(Node(x_lane[-1], 64 | HouseholderPerm, 65 | {'fixed': False, 'n_reflections': c['ndim_x']}, 66 | name=f'perm_{i}')) 67 | 68 | x_lane.append(Node(x_lane[-1], 69 | HierarchicalAffineCouplingBlock, 70 | {'c_internal': [c['hidden_layer_sizes'], c['hidden_layer_sizes']//2, c['hidden_layer_sizes']//4, c['hidden_layer_sizes']//8], 71 | 'max_splits': c['recursion_depth']}, 72 | name=f'hac_{i+1}')) 73 | 74 | x_lane.append(OutputNode(x_lane[-1], name='z')) 75 | 76 | model = ReversibleGraphNet(x_lane, verbose=False) 77 | model.to(c['device']) 78 | model.params_trainable = list(filter(lambda p: p.requires_grad, model.parameters())) 79 | 80 | 81 | def model_inverse(test_z): 82 | return model(test_z, rev=True) 83 | 84 | 85 | c['model'] = model 86 | c['model_inverse'] = model_inverse 87 | 88 | # create namedtuple from config dictionary 89 | c = namedtuple("Configuration",c.keys())(*c.values()) 90 | -------------------------------------------------------------------------------- /configs/plus_shape/unconditional_hint_4_2.py: -------------------------------------------------------------------------------- 1 | from collections import namedtuple 2 | import torch 3 | 4 | from FrEIA.framework import * 5 | from FrEIA.modules import * 6 | 7 | from data import PlusShapeModel as model 8 | from data import prepare_data_loaders 9 | n_parameters = model.n_parameters 10 | 11 | 12 | ###################################################### 13 | ### TRAINING CONFIGURATION AND HYPERPARAMETERS ### 14 | ###################################################### 15 | 16 | c = { 17 | # GENERAL STUFF 18 | 'suffix': f'{model.name}_unconditional_hint-4-2', # identifier for trained models and outputs 19 | 'device': 'cuda', # 'cuda' for GPU, 'cpu' for CPU 20 | 'interactive_visualization': True, # requires visdom package to be installed 21 | 22 | # DATA 23 | 'ndim_x': n_parameters, 24 | 'ndim_y': 0, 25 | 'ndim_z': n_parameters, 26 | 'data_model': model(), 27 | 'vis_y_target': None, 28 | 29 | # MODEL ARCHITECTURE 30 | 'n_blocks': 4, 31 | 'hidden_layer_sizes': 327, # 2M 32 | 'init_scale': 0.005, 33 | 34 | # TRAINING HYPERPARAMETERS 35 | 'n_epochs': 50, # total number of epochs 36 | 'max_batches_per_epoch': 100, # iterations per epoch (or if training data exhausted) 37 | 'batch_size': 10000, 38 | 'n_test': 100000, 39 | 'n_train': 1000000, 40 | 41 | 'lr_init': 0.01, # initial learning rate 42 | 'pre_low_lr': 3, # number of epochs at the start with very low learning rate 43 | 'final_decay': 0.01, # fraction of the learning rate to reach at final epoch 44 | 'l2_weight_reg': 1.86e-05, # strength of the weight regularization 45 | 'adam_betas': (0.9, 0.95), 46 | } 47 | 48 | # DATA LOADERS 49 | train_loader, test_loader = prepare_data_loaders(c['data_model'], c['n_train'], c['n_test'], c['batch_size']) 50 | c['train_loader'] = train_loader 51 | c['test_loader'] = test_loader 52 | 53 | 54 | ############################## 55 | ### MODEL ARCHITECTURE ### 56 | ############################## 57 | 58 | x_lane = [InputNode(c['ndim_x'], name='x')] 59 | 60 | for i in range(c['n_blocks']): 61 | if i > 0: 62 | x_lane.append(Node(x_lane[-1], 63 | HouseholderPerm, 64 | {'fixed': False, 'n_reflections': c['ndim_x']}, 65 | name=f'perm_{i}')) 66 | 67 | x_lane.append(Node(x_lane[-1], 68 | HierarchicalAffineCouplingBlock, 69 | {'c_internal': [c['hidden_layer_sizes'], c['hidden_layer_sizes']//2, c['hidden_layer_sizes']//4], 70 | 'max_splits': 2}, 71 | name=f'hac_{i+1}')) 72 | 73 | x_lane.append(OutputNode(x_lane[-1], name='z')) 74 | 75 | model = ReversibleGraphNet(x_lane, verbose=False) 76 | model.to(c['device']) 77 | model.params_trainable = list(filter(lambda p: p.requires_grad, model.parameters())) 78 | 79 | 80 | def model_inverse(test_z): 81 | return model(test_z, rev=True) 82 | 83 | 84 | c['model'] = model 85 | c['model_inverse'] = model_inverse 86 | 87 | # create namedtuple from config dictionary 88 | c = namedtuple("Configuration",c.keys())(*c.values()) 89 | -------------------------------------------------------------------------------- /configs/plus_shape/unconditional_hint_4_2_big.py: -------------------------------------------------------------------------------- 1 | from collections import namedtuple 2 | import torch 3 | 4 | from FrEIA.framework import * 5 | from FrEIA.modules import * 6 | 7 | from data import PlusShapeModel as model 8 | from data import prepare_data_loaders 9 | n_parameters = model.n_parameters 10 | 11 | 12 | ###################################################### 13 | ### TRAINING CONFIGURATION AND HYPERPARAMETERS ### 14 | ###################################################### 15 | 16 | c = { 17 | # GENERAL STUFF 18 | 'suffix': f'{model.name}_unconditional_hint-4-2-big', # identifier for trained models and outputs 19 | 'device': 'cuda', # 'cuda' for GPU, 'cpu' for CPU 20 | 'interactive_visualization': True, # requires visdom package to be installed 21 | 22 | # DATA 23 | 'ndim_x': n_parameters, 24 | 'ndim_y': 0, 25 | 'ndim_z': n_parameters, 26 | 'data_model': model(), 27 | 'vis_y_target': None, 28 | 29 | # MODEL ARCHITECTURE 30 | 'n_blocks': 4, 31 | 'hidden_layer_sizes': 512, 32 | 'init_scale': 0.005, 33 | 'recursion_depth': 2, 34 | 35 | # TRAINING HYPERPARAMETERS 36 | 'n_epochs': 50, # total number of epochs 37 | 'max_batches_per_epoch': 100, # iterations per epoch (or if training data exhausted) 38 | 'batch_size': 10000, 39 | 'n_test': 100000, 40 | 'n_train': 1000000, 41 | 42 | 'lr_init': 0.01, # initial learning rate 43 | 'pre_low_lr': 3, # number of epochs at the start with very low learning rate 44 | 'final_decay': 0.01, # fraction of the learning rate to reach at final epoch 45 | 'l2_weight_reg': 1.86e-05, # strength of the weight regularization 46 | 'adam_betas': (0.9, 0.95), 47 | } 48 | 49 | # DATA LOADERS 50 | train_loader, test_loader = prepare_data_loaders(c['data_model'], c['n_train'], c['n_test'], c['batch_size']) 51 | c['train_loader'] = train_loader 52 | c['test_loader'] = test_loader 53 | 54 | 55 | ############################## 56 | ### MODEL ARCHITECTURE ### 57 | ############################## 58 | 59 | x_lane = [InputNode(c['ndim_x'], name='x')] 60 | 61 | for i in range(c['n_blocks']): 62 | if i > 0: 63 | x_lane.append(Node(x_lane[-1], 64 | HouseholderPerm, 65 | {'fixed': False, 'n_reflections': c['ndim_x']}, 66 | name=f'perm_{i}')) 67 | 68 | x_lane.append(Node(x_lane[-1], 69 | HierarchicalAffineCouplingBlock, 70 | {'c_internal': [c['hidden_layer_sizes'], c['hidden_layer_sizes']//2, c['hidden_layer_sizes']//4, c['hidden_layer_sizes']//8], 71 | 'max_splits': c['recursion_depth']}, 72 | name=f'hac_{i+1}')) 73 | 74 | x_lane.append(OutputNode(x_lane[-1], name='z')) 75 | 76 | model = ReversibleGraphNet(x_lane, verbose=False) 77 | model.to(c['device']) 78 | model.params_trainable = list(filter(lambda p: p.requires_grad, model.parameters())) 79 | 80 | 81 | def model_inverse(test_z): 82 | return model(test_z, rev=True) 83 | 84 | 85 | c['model'] = model 86 | c['model_inverse'] = model_inverse 87 | 88 | # create namedtuple from config dictionary 89 | c = namedtuple("Configuration",c.keys())(*c.values()) 90 | -------------------------------------------------------------------------------- /configs/plus_shape/unconditional_hint_4_2_small.py: -------------------------------------------------------------------------------- 1 | from collections import namedtuple 2 | import torch 3 | 4 | from FrEIA.framework import * 5 | from FrEIA.modules import * 6 | 7 | from data import PlusShapeModel as model 8 | from data import prepare_data_loaders 9 | n_parameters = model.n_parameters 10 | 11 | 12 | ###################################################### 13 | ### TRAINING CONFIGURATION AND HYPERPARAMETERS ### 14 | ###################################################### 15 | 16 | c = { 17 | # GENERAL STUFF 18 | 'suffix': f'{model.name}_unconditional_hint-4-2-small', # identifier for trained models and outputs 19 | 'device': 'cuda', # 'cuda' for GPU, 'cpu' for CPU 20 | 'interactive_visualization': True, # requires visdom package to be installed 21 | 22 | # DATA 23 | 'ndim_x': n_parameters, 24 | 'ndim_y': 0, 25 | 'ndim_z': n_parameters, 26 | 'data_model': model(), 27 | 'vis_y_target': None, 28 | 29 | # MODEL ARCHITECTURE 30 | 'n_blocks': 4, 31 | 'hidden_layer_sizes': 128, 32 | 'init_scale': 0.005, 33 | 'recursion_depth': 2, 34 | 35 | # TRAINING HYPERPARAMETERS 36 | 'n_epochs': 50, # total number of epochs 37 | 'max_batches_per_epoch': 100, # iterations per epoch (or if training data exhausted) 38 | 'batch_size': 10000, 39 | 'n_test': 100000, 40 | 'n_train': 1000000, 41 | 42 | 'lr_init': 0.01, # initial learning rate 43 | 'pre_low_lr': 3, # number of epochs at the start with very low learning rate 44 | 'final_decay': 0.01, # fraction of the learning rate to reach at final epoch 45 | 'l2_weight_reg': 1.86e-05, # strength of the weight regularization 46 | 'adam_betas': (0.9, 0.95), 47 | } 48 | 49 | # DATA LOADERS 50 | train_loader, test_loader = prepare_data_loaders(c['data_model'], c['n_train'], c['n_test'], c['batch_size']) 51 | c['train_loader'] = train_loader 52 | c['test_loader'] = test_loader 53 | 54 | 55 | ############################## 56 | ### MODEL ARCHITECTURE ### 57 | ############################## 58 | 59 | x_lane = [InputNode(c['ndim_x'], name='x')] 60 | 61 | for i in range(c['n_blocks']): 62 | if i > 0: 63 | x_lane.append(Node(x_lane[-1], 64 | HouseholderPerm, 65 | {'fixed': False, 'n_reflections': c['ndim_x']}, 66 | name=f'perm_{i}')) 67 | 68 | x_lane.append(Node(x_lane[-1], 69 | HierarchicalAffineCouplingBlock, 70 | {'c_internal': [c['hidden_layer_sizes'], c['hidden_layer_sizes']//2, c['hidden_layer_sizes']//4, c['hidden_layer_sizes']//8], 71 | 'max_splits': c['recursion_depth']}, 72 | name=f'hac_{i+1}')) 73 | 74 | x_lane.append(OutputNode(x_lane[-1], name='z')) 75 | 76 | model = ReversibleGraphNet(x_lane, verbose=False) 77 | model.to(c['device']) 78 | model.params_trainable = list(filter(lambda p: p.requires_grad, model.parameters())) 79 | 80 | 81 | def model_inverse(test_z): 82 | return model(test_z, rev=True) 83 | 84 | 85 | c['model'] = model 86 | c['model_inverse'] = model_inverse 87 | 88 | # create namedtuple from config dictionary 89 | c = namedtuple("Configuration",c.keys())(*c.values()) 90 | -------------------------------------------------------------------------------- /configs/plus_shape/unconditional_hint_4_3.py: -------------------------------------------------------------------------------- 1 | from collections import namedtuple 2 | import torch 3 | 4 | from FrEIA.framework import * 5 | from FrEIA.modules import * 6 | 7 | from data import PlusShapeModel as model 8 | from data import prepare_data_loaders 9 | n_parameters = model.n_parameters 10 | 11 | 12 | ###################################################### 13 | ### TRAINING CONFIGURATION AND HYPERPARAMETERS ### 14 | ###################################################### 15 | 16 | c = { 17 | # GENERAL STUFF 18 | 'suffix': f'{model.name}_unconditional_hint-4-3', # identifier for trained models and outputs 19 | 'device': 'cuda', # 'cuda' for GPU, 'cpu' for CPU 20 | 'interactive_visualization': True, # requires visdom package to be installed 21 | 22 | # DATA 23 | 'ndim_x': n_parameters, 24 | 'ndim_y': 0, 25 | 'ndim_z': n_parameters, 26 | 'data_model': model(), 27 | 'vis_y_target': None, 28 | 29 | # MODEL ARCHITECTURE 30 | 'n_blocks': 4, 31 | 'hidden_layer_sizes': 314, # 2M 32 | 'init_scale': 0.005, 33 | 34 | # TRAINING HYPERPARAMETERS 35 | 'n_epochs': 50, # total number of epochs 36 | 'max_batches_per_epoch': 100, # iterations per epoch (or if training data exhausted) 37 | 'batch_size': 10000, 38 | 'n_test': 100000, 39 | 'n_train': 1000000, 40 | 41 | 'lr_init': 0.01, # initial learning rate 42 | 'pre_low_lr': 3, # number of epochs at the start with very low learning rate 43 | 'final_decay': 0.01, # fraction of the learning rate to reach at final epoch 44 | 'l2_weight_reg': 1.86e-05, # strength of the weight regularization 45 | 'adam_betas': (0.9, 0.95), 46 | } 47 | 48 | # DATA LOADERS 49 | train_loader, test_loader = prepare_data_loaders(c['data_model'], c['n_train'], c['n_test'], c['batch_size']) 50 | c['train_loader'] = train_loader 51 | c['test_loader'] = test_loader 52 | 53 | 54 | ############################## 55 | ### MODEL ARCHITECTURE ### 56 | ############################## 57 | 58 | x_lane = [InputNode(c['ndim_x'], name='x')] 59 | 60 | for i in range(c['n_blocks']): 61 | if i > 0: 62 | x_lane.append(Node(x_lane[-1], 63 | HouseholderPerm, 64 | {'fixed': False, 'n_reflections': c['ndim_x']}, 65 | name=f'perm_{i}')) 66 | 67 | x_lane.append(Node(x_lane[-1], 68 | HierarchicalAffineCouplingBlock, 69 | {'c_internal': [c['hidden_layer_sizes'], c['hidden_layer_sizes']//2, c['hidden_layer_sizes']//4, c['hidden_layer_sizes']//8], 70 | 'max_splits': 3}, 71 | name=f'hac_{i+1}')) 72 | 73 | x_lane.append(OutputNode(x_lane[-1], name='z')) 74 | 75 | model = ReversibleGraphNet(x_lane, verbose=False) 76 | model.to(c['device']) 77 | model.params_trainable = list(filter(lambda p: p.requires_grad, model.parameters())) 78 | 79 | 80 | def model_inverse(test_z): 81 | return model(test_z, rev=True) 82 | 83 | 84 | c['model'] = model 85 | c['model_inverse'] = model_inverse 86 | 87 | # create namedtuple from config dictionary 88 | c = namedtuple("Configuration",c.keys())(*c.values()) 89 | -------------------------------------------------------------------------------- /configs/plus_shape/unconditional_hint_4_3_big.py: -------------------------------------------------------------------------------- 1 | from collections import namedtuple 2 | import torch 3 | 4 | from FrEIA.framework import * 5 | from FrEIA.modules import * 6 | 7 | from data import PlusShapeModel as model 8 | from data import prepare_data_loaders 9 | n_parameters = model.n_parameters 10 | 11 | 12 | ###################################################### 13 | ### TRAINING CONFIGURATION AND HYPERPARAMETERS ### 14 | ###################################################### 15 | 16 | c = { 17 | # GENERAL STUFF 18 | 'suffix': f'{model.name}_unconditional_hint-4-3-big', # identifier for trained models and outputs 19 | 'device': 'cuda', # 'cuda' for GPU, 'cpu' for CPU 20 | 'interactive_visualization': True, # requires visdom package to be installed 21 | 22 | # DATA 23 | 'ndim_x': n_parameters, 24 | 'ndim_y': 0, 25 | 'ndim_z': n_parameters, 26 | 'data_model': model(), 27 | 'vis_y_target': None, 28 | 29 | # MODEL ARCHITECTURE 30 | 'n_blocks': 4, 31 | 'hidden_layer_sizes': 512, 32 | 'init_scale': 0.005, 33 | 'recursion_depth': 3, 34 | 35 | # TRAINING HYPERPARAMETERS 36 | 'n_epochs': 50, # total number of epochs 37 | 'max_batches_per_epoch': 100, # iterations per epoch (or if training data exhausted) 38 | 'batch_size': 10000, 39 | 'n_test': 100000, 40 | 'n_train': 1000000, 41 | 42 | 'lr_init': 0.01, # initial learning rate 43 | 'pre_low_lr': 3, # number of epochs at the start with very low learning rate 44 | 'final_decay': 0.01, # fraction of the learning rate to reach at final epoch 45 | 'l2_weight_reg': 1.86e-05, # strength of the weight regularization 46 | 'adam_betas': (0.9, 0.95), 47 | } 48 | 49 | # DATA LOADERS 50 | train_loader, test_loader = prepare_data_loaders(c['data_model'], c['n_train'], c['n_test'], c['batch_size']) 51 | c['train_loader'] = train_loader 52 | c['test_loader'] = test_loader 53 | 54 | 55 | ############################## 56 | ### MODEL ARCHITECTURE ### 57 | ############################## 58 | 59 | x_lane = [InputNode(c['ndim_x'], name='x')] 60 | 61 | for i in range(c['n_blocks']): 62 | if i > 0: 63 | x_lane.append(Node(x_lane[-1], 64 | HouseholderPerm, 65 | {'fixed': False, 'n_reflections': c['ndim_x']}, 66 | name=f'perm_{i}')) 67 | 68 | x_lane.append(Node(x_lane[-1], 69 | HierarchicalAffineCouplingBlock, 70 | {'c_internal': [c['hidden_layer_sizes'], c['hidden_layer_sizes']//2, c['hidden_layer_sizes']//4, c['hidden_layer_sizes']//8], 71 | 'max_splits': c['recursion_depth']}, 72 | name=f'hac_{i+1}')) 73 | 74 | x_lane.append(OutputNode(x_lane[-1], name='z')) 75 | 76 | model = ReversibleGraphNet(x_lane, verbose=False) 77 | model.to(c['device']) 78 | model.params_trainable = list(filter(lambda p: p.requires_grad, model.parameters())) 79 | 80 | 81 | def model_inverse(test_z): 82 | return model(test_z, rev=True) 83 | 84 | 85 | c['model'] = model 86 | c['model_inverse'] = model_inverse 87 | 88 | # create namedtuple from config dictionary 89 | c = namedtuple("Configuration",c.keys())(*c.values()) 90 | -------------------------------------------------------------------------------- /configs/plus_shape/unconditional_hint_4_3_constwidth.py: -------------------------------------------------------------------------------- 1 | from collections import namedtuple 2 | import torch 3 | 4 | from FrEIA.framework import * 5 | from FrEIA.modules import * 6 | 7 | from data import PlusShapeModel as model 8 | from data import prepare_data_loaders 9 | n_parameters = model.n_parameters 10 | 11 | 12 | ###################################################### 13 | ### TRAINING CONFIGURATION AND HYPERPARAMETERS ### 14 | ###################################################### 15 | 16 | c = { 17 | # GENERAL STUFF 18 | 'suffix': f'{model.name}_unconditional_hint-4-3-constwidth', # identifier for trained models and outputs 19 | 'device': 'cuda', # 'cuda' for GPU, 'cpu' for CPU 20 | 'interactive_visualization': True, # requires visdom package to be installed 21 | 22 | # DATA 23 | 'ndim_x': n_parameters, 24 | 'ndim_y': 0, 25 | 'ndim_z': n_parameters, 26 | 'data_model': model(), 27 | 'vis_y_target': None, 28 | 29 | # MODEL ARCHITECTURE 30 | 'n_blocks': 4, 31 | 'hidden_layer_sizes': 116, # 2M 32 | 'init_scale': 0.005, 33 | 'recursion_depth': 3, 34 | 35 | # TRAINING HYPERPARAMETERS 36 | 'n_epochs': 50, # total number of epochs 37 | 'max_batches_per_epoch': 100, # iterations per epoch (or if training data exhausted) 38 | 'batch_size': 10000, 39 | 'n_test': 100000, 40 | 'n_train': 1000000, 41 | 42 | 'lr_init': 0.01, # initial learning rate 43 | 'pre_low_lr': 1, # number of epochs at the start with very low learning rate 44 | 'final_decay': 0.01, # fraction of the learning rate to reach at final epoch 45 | 'l2_weight_reg': 1.86e-05, # strength of the weight regularization 46 | 'adam_betas': (0.9, 0.95), 47 | } 48 | 49 | # DATA LOADERS 50 | train_loader, test_loader = prepare_data_loaders(c['data_model'], c['n_train'], c['n_test'], c['batch_size']) 51 | c['train_loader'] = train_loader 52 | c['test_loader'] = test_loader 53 | 54 | 55 | ############################## 56 | ### MODEL ARCHITECTURE ### 57 | ############################## 58 | 59 | x_lane = [InputNode(c['ndim_x'], name='x')] 60 | 61 | for i in range(c['n_blocks']): 62 | if i > 0: 63 | x_lane.append(Node(x_lane[-1], 64 | HouseholderPerm, 65 | {'fixed': True, 'n_reflections': c['ndim_x']}, 66 | name=f'perm_{i}')) 67 | 68 | x_lane.append(Node(x_lane[-1], 69 | HierarchicalAffineCouplingBlock, 70 | {'c_internal': [c['hidden_layer_sizes']], 71 | 'max_splits': c['recursion_depth']}, 72 | name=f'hac_{i+1}')) 73 | 74 | x_lane.append(OutputNode(x_lane[-1], name='z')) 75 | 76 | model = ReversibleGraphNet(x_lane, verbose=False) 77 | model.to(c['device']) 78 | model.params_trainable = list(filter(lambda p: p.requires_grad, model.parameters())) 79 | 80 | 81 | def model_inverse(test_z): 82 | return model(test_z, rev=True) 83 | 84 | 85 | c['model'] = model 86 | c['model_inverse'] = model_inverse 87 | 88 | # create namedtuple from config dictionary 89 | c = namedtuple("Configuration",c.keys())(*c.values()) 90 | -------------------------------------------------------------------------------- /configs/plus_shape/unconditional_hint_4_3_lessshrink.py: -------------------------------------------------------------------------------- 1 | from collections import namedtuple 2 | from numpy import sqrt 3 | import torch 4 | 5 | from FrEIA.framework import * 6 | from FrEIA.modules import * 7 | 8 | from data import PlusShapeModel as model 9 | from data import prepare_data_loaders 10 | n_parameters = model.n_parameters 11 | 12 | 13 | ###################################################### 14 | ### TRAINING CONFIGURATION AND HYPERPARAMETERS ### 15 | ###################################################### 16 | 17 | c = { 18 | # GENERAL STUFF 19 | 'suffix': f'{model.name}_unconditional_hint-4-3-lessshrink', # identifier for trained models and outputs 20 | 'device': 'cuda', # 'cuda' for GPU, 'cpu' for CPU 21 | 'interactive_visualization': True, # requires visdom package to be installed 22 | 23 | # DATA 24 | 'ndim_x': n_parameters, 25 | 'ndim_y': 0, 26 | 'ndim_z': n_parameters, 27 | 'data_model': model(), 28 | 'vis_y_target': None, 29 | 30 | # MODEL ARCHITECTURE 31 | 'n_blocks': 4, 32 | 'hidden_layer_sizes': 219, # 2M 33 | 'init_scale': 0.005, 34 | 'recursion_depth': 3, 35 | 36 | # TRAINING HYPERPARAMETERS 37 | 'n_epochs': 50, # total number of epochs 38 | 'max_batches_per_epoch': 100, # iterations per epoch (or if training data exhausted) 39 | 'batch_size': 10000, 40 | 'n_test': 100000, 41 | 'n_train': 1000000, 42 | 43 | 'lr_init': 0.01, # initial learning rate 44 | 'pre_low_lr': 1, # number of epochs at the start with very low learning rate 45 | 'final_decay': 0.01, # fraction of the learning rate to reach at final epoch 46 | 'l2_weight_reg': 1.86e-05, # strength of the weight regularization 47 | 'adam_betas': (0.9, 0.95), 48 | } 49 | 50 | # DATA LOADERS 51 | train_loader, test_loader = prepare_data_loaders(c['data_model'], c['n_train'], c['n_test'], c['batch_size']) 52 | c['train_loader'] = train_loader 53 | c['test_loader'] = test_loader 54 | 55 | 56 | ############################## 57 | ### MODEL ARCHITECTURE ### 58 | ############################## 59 | 60 | x_lane = [InputNode(c['ndim_x'], name='x')] 61 | 62 | for i in range(c['n_blocks']): 63 | if i > 0: 64 | x_lane.append(Node(x_lane[-1], 65 | HouseholderPerm, 66 | {'fixed': True, 'n_reflections': c['ndim_x']}, 67 | name=f'perm_{i}')) 68 | 69 | x_lane.append(Node(x_lane[-1], 70 | HierarchicalAffineCouplingBlock, 71 | {'c_internal': [c['hidden_layer_sizes'], int(c['hidden_layer_sizes']//sqrt(2)), c['hidden_layer_sizes']//2, int(c['hidden_layer_sizes']//(2*sqrt(2)))], 72 | 'max_splits': c['recursion_depth']}, 73 | name=f'hac_{i+1}')) 74 | 75 | x_lane.append(OutputNode(x_lane[-1], name='z')) 76 | 77 | model = ReversibleGraphNet(x_lane, verbose=False) 78 | model.to(c['device']) 79 | model.params_trainable = list(filter(lambda p: p.requires_grad, model.parameters())) 80 | 81 | 82 | def model_inverse(test_z): 83 | return model(test_z, rev=True) 84 | 85 | 86 | c['model'] = model 87 | c['model_inverse'] = model_inverse 88 | 89 | # create namedtuple from config dictionary 90 | c = namedtuple("Configuration",c.keys())(*c.values()) 91 | -------------------------------------------------------------------------------- /configs/plus_shape/unconditional_hint_4_3_reshuffle.py: -------------------------------------------------------------------------------- 1 | from collections import namedtuple 2 | import torch 3 | 4 | from FrEIA.framework import * 5 | from FrEIA.modules import * 6 | 7 | from data import PlusShapeModel as model 8 | from data import prepare_data_loaders 9 | n_parameters = model.n_parameters 10 | 11 | 12 | ###################################################### 13 | ### TRAINING CONFIGURATION AND HYPERPARAMETERS ### 14 | ###################################################### 15 | 16 | c = { 17 | # GENERAL STUFF 18 | 'suffix': f'{model.name}_unconditional_hint-4-3-reshuffle', # identifier for trained models and outputs 19 | 'device': 'cuda', # 'cuda' for GPU, 'cpu' for CPU 20 | 'interactive_visualization': True, # requires visdom package to be installed 21 | 22 | # DATA 23 | 'ndim_x': n_parameters, 24 | 'ndim_y': 0, 25 | 'ndim_z': n_parameters, 26 | 'data_model': model(), 27 | 'vis_y_target': None, 28 | 29 | # MODEL ARCHITECTURE 30 | 'n_blocks': 4, 31 | 'hidden_layer_sizes': 314, # 2M 32 | 'init_scale': 0.005, 33 | 'recursion_depth': 3, 34 | 35 | # TRAINING HYPERPARAMETERS 36 | 'n_epochs': 50, # total number of epochs 37 | 'max_batches_per_epoch': 100, # iterations per epoch (or if training data exhausted) 38 | 'batch_size': 10000, 39 | 'n_test': 100000, 40 | 'n_train': 1000000, 41 | 42 | 'lr_init': 0.01, # initial learning rate 43 | 'pre_low_lr': 1, # number of epochs at the start with very low learning rate 44 | 'final_decay': 0.01, # fraction of the learning rate to reach at final epoch 45 | 'l2_weight_reg': 1.86e-05, # strength of the weight regularization 46 | 'adam_betas': (0.9, 0.95), 47 | } 48 | 49 | # DATA LOADERS 50 | train_loader, test_loader = prepare_data_loaders(c['data_model'], c['n_train'], c['n_test'], c['batch_size']) 51 | c['train_loader'] = train_loader 52 | c['test_loader'] = test_loader 53 | 54 | 55 | ############################## 56 | ### MODEL ARCHITECTURE ### 57 | ############################## 58 | 59 | x_lane = [InputNode(c['ndim_x'], name='x')] 60 | 61 | for i in range(c['n_blocks']): 62 | if i > 0: 63 | x_lane.append(Node(x_lane[-1], 64 | HouseholderPerm, 65 | {'fixed': True, 'n_reflections': c['ndim_x']}, 66 | name=f'perm_{i}')) 67 | 68 | x_lane.append(Node(x_lane[-1], 69 | HierarchicalAffineCouplingBlock, 70 | {'c_internal': [c['hidden_layer_sizes'], c['hidden_layer_sizes']//2, c['hidden_layer_sizes']//4, c['hidden_layer_sizes']//8], 71 | 'max_splits': c['recursion_depth'], 72 | 'reshuffle': True}, 73 | name=f'hac_{i+1}')) 74 | 75 | x_lane.append(OutputNode(x_lane[-1], name='z')) 76 | 77 | model = ReversibleGraphNet(x_lane, verbose=False) 78 | model.to(c['device']) 79 | model.params_trainable = list(filter(lambda p: p.requires_grad, model.parameters())) 80 | 81 | 82 | def model_inverse(test_z): 83 | return model(test_z, rev=True) 84 | 85 | 86 | c['model'] = model 87 | c['model_inverse'] = model_inverse 88 | 89 | # create namedtuple from config dictionary 90 | c = namedtuple("Configuration",c.keys())(*c.values()) 91 | -------------------------------------------------------------------------------- /configs/plus_shape/unconditional_hint_4_3_small.py: -------------------------------------------------------------------------------- 1 | from collections import namedtuple 2 | import torch 3 | 4 | from FrEIA.framework import * 5 | from FrEIA.modules import * 6 | 7 | from data import PlusShapeModel as model 8 | from data import prepare_data_loaders 9 | n_parameters = model.n_parameters 10 | 11 | 12 | ###################################################### 13 | ### TRAINING CONFIGURATION AND HYPERPARAMETERS ### 14 | ###################################################### 15 | 16 | c = { 17 | # GENERAL STUFF 18 | 'suffix': f'{model.name}_unconditional_hint-4-3-small', # identifier for trained models and outputs 19 | 'device': 'cuda', # 'cuda' for GPU, 'cpu' for CPU 20 | 'interactive_visualization': True, # requires visdom package to be installed 21 | 22 | # DATA 23 | 'ndim_x': n_parameters, 24 | 'ndim_y': 0, 25 | 'ndim_z': n_parameters, 26 | 'data_model': model(), 27 | 'vis_y_target': None, 28 | 29 | # MODEL ARCHITECTURE 30 | 'n_blocks': 4, 31 | 'hidden_layer_sizes': 128, 32 | 'init_scale': 0.005, 33 | 'recursion_depth': 3, 34 | 35 | # TRAINING HYPERPARAMETERS 36 | 'n_epochs': 50, # total number of epochs 37 | 'max_batches_per_epoch': 100, # iterations per epoch (or if training data exhausted) 38 | 'batch_size': 10000, 39 | 'n_test': 100000, 40 | 'n_train': 1000000, 41 | 42 | 'lr_init': 0.01, # initial learning rate 43 | 'pre_low_lr': 3, # number of epochs at the start with very low learning rate 44 | 'final_decay': 0.01, # fraction of the learning rate to reach at final epoch 45 | 'l2_weight_reg': 1.86e-05, # strength of the weight regularization 46 | 'adam_betas': (0.9, 0.95), 47 | } 48 | 49 | # DATA LOADERS 50 | train_loader, test_loader = prepare_data_loaders(c['data_model'], c['n_train'], c['n_test'], c['batch_size']) 51 | c['train_loader'] = train_loader 52 | c['test_loader'] = test_loader 53 | 54 | 55 | ############################## 56 | ### MODEL ARCHITECTURE ### 57 | ############################## 58 | 59 | x_lane = [InputNode(c['ndim_x'], name='x')] 60 | 61 | for i in range(c['n_blocks']): 62 | if i > 0: 63 | x_lane.append(Node(x_lane[-1], 64 | HouseholderPerm, 65 | {'fixed': False, 'n_reflections': c['ndim_x']}, 66 | name=f'perm_{i}')) 67 | 68 | x_lane.append(Node(x_lane[-1], 69 | HierarchicalAffineCouplingBlock, 70 | {'c_internal': [c['hidden_layer_sizes'], c['hidden_layer_sizes']//2, c['hidden_layer_sizes']//4, c['hidden_layer_sizes']//8], 71 | 'max_splits': c['recursion_depth']}, 72 | name=f'hac_{i+1}')) 73 | 74 | x_lane.append(OutputNode(x_lane[-1], name='z')) 75 | 76 | model = ReversibleGraphNet(x_lane, verbose=False) 77 | model.to(c['device']) 78 | model.params_trainable = list(filter(lambda p: p.requires_grad, model.parameters())) 79 | 80 | 81 | def model_inverse(test_z): 82 | return model(test_z, rev=True) 83 | 84 | 85 | c['model'] = model 86 | c['model_inverse'] = model_inverse 87 | 88 | # create namedtuple from config dictionary 89 | c = namedtuple("Configuration",c.keys())(*c.values()) 90 | -------------------------------------------------------------------------------- /configs/plus_shape/unconditional_hint_4_full.py: -------------------------------------------------------------------------------- 1 | from collections import namedtuple 2 | import torch 3 | 4 | from FrEIA.framework import * 5 | from FrEIA.modules import * 6 | 7 | from data import PlusShapeModel as model 8 | from data import prepare_data_loaders 9 | n_parameters = model.n_parameters 10 | 11 | 12 | ###################################################### 13 | ### TRAINING CONFIGURATION AND HYPERPARAMETERS ### 14 | ###################################################### 15 | 16 | c = { 17 | # GENERAL STUFF 18 | 'suffix': f'{model.name}_unconditional_hint-4-full', # identifier for trained models and outputs 19 | 'device': 'cuda', # 'cuda' for GPU, 'cpu' for CPU 20 | 'interactive_visualization': True, # requires visdom package to be installed 21 | 22 | # DATA 23 | 'ndim_x': n_parameters, 24 | 'ndim_y': 0, 25 | 'ndim_z': n_parameters, 26 | 'data_model': model(), 27 | 'vis_y_target': None, 28 | 29 | # MODEL ARCHITECTURE 30 | 'n_blocks': 4, 31 | 'hidden_layer_sizes': 263, # 299, # 2M 32 | 'init_scale': 0.005, 33 | 34 | # TRAINING HYPERPARAMETERS 35 | 'n_epochs': 50, # total number of epochs 36 | 'max_batches_per_epoch': 100, # iterations per epoch (or if training data exhausted) 37 | 'batch_size': 10000, 38 | 'n_test': 100000, 39 | 'n_train': 1000000, 40 | 41 | 'lr_init': 0.01, # initial learning rate 42 | 'pre_low_lr': 3, # number of epochs at the start with very low learning rate 43 | 'final_decay': 0.01, # fraction of the learning rate to reach at final epoch 44 | 'l2_weight_reg': 1.86e-05, # strength of the weight regularization 45 | 'adam_betas': (0.9, 0.95), 46 | } 47 | 48 | # DATA LOADERS 49 | train_loader, test_loader = prepare_data_loaders(c['data_model'], c['n_train'], c['n_test'], c['batch_size']) 50 | c['train_loader'] = train_loader 51 | c['test_loader'] = test_loader 52 | 53 | 54 | ############################## 55 | ### MODEL ARCHITECTURE ### 56 | ############################## 57 | 58 | x_lane = [InputNode(c['ndim_x'], name='x')] 59 | 60 | for i in range(c['n_blocks']): 61 | if i > 0: 62 | x_lane.append(Node(x_lane[-1], 63 | HouseholderPerm, 64 | {'fixed': True, 'n_reflections': c['ndim_x']}, 65 | name=f'perm_{i}')) 66 | 67 | x_lane.append(Node(x_lane[-1], 68 | HierarchicalAffineCouplingBlock, 69 | {'c_internal': [c['hidden_layer_sizes'], c['hidden_layer_sizes']//2, c['hidden_layer_sizes']//4, c['hidden_layer_sizes']//8, c['hidden_layer_sizes']//8]}, 70 | name=f'hac_{i+1}')) 71 | 72 | x_lane.append(OutputNode(x_lane[-1], name='z')) 73 | 74 | model = ReversibleGraphNet(x_lane, verbose=False) 75 | model.to(c['device']) 76 | model.params_trainable = list(filter(lambda p: p.requires_grad, model.parameters())) 77 | 78 | 79 | def model_inverse(test_z): 80 | return model(test_z, rev=True) 81 | 82 | 83 | c['model'] = model 84 | c['model_inverse'] = model_inverse 85 | 86 | # create namedtuple from config dictionary 87 | c = namedtuple("Configuration",c.keys())(*c.values()) 88 | -------------------------------------------------------------------------------- /configs/plus_shape/unconditional_hint_8_0_big.py: -------------------------------------------------------------------------------- 1 | from collections import namedtuple 2 | import torch 3 | 4 | from FrEIA.framework import * 5 | from FrEIA.modules import * 6 | 7 | from data import PlusShapeModel as model 8 | from data import prepare_data_loaders 9 | n_parameters = model.n_parameters 10 | 11 | 12 | ###################################################### 13 | ### TRAINING CONFIGURATION AND HYPERPARAMETERS ### 14 | ###################################################### 15 | 16 | c = { 17 | # GENERAL STUFF 18 | 'suffix': f'{model.name}_unconditional_hint-8-0-big', # identifier for trained models and outputs 19 | 'device': 'cuda', # 'cuda' for GPU, 'cpu' for CPU 20 | 'interactive_visualization': True, # requires visdom package to be installed 21 | 22 | # DATA 23 | 'ndim_x': n_parameters, 24 | 'ndim_y': 0, 25 | 'ndim_z': n_parameters, 26 | 'data_model': model(), 27 | 'vis_y_target': None, 28 | 29 | # MODEL ARCHITECTURE 30 | 'n_blocks': 8, 31 | 'hidden_layer_sizes': 512, 32 | 'init_scale': 0.005, 33 | 'recursion_depth': 0, 34 | 35 | # TRAINING HYPERPARAMETERS 36 | 'n_epochs': 50, # total number of epochs 37 | 'max_batches_per_epoch': 100, # iterations per epoch (or if training data exhausted) 38 | 'batch_size': 10000, 39 | 'n_test': 100000, 40 | 'n_train': 1000000, 41 | 42 | 'lr_init': 0.01, # initial learning rate 43 | 'pre_low_lr': 3, # number of epochs at the start with very low learning rate 44 | 'final_decay': 0.01, # fraction of the learning rate to reach at final epoch 45 | 'l2_weight_reg': 1.86e-05, # strength of the weight regularization 46 | 'adam_betas': (0.9, 0.95), 47 | } 48 | 49 | # DATA LOADERS 50 | train_loader, test_loader = prepare_data_loaders(c['data_model'], c['n_train'], c['n_test'], c['batch_size']) 51 | c['train_loader'] = train_loader 52 | c['test_loader'] = test_loader 53 | 54 | 55 | ############################## 56 | ### MODEL ARCHITECTURE ### 57 | ############################## 58 | 59 | x_lane = [InputNode(c['ndim_x'], name='x')] 60 | 61 | for i in range(c['n_blocks']): 62 | if i > 0: 63 | x_lane.append(Node(x_lane[-1], 64 | HouseholderPerm, 65 | {'fixed': False, 'n_reflections': c['ndim_x']}, 66 | name=f'perm_{i}')) 67 | 68 | x_lane.append(Node(x_lane[-1], 69 | HierarchicalAffineCouplingBlock, 70 | {'c_internal': [c['hidden_layer_sizes'], c['hidden_layer_sizes']//2, c['hidden_layer_sizes']//4, c['hidden_layer_sizes']//8], 71 | 'max_splits': c['recursion_depth']}, 72 | name=f'hac_{i+1}')) 73 | 74 | x_lane.append(OutputNode(x_lane[-1], name='z')) 75 | 76 | model = ReversibleGraphNet(x_lane, verbose=False) 77 | model.to(c['device']) 78 | model.params_trainable = list(filter(lambda p: p.requires_grad, model.parameters())) 79 | 80 | 81 | def model_inverse(test_z): 82 | return model(test_z, rev=True) 83 | 84 | 85 | c['model'] = model 86 | c['model_inverse'] = model_inverse 87 | 88 | # create namedtuple from config dictionary 89 | c = namedtuple("Configuration",c.keys())(*c.values()) 90 | -------------------------------------------------------------------------------- /configs/plus_shape/unconditional_hint_8_0_small.py: -------------------------------------------------------------------------------- 1 | from collections import namedtuple 2 | import torch 3 | 4 | from FrEIA.framework import * 5 | from FrEIA.modules import * 6 | 7 | from data import PlusShapeModel as model 8 | from data import prepare_data_loaders 9 | n_parameters = model.n_parameters 10 | 11 | 12 | ###################################################### 13 | ### TRAINING CONFIGURATION AND HYPERPARAMETERS ### 14 | ###################################################### 15 | 16 | c = { 17 | # GENERAL STUFF 18 | 'suffix': f'{model.name}_unconditional_hint-8-0-small', # identifier for trained models and outputs 19 | 'device': 'cuda', # 'cuda' for GPU, 'cpu' for CPU 20 | 'interactive_visualization': True, # requires visdom package to be installed 21 | 22 | # DATA 23 | 'ndim_x': n_parameters, 24 | 'ndim_y': 0, 25 | 'ndim_z': n_parameters, 26 | 'data_model': model(), 27 | 'vis_y_target': None, 28 | 29 | # MODEL ARCHITECTURE 30 | 'n_blocks': 8, 31 | 'hidden_layer_sizes': 128, 32 | 'init_scale': 0.005, 33 | 'recursion_depth': 0, 34 | 35 | # TRAINING HYPERPARAMETERS 36 | 'n_epochs': 50, # total number of epochs 37 | 'max_batches_per_epoch': 100, # iterations per epoch (or if training data exhausted) 38 | 'batch_size': 10000, 39 | 'n_test': 100000, 40 | 'n_train': 1000000, 41 | 42 | 'lr_init': 0.01, # initial learning rate 43 | 'pre_low_lr': 3, # number of epochs at the start with very low learning rate 44 | 'final_decay': 0.01, # fraction of the learning rate to reach at final epoch 45 | 'l2_weight_reg': 1.86e-05, # strength of the weight regularization 46 | 'adam_betas': (0.9, 0.95), 47 | } 48 | 49 | # DATA LOADERS 50 | train_loader, test_loader = prepare_data_loaders(c['data_model'], c['n_train'], c['n_test'], c['batch_size']) 51 | c['train_loader'] = train_loader 52 | c['test_loader'] = test_loader 53 | 54 | 55 | ############################## 56 | ### MODEL ARCHITECTURE ### 57 | ############################## 58 | 59 | x_lane = [InputNode(c['ndim_x'], name='x')] 60 | 61 | for i in range(c['n_blocks']): 62 | if i > 0: 63 | x_lane.append(Node(x_lane[-1], 64 | HouseholderPerm, 65 | {'fixed': False, 'n_reflections': c['ndim_x']}, 66 | name=f'perm_{i}')) 67 | 68 | x_lane.append(Node(x_lane[-1], 69 | HierarchicalAffineCouplingBlock, 70 | {'c_internal': [c['hidden_layer_sizes'], c['hidden_layer_sizes']//2, c['hidden_layer_sizes']//4, c['hidden_layer_sizes']//8], 71 | 'max_splits': c['recursion_depth']}, 72 | name=f'hac_{i+1}')) 73 | 74 | x_lane.append(OutputNode(x_lane[-1], name='z')) 75 | 76 | model = ReversibleGraphNet(x_lane, verbose=False) 77 | model.to(c['device']) 78 | model.params_trainable = list(filter(lambda p: p.requires_grad, model.parameters())) 79 | 80 | 81 | def model_inverse(test_z): 82 | return model(test_z, rev=True) 83 | 84 | 85 | c['model'] = model 86 | c['model_inverse'] = model_inverse 87 | 88 | # create namedtuple from config dictionary 89 | c = namedtuple("Configuration",c.keys())(*c.values()) 90 | -------------------------------------------------------------------------------- /configs/plus_shape/unconditional_hint_8_1.py: -------------------------------------------------------------------------------- 1 | from collections import namedtuple 2 | import torch 3 | 4 | from FrEIA.framework import * 5 | from FrEIA.modules import * 6 | 7 | from data import PlusShapeModel as model 8 | from data import prepare_data_loaders 9 | n_parameters = model.n_parameters 10 | 11 | 12 | ###################################################### 13 | ### TRAINING CONFIGURATION AND HYPERPARAMETERS ### 14 | ###################################################### 15 | 16 | c = { 17 | # GENERAL STUFF 18 | 'suffix': f'{model.name}_unconditional_hint-8-1', # identifier for trained models and outputs 19 | 'device': 'cuda', # 'cuda' for GPU, 'cpu' for CPU 20 | 'interactive_visualization': True, # requires visdom package to be installed 21 | 22 | # DATA 23 | 'ndim_x': n_parameters, 24 | 'ndim_y': 0, 25 | 'ndim_z': n_parameters, 26 | 'data_model': model(), 27 | 'vis_y_target': None, 28 | 29 | # MODEL ARCHITECTURE 30 | 'n_blocks': 8, 31 | 'hidden_layer_sizes': 242, # 2M 32 | 'init_scale': 0.005, 33 | 34 | # TRAINING HYPERPARAMETERS 35 | 'n_epochs': 50, # total number of epochs 36 | 'max_batches_per_epoch': 100, # iterations per epoch (or if training data exhausted) 37 | 'batch_size': 10000, 38 | 'n_test': 100000, 39 | 'n_train': 1000000, 40 | 41 | 'lr_init': 0.01, # initial learning rate 42 | 'pre_low_lr': 3, # number of epochs at the start with very low learning rate 43 | 'final_decay': 0.01, # fraction of the learning rate to reach at final epoch 44 | 'l2_weight_reg': 1.86e-05, # strength of the weight regularization 45 | 'adam_betas': (0.9, 0.95), 46 | } 47 | 48 | # DATA LOADERS 49 | train_loader, test_loader = prepare_data_loaders(c['data_model'], c['n_train'], c['n_test'], c['batch_size']) 50 | c['train_loader'] = train_loader 51 | c['test_loader'] = test_loader 52 | 53 | 54 | ############################## 55 | ### MODEL ARCHITECTURE ### 56 | ############################## 57 | 58 | x_lane = [InputNode(c['ndim_x'], name='x')] 59 | 60 | for i in range(c['n_blocks']): 61 | if i > 0: 62 | x_lane.append(Node(x_lane[-1], 63 | HouseholderPerm, 64 | {'fixed': True, 'n_reflections': c['ndim_x']}, 65 | name=f'perm_{i}')) 66 | 67 | x_lane.append(Node(x_lane[-1], 68 | HierarchicalAffineCouplingBlock, 69 | {'c_internal': [c['hidden_layer_sizes'], c['hidden_layer_sizes']//2], 70 | 'max_splits': 1}, 71 | name=f'hac_{i+1}')) 72 | 73 | x_lane.append(OutputNode(x_lane[-1], name='z')) 74 | 75 | model = ReversibleGraphNet(x_lane, verbose=False) 76 | model.to(c['device']) 77 | model.params_trainable = list(filter(lambda p: p.requires_grad, model.parameters())) 78 | 79 | 80 | def model_inverse(test_z): 81 | return model(test_z, rev=True) 82 | 83 | 84 | c['model'] = model 85 | c['model_inverse'] = model_inverse 86 | 87 | # create namedtuple from config dictionary 88 | c = namedtuple("Configuration",c.keys())(*c.values()) 89 | -------------------------------------------------------------------------------- /configs/plus_shape/unconditional_hint_8_1_big.py: -------------------------------------------------------------------------------- 1 | from collections import namedtuple 2 | import torch 3 | 4 | from FrEIA.framework import * 5 | from FrEIA.modules import * 6 | 7 | from data import PlusShapeModel as model 8 | from data import prepare_data_loaders 9 | n_parameters = model.n_parameters 10 | 11 | 12 | ###################################################### 13 | ### TRAINING CONFIGURATION AND HYPERPARAMETERS ### 14 | ###################################################### 15 | 16 | c = { 17 | # GENERAL STUFF 18 | 'suffix': f'{model.name}_unconditional_hint-8-1-big', # identifier for trained models and outputs 19 | 'device': 'cuda', # 'cuda' for GPU, 'cpu' for CPU 20 | 'interactive_visualization': True, # requires visdom package to be installed 21 | 22 | # DATA 23 | 'ndim_x': n_parameters, 24 | 'ndim_y': 0, 25 | 'ndim_z': n_parameters, 26 | 'data_model': model(), 27 | 'vis_y_target': None, 28 | 29 | # MODEL ARCHITECTURE 30 | 'n_blocks': 8, 31 | 'hidden_layer_sizes': 512, 32 | 'init_scale': 0.005, 33 | 'recursion_depth': 1, 34 | 35 | # TRAINING HYPERPARAMETERS 36 | 'n_epochs': 50, # total number of epochs 37 | 'max_batches_per_epoch': 100, # iterations per epoch (or if training data exhausted) 38 | 'batch_size': 10000, 39 | 'n_test': 100000, 40 | 'n_train': 1000000, 41 | 42 | 'lr_init': 0.01, # initial learning rate 43 | 'pre_low_lr': 3, # number of epochs at the start with very low learning rate 44 | 'final_decay': 0.01, # fraction of the learning rate to reach at final epoch 45 | 'l2_weight_reg': 1.86e-05, # strength of the weight regularization 46 | 'adam_betas': (0.9, 0.95), 47 | } 48 | 49 | # DATA LOADERS 50 | train_loader, test_loader = prepare_data_loaders(c['data_model'], c['n_train'], c['n_test'], c['batch_size']) 51 | c['train_loader'] = train_loader 52 | c['test_loader'] = test_loader 53 | 54 | 55 | ############################## 56 | ### MODEL ARCHITECTURE ### 57 | ############################## 58 | 59 | x_lane = [InputNode(c['ndim_x'], name='x')] 60 | 61 | for i in range(c['n_blocks']): 62 | if i > 0: 63 | x_lane.append(Node(x_lane[-1], 64 | HouseholderPerm, 65 | {'fixed': False, 'n_reflections': c['ndim_x']}, 66 | name=f'perm_{i}')) 67 | 68 | x_lane.append(Node(x_lane[-1], 69 | HierarchicalAffineCouplingBlock, 70 | {'c_internal': [c['hidden_layer_sizes'], c['hidden_layer_sizes']//2, c['hidden_layer_sizes']//4, c['hidden_layer_sizes']//8], 71 | 'max_splits': c['recursion_depth']}, 72 | name=f'hac_{i+1}')) 73 | 74 | x_lane.append(OutputNode(x_lane[-1], name='z')) 75 | 76 | model = ReversibleGraphNet(x_lane, verbose=False) 77 | model.to(c['device']) 78 | model.params_trainable = list(filter(lambda p: p.requires_grad, model.parameters())) 79 | 80 | 81 | def model_inverse(test_z): 82 | return model(test_z, rev=True) 83 | 84 | 85 | c['model'] = model 86 | c['model_inverse'] = model_inverse 87 | 88 | # create namedtuple from config dictionary 89 | c = namedtuple("Configuration",c.keys())(*c.values()) 90 | -------------------------------------------------------------------------------- /configs/plus_shape/unconditional_hint_8_1_small.py: -------------------------------------------------------------------------------- 1 | from collections import namedtuple 2 | import torch 3 | 4 | from FrEIA.framework import * 5 | from FrEIA.modules import * 6 | 7 | from data import PlusShapeModel as model 8 | from data import prepare_data_loaders 9 | n_parameters = model.n_parameters 10 | 11 | 12 | ###################################################### 13 | ### TRAINING CONFIGURATION AND HYPERPARAMETERS ### 14 | ###################################################### 15 | 16 | c = { 17 | # GENERAL STUFF 18 | 'suffix': f'{model.name}_unconditional_hint-8-1-small', # identifier for trained models and outputs 19 | 'device': 'cuda', # 'cuda' for GPU, 'cpu' for CPU 20 | 'interactive_visualization': True, # requires visdom package to be installed 21 | 22 | # DATA 23 | 'ndim_x': n_parameters, 24 | 'ndim_y': 0, 25 | 'ndim_z': n_parameters, 26 | 'data_model': model(), 27 | 'vis_y_target': None, 28 | 29 | # MODEL ARCHITECTURE 30 | 'n_blocks': 8, 31 | 'hidden_layer_sizes': 128, 32 | 'init_scale': 0.005, 33 | 'recursion_depth': 1, 34 | 35 | # TRAINING HYPERPARAMETERS 36 | 'n_epochs': 50, # total number of epochs 37 | 'max_batches_per_epoch': 100, # iterations per epoch (or if training data exhausted) 38 | 'batch_size': 10000, 39 | 'n_test': 100000, 40 | 'n_train': 1000000, 41 | 42 | 'lr_init': 0.01, # initial learning rate 43 | 'pre_low_lr': 3, # number of epochs at the start with very low learning rate 44 | 'final_decay': 0.01, # fraction of the learning rate to reach at final epoch 45 | 'l2_weight_reg': 1.86e-05, # strength of the weight regularization 46 | 'adam_betas': (0.9, 0.95), 47 | } 48 | 49 | # DATA LOADERS 50 | train_loader, test_loader = prepare_data_loaders(c['data_model'], c['n_train'], c['n_test'], c['batch_size']) 51 | c['train_loader'] = train_loader 52 | c['test_loader'] = test_loader 53 | 54 | 55 | ############################## 56 | ### MODEL ARCHITECTURE ### 57 | ############################## 58 | 59 | x_lane = [InputNode(c['ndim_x'], name='x')] 60 | 61 | for i in range(c['n_blocks']): 62 | if i > 0: 63 | x_lane.append(Node(x_lane[-1], 64 | HouseholderPerm, 65 | {'fixed': False, 'n_reflections': c['ndim_x']}, 66 | name=f'perm_{i}')) 67 | 68 | x_lane.append(Node(x_lane[-1], 69 | HierarchicalAffineCouplingBlock, 70 | {'c_internal': [c['hidden_layer_sizes'], c['hidden_layer_sizes']//2, c['hidden_layer_sizes']//4, c['hidden_layer_sizes']//8], 71 | 'max_splits': c['recursion_depth']}, 72 | name=f'hac_{i+1}')) 73 | 74 | x_lane.append(OutputNode(x_lane[-1], name='z')) 75 | 76 | model = ReversibleGraphNet(x_lane, verbose=False) 77 | model.to(c['device']) 78 | model.params_trainable = list(filter(lambda p: p.requires_grad, model.parameters())) 79 | 80 | 81 | def model_inverse(test_z): 82 | return model(test_z, rev=True) 83 | 84 | 85 | c['model'] = model 86 | c['model_inverse'] = model_inverse 87 | 88 | # create namedtuple from config dictionary 89 | c = namedtuple("Configuration",c.keys())(*c.values()) 90 | -------------------------------------------------------------------------------- /configs/plus_shape/unconditional_hint_8_2.py: -------------------------------------------------------------------------------- 1 | from collections import namedtuple 2 | import torch 3 | 4 | from FrEIA.framework import * 5 | from FrEIA.modules import * 6 | 7 | from data import PlusShapeModel as model 8 | from data import prepare_data_loaders 9 | n_parameters = model.n_parameters 10 | 11 | 12 | ###################################################### 13 | ### TRAINING CONFIGURATION AND HYPERPARAMETERS ### 14 | ###################################################### 15 | 16 | c = { 17 | # GENERAL STUFF 18 | 'suffix': f'{model.name}_unconditional_hint-8-2', # identifier for trained models and outputs 19 | 'device': 'cuda', # 'cuda' for GPU, 'cpu' for CPU 20 | 'interactive_visualization': True, # requires visdom package to be installed 21 | 22 | # DATA 23 | 'ndim_x': n_parameters, 24 | 'ndim_y': 0, 25 | 'ndim_z': n_parameters, 26 | 'data_model': model(), 27 | 'vis_y_target': None, 28 | 29 | # MODEL ARCHITECTURE 30 | 'n_blocks': 8, 31 | 'hidden_layer_sizes': 220, # 2M 32 | 'init_scale': 0.005, 33 | 34 | # TRAINING HYPERPARAMETERS 35 | 'n_epochs': 50, # total number of epochs 36 | 'max_batches_per_epoch': 100, # iterations per epoch (or if training data exhausted) 37 | 'batch_size': 10000, 38 | 'n_test': 100000, 39 | 'n_train': 1000000, 40 | 41 | 'lr_init': 0.01, # initial learning rate 42 | 'pre_low_lr': 3, # number of epochs at the start with very low learning rate 43 | 'final_decay': 0.01, # fraction of the learning rate to reach at final epoch 44 | 'l2_weight_reg': 1.86e-05, # strength of the weight regularization 45 | 'adam_betas': (0.9, 0.95), 46 | } 47 | 48 | # DATA LOADERS 49 | train_loader, test_loader = prepare_data_loaders(c['data_model'], c['n_train'], c['n_test'], c['batch_size']) 50 | c['train_loader'] = train_loader 51 | c['test_loader'] = test_loader 52 | 53 | 54 | ############################## 55 | ### MODEL ARCHITECTURE ### 56 | ############################## 57 | 58 | x_lane = [InputNode(c['ndim_x'], name='x')] 59 | 60 | for i in range(c['n_blocks']): 61 | if i > 0: 62 | x_lane.append(Node(x_lane[-1], 63 | HouseholderPerm, 64 | {'fixed': True, 'n_reflections': c['ndim_x']}, 65 | name=f'perm_{i}')) 66 | 67 | x_lane.append(Node(x_lane[-1], 68 | HierarchicalAffineCouplingBlock, 69 | {'c_internal': [c['hidden_layer_sizes'], c['hidden_layer_sizes']//2, c['hidden_layer_sizes']//4], 70 | 'max_splits': 2}, 71 | name=f'hac_{i+1}')) 72 | 73 | x_lane.append(OutputNode(x_lane[-1], name='z')) 74 | 75 | model = ReversibleGraphNet(x_lane, verbose=False) 76 | model.to(c['device']) 77 | model.params_trainable = list(filter(lambda p: p.requires_grad, model.parameters())) 78 | 79 | 80 | def model_inverse(test_z): 81 | return model(test_z, rev=True) 82 | 83 | 84 | c['model'] = model 85 | c['model_inverse'] = model_inverse 86 | 87 | # create namedtuple from config dictionary 88 | c = namedtuple("Configuration",c.keys())(*c.values()) 89 | -------------------------------------------------------------------------------- /configs/plus_shape/unconditional_hint_8_2_big.py: -------------------------------------------------------------------------------- 1 | from collections import namedtuple 2 | import torch 3 | 4 | from FrEIA.framework import * 5 | from FrEIA.modules import * 6 | 7 | from data import PlusShapeModel as model 8 | from data import prepare_data_loaders 9 | n_parameters = model.n_parameters 10 | 11 | 12 | ###################################################### 13 | ### TRAINING CONFIGURATION AND HYPERPARAMETERS ### 14 | ###################################################### 15 | 16 | c = { 17 | # GENERAL STUFF 18 | 'suffix': f'{model.name}_unconditional_hint-8-2-big', # identifier for trained models and outputs 19 | 'device': 'cuda', # 'cuda' for GPU, 'cpu' for CPU 20 | 'interactive_visualization': True, # requires visdom package to be installed 21 | 22 | # DATA 23 | 'ndim_x': n_parameters, 24 | 'ndim_y': 0, 25 | 'ndim_z': n_parameters, 26 | 'data_model': model(), 27 | 'vis_y_target': None, 28 | 29 | # MODEL ARCHITECTURE 30 | 'n_blocks': 8, 31 | 'hidden_layer_sizes': 512, 32 | 'init_scale': 0.005, 33 | 'recursion_depth': 2, 34 | 35 | # TRAINING HYPERPARAMETERS 36 | 'n_epochs': 50, # total number of epochs 37 | 'max_batches_per_epoch': 100, # iterations per epoch (or if training data exhausted) 38 | 'batch_size': 10000, 39 | 'n_test': 100000, 40 | 'n_train': 1000000, 41 | 42 | 'lr_init': 0.01, # initial learning rate 43 | 'pre_low_lr': 3, # number of epochs at the start with very low learning rate 44 | 'final_decay': 0.01, # fraction of the learning rate to reach at final epoch 45 | 'l2_weight_reg': 1.86e-05, # strength of the weight regularization 46 | 'adam_betas': (0.9, 0.95), 47 | } 48 | 49 | # DATA LOADERS 50 | train_loader, test_loader = prepare_data_loaders(c['data_model'], c['n_train'], c['n_test'], c['batch_size']) 51 | c['train_loader'] = train_loader 52 | c['test_loader'] = test_loader 53 | 54 | 55 | ############################## 56 | ### MODEL ARCHITECTURE ### 57 | ############################## 58 | 59 | x_lane = [InputNode(c['ndim_x'], name='x')] 60 | 61 | for i in range(c['n_blocks']): 62 | if i > 0: 63 | x_lane.append(Node(x_lane[-1], 64 | HouseholderPerm, 65 | {'fixed': False, 'n_reflections': c['ndim_x']}, 66 | name=f'perm_{i}')) 67 | 68 | x_lane.append(Node(x_lane[-1], 69 | HierarchicalAffineCouplingBlock, 70 | {'c_internal': [c['hidden_layer_sizes'], c['hidden_layer_sizes']//2, c['hidden_layer_sizes']//4, c['hidden_layer_sizes']//8], 71 | 'max_splits': c['recursion_depth']}, 72 | name=f'hac_{i+1}')) 73 | 74 | x_lane.append(OutputNode(x_lane[-1], name='z')) 75 | 76 | model = ReversibleGraphNet(x_lane, verbose=False) 77 | model.to(c['device']) 78 | model.params_trainable = list(filter(lambda p: p.requires_grad, model.parameters())) 79 | 80 | 81 | def model_inverse(test_z): 82 | return model(test_z, rev=True) 83 | 84 | 85 | c['model'] = model 86 | c['model_inverse'] = model_inverse 87 | 88 | # create namedtuple from config dictionary 89 | c = namedtuple("Configuration",c.keys())(*c.values()) 90 | -------------------------------------------------------------------------------- /configs/plus_shape/unconditional_hint_8_2_small.py: -------------------------------------------------------------------------------- 1 | from collections import namedtuple 2 | import torch 3 | 4 | from FrEIA.framework import * 5 | from FrEIA.modules import * 6 | 7 | from data import PlusShapeModel as model 8 | from data import prepare_data_loaders 9 | n_parameters = model.n_parameters 10 | 11 | 12 | ###################################################### 13 | ### TRAINING CONFIGURATION AND HYPERPARAMETERS ### 14 | ###################################################### 15 | 16 | c = { 17 | # GENERAL STUFF 18 | 'suffix': f'{model.name}_unconditional_hint-8-2-small', # identifier for trained models and outputs 19 | 'device': 'cuda', # 'cuda' for GPU, 'cpu' for CPU 20 | 'interactive_visualization': True, # requires visdom package to be installed 21 | 22 | # DATA 23 | 'ndim_x': n_parameters, 24 | 'ndim_y': 0, 25 | 'ndim_z': n_parameters, 26 | 'data_model': model(), 27 | 'vis_y_target': None, 28 | 29 | # MODEL ARCHITECTURE 30 | 'n_blocks': 8, 31 | 'hidden_layer_sizes': 128, 32 | 'init_scale': 0.005, 33 | 'recursion_depth': 2, 34 | 35 | # TRAINING HYPERPARAMETERS 36 | 'n_epochs': 50, # total number of epochs 37 | 'max_batches_per_epoch': 100, # iterations per epoch (or if training data exhausted) 38 | 'batch_size': 10000, 39 | 'n_test': 100000, 40 | 'n_train': 1000000, 41 | 42 | 'lr_init': 0.01, # initial learning rate 43 | 'pre_low_lr': 3, # number of epochs at the start with very low learning rate 44 | 'final_decay': 0.01, # fraction of the learning rate to reach at final epoch 45 | 'l2_weight_reg': 1.86e-05, # strength of the weight regularization 46 | 'adam_betas': (0.9, 0.95), 47 | } 48 | 49 | # DATA LOADERS 50 | train_loader, test_loader = prepare_data_loaders(c['data_model'], c['n_train'], c['n_test'], c['batch_size']) 51 | c['train_loader'] = train_loader 52 | c['test_loader'] = test_loader 53 | 54 | 55 | ############################## 56 | ### MODEL ARCHITECTURE ### 57 | ############################## 58 | 59 | x_lane = [InputNode(c['ndim_x'], name='x')] 60 | 61 | for i in range(c['n_blocks']): 62 | if i > 0: 63 | x_lane.append(Node(x_lane[-1], 64 | HouseholderPerm, 65 | {'fixed': False, 'n_reflections': c['ndim_x']}, 66 | name=f'perm_{i}')) 67 | 68 | x_lane.append(Node(x_lane[-1], 69 | HierarchicalAffineCouplingBlock, 70 | {'c_internal': [c['hidden_layer_sizes'], c['hidden_layer_sizes']//2, c['hidden_layer_sizes']//4, c['hidden_layer_sizes']//8], 71 | 'max_splits': c['recursion_depth']}, 72 | name=f'hac_{i+1}')) 73 | 74 | x_lane.append(OutputNode(x_lane[-1], name='z')) 75 | 76 | model = ReversibleGraphNet(x_lane, verbose=False) 77 | model.to(c['device']) 78 | model.params_trainable = list(filter(lambda p: p.requires_grad, model.parameters())) 79 | 80 | 81 | def model_inverse(test_z): 82 | return model(test_z, rev=True) 83 | 84 | 85 | c['model'] = model 86 | c['model_inverse'] = model_inverse 87 | 88 | # create namedtuple from config dictionary 89 | c = namedtuple("Configuration",c.keys())(*c.values()) 90 | -------------------------------------------------------------------------------- /configs/plus_shape/unconditional_hint_8_full.py: -------------------------------------------------------------------------------- 1 | from collections import namedtuple 2 | import torch 3 | 4 | from FrEIA.framework import * 5 | from FrEIA.modules import * 6 | 7 | from data import PlusShapeModel as model 8 | from data import prepare_data_loaders 9 | n_parameters = model.n_parameters 10 | 11 | 12 | ###################################################### 13 | ### TRAINING CONFIGURATION AND HYPERPARAMETERS ### 14 | ###################################################### 15 | 16 | c = { 17 | # GENERAL STUFF 18 | 'suffix': f'{model.name}_unconditional_hint-8-full', # identifier for trained models and outputs 19 | 'device': 'cuda', # 'cuda' for GPU, 'cpu' for CPU 20 | 'interactive_visualization': True, # requires visdom package to be installed 21 | 22 | # DATA 23 | 'ndim_x': n_parameters, 24 | 'ndim_y': 0, 25 | 'ndim_z': n_parameters, 26 | 'data_model': model(), 27 | 'vis_y_target': None, 28 | 29 | # MODEL ARCHITECTURE 30 | 'n_blocks': 8, 31 | 'hidden_layer_sizes': 176, # 2M 32 | 'init_scale': 0.005, 33 | 34 | # TRAINING HYPERPARAMETERS 35 | 'n_epochs': 50, # total number of epochs 36 | 'max_batches_per_epoch': 100, # iterations per epoch (or if training data exhausted) 37 | 'batch_size': 10000, 38 | 'n_test': 100000, 39 | 'n_train': 1000000, 40 | 41 | 'lr_init': 0.01, # initial learning rate 42 | 'pre_low_lr': 3, # number of epochs at the start with very low learning rate 43 | 'final_decay': 0.01, # fraction of the learning rate to reach at final epoch 44 | 'l2_weight_reg': 1.86e-05, # strength of the weight regularization 45 | 'adam_betas': (0.9, 0.95), 46 | } 47 | 48 | # DATA LOADERS 49 | train_loader, test_loader = prepare_data_loaders(c['data_model'], c['n_train'], c['n_test'], c['batch_size']) 50 | c['train_loader'] = train_loader 51 | c['test_loader'] = test_loader 52 | 53 | 54 | ############################## 55 | ### MODEL ARCHITECTURE ### 56 | ############################## 57 | 58 | x_lane = [InputNode(c['ndim_x'], name='x')] 59 | 60 | for i in range(c['n_blocks']): 61 | if i > 0: 62 | x_lane.append(Node(x_lane[-1], 63 | HouseholderPerm, 64 | {'fixed': True, 'n_reflections': c['ndim_x']}, 65 | name=f'perm_{i}')) 66 | 67 | x_lane.append(Node(x_lane[-1], 68 | HierarchicalAffineCouplingBlock, 69 | {'c_internal': [c['hidden_layer_sizes'], c['hidden_layer_sizes']//2, c['hidden_layer_sizes']//4, c['hidden_layer_sizes']//8, c['hidden_layer_sizes']//8]}, 70 | name=f'hac_{i+1}')) 71 | 72 | x_lane.append(OutputNode(x_lane[-1], name='z')) 73 | 74 | model = ReversibleGraphNet(x_lane, verbose=False) 75 | model.to(c['device']) 76 | model.params_trainable = list(filter(lambda p: p.requires_grad, model.parameters())) 77 | 78 | 79 | def model_inverse(test_z): 80 | return model(test_z, rev=True) 81 | 82 | 83 | c['model'] = model 84 | c['model_inverse'] = model_inverse 85 | 86 | # create namedtuple from config dictionary 87 | c = namedtuple("Configuration",c.keys())(*c.values()) 88 | -------------------------------------------------------------------------------- /configs/plus_shape/unconditional_inn_1.py: -------------------------------------------------------------------------------- 1 | from collections import namedtuple 2 | import torch 3 | 4 | from FrEIA.framework import * 5 | from FrEIA.modules import * 6 | 7 | from data import PlusShapeModel as model 8 | from data import prepare_data_loaders 9 | n_parameters = model.n_parameters 10 | 11 | 12 | ###################################################### 13 | ### TRAINING CONFIGURATION AND HYPERPARAMETERS ### 14 | ###################################################### 15 | 16 | c = { 17 | # GENERAL STUFF 18 | 'suffix': f'{model.name}_unconditional_inn-1', # identifier for trained models and outputs 19 | 'device': 'cuda', # 'cuda' for GPU, 'cpu' for CPU 20 | 'interactive_visualization': True, # requires visdom package to be installed 21 | 22 | # DATA 23 | 'ndim_x': n_parameters, 24 | 'ndim_y': 0, 25 | 'ndim_z': n_parameters, 26 | 'data_model': model(), 27 | 'vis_y_target': None, 28 | 29 | # MODEL ARCHITECTURE 30 | 'n_blocks': 1, 31 | 'hidden_layer_sizes': 325, # 500k 32 | 'init_scale': 0.005, 33 | 34 | # TRAINING HYPERPARAMETERS 35 | 'n_epochs': 50, # total number of epochs 36 | 'max_batches_per_epoch': 100, # iterations per epoch (or if training data exhausted) 37 | 'batch_size': 10000, 38 | 'n_test': 100000, 39 | 'n_train': 1000000, 40 | 41 | 'lr_init': 0.01, # initial learning rate 42 | 'pre_low_lr': 3, # number of epochs at the start with very low learning rate 43 | 'final_decay': 0.01, # fraction of the learning rate to reach at final epoch 44 | 'l2_weight_reg': 1.86e-05, # strength of the weight regularization 45 | 'adam_betas': (0.9, 0.95), 46 | } 47 | 48 | # DATA LOADERS 49 | train_loader, test_loader = prepare_data_loaders(c['data_model'], c['n_train'], c['n_test'], c['batch_size']) 50 | c['train_loader'] = train_loader 51 | c['test_loader'] = test_loader 52 | 53 | 54 | ############################## 55 | ### MODEL ARCHITECTURE ### 56 | ############################## 57 | 58 | nodes = [InputNode(c['ndim_x'], name='x')] 59 | 60 | for i in range(c['n_blocks']): 61 | nodes.append(Node(nodes[-1], 62 | HouseholderPerm, 63 | {'fixed': False, 'n_reflections': c['ndim_x']}, 64 | name=f'perm_{i+1}')) 65 | nodes.append(Node(nodes[-1], 66 | AffineCoupling, 67 | {'F_class': F_fully_connected, 68 | 'F_args': {'internal_size': c['hidden_layer_sizes']}}, 69 | name=f'ac_{i+1}')) 70 | 71 | nodes.append(OutputNode(nodes[-1], name='z')) 72 | 73 | model = ReversibleGraphNet(nodes, verbose=False) 74 | model.to(c['device']) 75 | model.params_trainable = list(filter(lambda p: p.requires_grad, model.parameters())) 76 | 77 | 78 | def model_inverse(test_z): 79 | return model([test_z], rev=True) 80 | 81 | 82 | c['model'] = model 83 | c['model_inverse'] = model_inverse 84 | 85 | # create namedtuple from config dictionary 86 | c = namedtuple("Configuration",c.keys())(*c.values()) 87 | -------------------------------------------------------------------------------- /configs/plus_shape/unconditional_inn_16.py: -------------------------------------------------------------------------------- 1 | from collections import namedtuple 2 | import torch 3 | 4 | from FrEIA.framework import * 5 | from FrEIA.modules import * 6 | 7 | from data import PlusShapeModel as model 8 | from data import prepare_data_loaders 9 | n_parameters = model.n_parameters 10 | 11 | 12 | ###################################################### 13 | ### TRAINING CONFIGURATION AND HYPERPARAMETERS ### 14 | ###################################################### 15 | 16 | c = { 17 | # GENERAL STUFF 18 | 'suffix': f'{model.name}_unconditional_inn-16', # identifier for trained models and outputs 19 | 'device': 'cuda', # 'cuda' for GPU, 'cpu' for CPU 20 | 'interactive_visualization': True, # requires visdom package to be installed 21 | 22 | # DATA 23 | 'ndim_x': n_parameters, 24 | 'ndim_y': 0, 25 | 'ndim_z': n_parameters, 26 | 'data_model': model(), 27 | 'vis_y_target': None, 28 | 29 | # MODEL ARCHITECTURE 30 | 'n_blocks': 16, 31 | 'hidden_layer_sizes': 153, # 146, # 2M 32 | 'init_scale': 0.005, 33 | 34 | # TRAINING HYPERPARAMETERS 35 | 'n_epochs': 50, # total number of epochs 36 | 'max_batches_per_epoch': 100, # iterations per epoch (or if training data exhausted) 37 | 'batch_size': 10000, 38 | 'n_test': 100000, 39 | 'n_train': 1000000, 40 | 41 | 'lr_init': 0.01, # initial learning rate 42 | 'pre_low_lr': 3, # number of epochs at the start with very low learning rate 43 | 'final_decay': 0.01, # fraction of the learning rate to reach at final epoch 44 | 'l2_weight_reg': 1.86e-05, # strength of the weight regularization 45 | 'adam_betas': (0.9, 0.95), 46 | } 47 | 48 | # DATA LOADERS 49 | train_loader, test_loader = prepare_data_loaders(c['data_model'], c['n_train'], c['n_test'], c['batch_size']) 50 | c['train_loader'] = train_loader 51 | c['test_loader'] = test_loader 52 | 53 | 54 | ############################## 55 | ### MODEL ARCHITECTURE ### 56 | ############################## 57 | 58 | nodes = [InputNode(c['ndim_x'], name='x')] 59 | 60 | for i in range(c['n_blocks']): 61 | nodes.append(Node(nodes[-1], 62 | HouseholderPerm, 63 | {'fixed': True, 'n_reflections': c['ndim_x']}, 64 | name=f'perm_{i+1}')) 65 | nodes.append(Node(nodes[-1], 66 | AffineCoupling, 67 | {'F_class': F_fully_connected, 68 | 'F_args': {'internal_size': c['hidden_layer_sizes']}}, 69 | name=f'ac_{i+1}')) 70 | 71 | nodes.append(OutputNode(nodes[-1], name='z')) 72 | 73 | model = ReversibleGraphNet(nodes, verbose=False) 74 | model.to(c['device']) 75 | model.params_trainable = list(filter(lambda p: p.requires_grad, model.parameters())) 76 | 77 | 78 | def model_inverse(test_z): 79 | return model([test_z], rev=True) 80 | 81 | 82 | c['model'] = model 83 | c['model_inverse'] = model_inverse 84 | 85 | # create namedtuple from config dictionary 86 | c = namedtuple("Configuration",c.keys())(*c.values()) 87 | -------------------------------------------------------------------------------- /configs/plus_shape/unconditional_inn_2.py: -------------------------------------------------------------------------------- 1 | from collections import namedtuple 2 | import torch 3 | 4 | from FrEIA.framework import * 5 | from FrEIA.modules import * 6 | 7 | from data import PlusShapeModel as model 8 | from data import prepare_data_loaders 9 | n_parameters = model.n_parameters 10 | 11 | 12 | ###################################################### 13 | ### TRAINING CONFIGURATION AND HYPERPARAMETERS ### 14 | ###################################################### 15 | 16 | c = { 17 | # GENERAL STUFF 18 | 'suffix': f'{model.name}_unconditional_inn-2', # identifier for trained models and outputs 19 | 'device': 'cuda', # 'cuda' for GPU, 'cpu' for CPU 20 | 'interactive_visualization': True, # requires visdom package to be installed 21 | 22 | # DATA 23 | 'ndim_x': n_parameters, 24 | 'ndim_y': 0, 25 | 'ndim_z': n_parameters, 26 | 'data_model': model(), 27 | 'vis_y_target': None, 28 | 29 | # MODEL ARCHITECTURE 30 | 'n_blocks': 2, 31 | 'hidden_layer_sizes': 220, # 500k 32 | 'init_scale': 0.005, 33 | 34 | # TRAINING HYPERPARAMETERS 35 | 'n_epochs': 50, # total number of epochs 36 | 'max_batches_per_epoch': 100, # iterations per epoch (or if training data exhausted) 37 | 'batch_size': 10000, 38 | 'n_test': 100000, 39 | 'n_train': 1000000, 40 | 41 | 'lr_init': 0.01, # initial learning rate 42 | 'pre_low_lr': 3, # number of epochs at the start with very low learning rate 43 | 'final_decay': 0.01, # fraction of the learning rate to reach at final epoch 44 | 'l2_weight_reg': 1.86e-05, # strength of the weight regularization 45 | 'adam_betas': (0.9, 0.95), 46 | } 47 | 48 | # DATA LOADERS 49 | train_loader, test_loader = prepare_data_loaders(c['data_model'], c['n_train'], c['n_test'], c['batch_size']) 50 | c['train_loader'] = train_loader 51 | c['test_loader'] = test_loader 52 | 53 | 54 | ############################## 55 | ### MODEL ARCHITECTURE ### 56 | ############################## 57 | 58 | nodes = [InputNode(c['ndim_x'], name='x')] 59 | 60 | for i in range(c['n_blocks']): 61 | nodes.append(Node(nodes[-1], 62 | HouseholderPerm, 63 | {'fixed': False, 'n_reflections': c['ndim_x']}, 64 | name=f'perm_{i+1}')) 65 | nodes.append(Node(nodes[-1], 66 | AffineCoupling, 67 | {'F_class': F_fully_connected, 68 | 'F_args': {'internal_size': c['hidden_layer_sizes']}}, 69 | name=f'ac_{i+1}')) 70 | 71 | nodes.append(OutputNode(nodes[-1], name='z')) 72 | 73 | model = ReversibleGraphNet(nodes, verbose=False) 74 | model.to(c['device']) 75 | model.params_trainable = list(filter(lambda p: p.requires_grad, model.parameters())) 76 | 77 | 78 | def model_inverse(test_z): 79 | return model([test_z], rev=True) 80 | 81 | 82 | c['model'] = model 83 | c['model_inverse'] = model_inverse 84 | 85 | # create namedtuple from config dictionary 86 | c = namedtuple("Configuration",c.keys())(*c.values()) 87 | -------------------------------------------------------------------------------- /configs/plus_shape/unconditional_inn_32.py: -------------------------------------------------------------------------------- 1 | from collections import namedtuple 2 | import torch 3 | 4 | from FrEIA.framework import * 5 | from FrEIA.modules import * 6 | 7 | from data import PlusShapeModel as model 8 | from data import prepare_data_loaders 9 | n_parameters = model.n_parameters 10 | 11 | 12 | ###################################################### 13 | ### TRAINING CONFIGURATION AND HYPERPARAMETERS ### 14 | ###################################################### 15 | 16 | c = { 17 | # GENERAL STUFF 18 | 'suffix': f'{model.name}_unconditional_inn-32', # identifier for trained models and outputs 19 | 'device': 'cuda', # 'cuda' for GPU, 'cpu' for CPU 20 | 'interactive_visualization': True, # requires visdom package to be installed 21 | 22 | # DATA 23 | 'ndim_x': n_parameters, 24 | 'ndim_y': 0, 25 | 'ndim_z': n_parameters, 26 | 'data_model': model(), 27 | 'vis_y_target': None, 28 | 29 | # MODEL ARCHITECTURE 30 | 'n_blocks': 32, 31 | 'hidden_layer_sizes': 102, # 2M 32 | 'init_scale': 0.005, 33 | 34 | # TRAINING HYPERPARAMETERS 35 | 'n_epochs': 50, # total number of epochs 36 | 'max_batches_per_epoch': 100, # iterations per epoch (or if training data exhausted) 37 | 'batch_size': 10000, 38 | 'n_test': 100000, 39 | 'n_train': 1000000, 40 | 41 | 'lr_init': 0.01, # initial learning rate 42 | 'pre_low_lr': 3, # number of epochs at the start with very low learning rate 43 | 'final_decay': 0.01, # fraction of the learning rate to reach at final epoch 44 | 'l2_weight_reg': 1.86e-05, # strength of the weight regularization 45 | 'adam_betas': (0.9, 0.95), 46 | } 47 | 48 | # DATA LOADERS 49 | train_loader, test_loader = prepare_data_loaders(c['data_model'], c['n_train'], c['n_test'], c['batch_size']) 50 | c['train_loader'] = train_loader 51 | c['test_loader'] = test_loader 52 | 53 | 54 | ############################## 55 | ### MODEL ARCHITECTURE ### 56 | ############################## 57 | 58 | nodes = [InputNode(c['ndim_x'], name='x')] 59 | 60 | for i in range(c['n_blocks']): 61 | nodes.append(Node(nodes[-1], 62 | HouseholderPerm, 63 | {'fixed': True, 'n_reflections': c['ndim_x']}, 64 | name=f'perm_{i+1}')) 65 | nodes.append(Node(nodes[-1], 66 | AffineCoupling, 67 | {'F_class': F_fully_connected, 68 | 'F_args': {'internal_size': c['hidden_layer_sizes']}}, 69 | name=f'ac_{i+1}')) 70 | 71 | nodes.append(OutputNode(nodes[-1], name='z')) 72 | 73 | model = ReversibleGraphNet(nodes, verbose=False) 74 | model.to(c['device']) 75 | model.params_trainable = list(filter(lambda p: p.requires_grad, model.parameters())) 76 | 77 | 78 | def model_inverse(test_z): 79 | return model([test_z], rev=True) 80 | 81 | 82 | c['model'] = model 83 | c['model_inverse'] = model_inverse 84 | 85 | # create namedtuple from config dictionary 86 | c = namedtuple("Configuration",c.keys())(*c.values()) 87 | -------------------------------------------------------------------------------- /configs/plus_shape/unconditional_inn_4.py: -------------------------------------------------------------------------------- 1 | from collections import namedtuple 2 | import torch 3 | 4 | from FrEIA.framework import * 5 | from FrEIA.modules import * 6 | 7 | from data import PlusShapeModel as model 8 | from data import prepare_data_loaders 9 | n_parameters = model.n_parameters 10 | 11 | 12 | ###################################################### 13 | ### TRAINING CONFIGURATION AND HYPERPARAMETERS ### 14 | ###################################################### 15 | 16 | c = { 17 | # GENERAL STUFF 18 | 'suffix': f'{model.name}_unconditional_inn-4', # identifier for trained models and outputs 19 | 'device': 'cuda', # 'cuda' for GPU, 'cpu' for CPU 20 | 'interactive_visualization': True, # requires visdom package to be installed 21 | 22 | # DATA 23 | 'ndim_x': n_parameters, 24 | 'ndim_y': 0, 25 | 'ndim_z': n_parameters, 26 | 'data_model': model(), 27 | 'vis_y_target': None, 28 | 29 | # MODEL ARCHITECTURE 30 | 'n_blocks': 4, 31 | 'hidden_layer_sizes': 329, # 326, # 2M 32 | 'init_scale': 0.005, 33 | 34 | # TRAINING HYPERPARAMETERS 35 | 'n_epochs': 50, # total number of epochs 36 | 'max_batches_per_epoch': 100, # iterations per epoch (or if training data exhausted) 37 | 'batch_size': 10000, 38 | 'n_test': 100000, 39 | 'n_train': 1000000, 40 | 41 | 'lr_init': 0.01, # initial learning rate 42 | 'pre_low_lr': 3, # number of epochs at the start with very low learning rate 43 | 'final_decay': 0.01, # fraction of the learning rate to reach at final epoch 44 | 'l2_weight_reg': 1.86e-05, # strength of the weight regularization 45 | 'adam_betas': (0.9, 0.95), 46 | } 47 | 48 | # DATA LOADERS 49 | train_loader, test_loader = prepare_data_loaders(c['data_model'], c['n_train'], c['n_test'], c['batch_size']) 50 | c['train_loader'] = train_loader 51 | c['test_loader'] = test_loader 52 | 53 | 54 | ############################## 55 | ### MODEL ARCHITECTURE ### 56 | ############################## 57 | 58 | nodes = [InputNode(c['ndim_x'], name='x')] 59 | 60 | for i in range(c['n_blocks']): 61 | nodes.append(Node(nodes[-1], 62 | HouseholderPerm, 63 | {'fixed': True, 'n_reflections': c['ndim_x']}, 64 | name=f'perm_{i+1}')) 65 | nodes.append(Node(nodes[-1], 66 | AffineCoupling, 67 | {'F_class': F_fully_connected, 68 | 'F_args': {'internal_size': c['hidden_layer_sizes']}}, 69 | name=f'ac_{i+1}')) 70 | 71 | nodes.append(OutputNode(nodes[-1], name='z')) 72 | 73 | model = ReversibleGraphNet(nodes, verbose=False) 74 | model.to(c['device']) 75 | model.params_trainable = list(filter(lambda p: p.requires_grad, model.parameters())) 76 | 77 | 78 | def model_inverse(test_z): 79 | return model([test_z], rev=True) 80 | 81 | 82 | c['model'] = model 83 | c['model_inverse'] = model_inverse 84 | 85 | # create namedtuple from config dictionary 86 | c = namedtuple("Configuration",c.keys())(*c.values()) 87 | -------------------------------------------------------------------------------- /configs/plus_shape/unconditional_inn_4_Q.py: -------------------------------------------------------------------------------- 1 | from collections import namedtuple 2 | import torch 3 | 4 | from FrEIA.framework import * 5 | from FrEIA.modules import * 6 | 7 | from data import PlusShapeModel as model 8 | from data import prepare_data_loaders 9 | n_parameters = model.n_parameters 10 | 11 | 12 | ###################################################### 13 | ### TRAINING CONFIGURATION AND HYPERPARAMETERS ### 14 | ###################################################### 15 | 16 | c = { 17 | # GENERAL STUFF 18 | 'suffix': f'{model.name}_unconditional_inn-4-Q', # identifier for trained models and outputs 19 | 'device': 'cuda', # 'cuda' for GPU, 'cpu' for CPU 20 | 'interactive_visualization': True, # requires visdom package to be installed 21 | 22 | # DATA 23 | 'ndim_x': n_parameters, 24 | 'ndim_y': 0, 25 | 'ndim_z': n_parameters, 26 | 'data_model': model(), 27 | 'vis_y_target': None, 28 | 29 | # MODEL ARCHITECTURE 30 | 'n_blocks': 4, 31 | 'hidden_layer_sizes': 326, # 329, # 2M 32 | 'init_scale': 0.005, 33 | 34 | # TRAINING HYPERPARAMETERS 35 | 'n_epochs': 50, # total number of epochs 36 | 'max_batches_per_epoch': 100, # iterations per epoch (or if training data exhausted) 37 | 'batch_size': 10000, 38 | 'n_test': 100000, 39 | 'n_train': 1000000, 40 | 41 | 'lr_init': 0.01, # initial learning rate 42 | 'pre_low_lr': 3, # number of epochs at the start with very low learning rate 43 | 'final_decay': 0.01, # fraction of the learning rate to reach at final epoch 44 | 'l2_weight_reg': 1.86e-05, # strength of the weight regularization 45 | 'adam_betas': (0.9, 0.95), 46 | } 47 | 48 | # DATA LOADERS 49 | train_loader, test_loader = prepare_data_loaders(c['data_model'], c['n_train'], c['n_test'], c['batch_size']) 50 | c['train_loader'] = train_loader 51 | c['test_loader'] = test_loader 52 | 53 | 54 | ############################## 55 | ### MODEL ARCHITECTURE ### 56 | ############################## 57 | 58 | nodes = [InputNode(c['ndim_x'], name='x')] 59 | 60 | for i in range(c['n_blocks']): 61 | nodes.append(Node(nodes[-1], 62 | HouseholderPerm, 63 | {'fixed': False, 'n_reflections': c['ndim_x']}, 64 | name=f'perm_{i+1}')) 65 | nodes.append(Node(nodes[-1], 66 | AffineCoupling, 67 | {'F_class': F_fully_connected, 68 | 'F_args': {'internal_size': c['hidden_layer_sizes']}}, 69 | name=f'ac_{i+1}')) 70 | 71 | nodes.append(OutputNode(nodes[-1], name='z')) 72 | 73 | model = ReversibleGraphNet(nodes, verbose=False) 74 | model.to(c['device']) 75 | model.params_trainable = list(filter(lambda p: p.requires_grad, model.parameters())) 76 | 77 | 78 | def model_inverse(test_z): 79 | return model([test_z], rev=True) 80 | 81 | 82 | c['model'] = model 83 | c['model_inverse'] = model_inverse 84 | 85 | # create namedtuple from config dictionary 86 | c = namedtuple("Configuration",c.keys())(*c.values()) 87 | -------------------------------------------------------------------------------- /configs/plus_shape/unconditional_inn_8.py: -------------------------------------------------------------------------------- 1 | from collections import namedtuple 2 | import torch 3 | 4 | from FrEIA.framework import * 5 | from FrEIA.modules import * 6 | 7 | from data import PlusShapeModel as model 8 | from data import prepare_data_loaders 9 | n_parameters = model.n_parameters 10 | 11 | 12 | ###################################################### 13 | ### TRAINING CONFIGURATION AND HYPERPARAMETERS ### 14 | ###################################################### 15 | 16 | c = { 17 | # GENERAL STUFF 18 | 'suffix': f'{model.name}_unconditional_inn-8', # identifier for trained models and outputs 19 | 'device': 'cuda', # 'cuda' for GPU, 'cpu' for CPU 20 | 'interactive_visualization': True, # requires visdom package to be installed 21 | 22 | # DATA 23 | 'ndim_x': n_parameters, 24 | 'ndim_y': 0, 25 | 'ndim_z': n_parameters, 26 | 'data_model': model(), 27 | 'vis_y_target': None, 28 | 29 | # MODEL ARCHITECTURE 30 | 'n_blocks': 8, 31 | 'hidden_layer_sizes': 226, # 220, # 2M 32 | 'init_scale': 0.005, 33 | 34 | # TRAINING HYPERPARAMETERS 35 | 'n_epochs': 50, # total number of epochs 36 | 'max_batches_per_epoch': 100, # iterations per epoch (or if training data exhausted) 37 | 'batch_size': 10000, 38 | 'n_test': 100000, 39 | 'n_train': 1000000, 40 | 41 | 'lr_init': 0.01, # initial learning rate 42 | 'pre_low_lr': 3, # number of epochs at the start with very low learning rate 43 | 'final_decay': 0.01, # fraction of the learning rate to reach at final epoch 44 | 'l2_weight_reg': 1.86e-05, # strength of the weight regularization 45 | 'adam_betas': (0.9, 0.95), 46 | } 47 | 48 | # DATA LOADERS 49 | train_loader, test_loader = prepare_data_loaders(c['data_model'], c['n_train'], c['n_test'], c['batch_size']) 50 | c['train_loader'] = train_loader 51 | c['test_loader'] = test_loader 52 | 53 | 54 | ############################## 55 | ### MODEL ARCHITECTURE ### 56 | ############################## 57 | 58 | nodes = [InputNode(c['ndim_x'], name='x')] 59 | 60 | for i in range(c['n_blocks']): 61 | nodes.append(Node(nodes[-1], 62 | HouseholderPerm, 63 | {'fixed': True, 'n_reflections': c['ndim_x']}, 64 | name=f'perm_{i+1}')) 65 | nodes.append(Node(nodes[-1], 66 | AffineCoupling, 67 | {'F_class': F_fully_connected, 68 | 'F_args': {'internal_size': c['hidden_layer_sizes']}}, 69 | name=f'ac_{i+1}')) 70 | 71 | nodes.append(OutputNode(nodes[-1], name='z')) 72 | 73 | model = ReversibleGraphNet(nodes, verbose=False) 74 | model.to(c['device']) 75 | model.params_trainable = list(filter(lambda p: p.requires_grad, model.parameters())) 76 | 77 | 78 | def model_inverse(test_z): 79 | return model([test_z], rev=True) 80 | 81 | 82 | c['model'] = model 83 | c['model_inverse'] = model_inverse 84 | 85 | # create namedtuple from config dictionary 86 | c = namedtuple("Configuration",c.keys())(*c.values()) 87 | -------------------------------------------------------------------------------- /configs/uci_data/gas_hint_4.py: -------------------------------------------------------------------------------- 1 | from collections import namedtuple 2 | import torch 3 | 4 | from FrEIA.framework import * 5 | from FrEIA.modules import * 6 | 7 | from data import prepare_uci_loaders, Gas as model 8 | n_parameters = model.n_parameters 9 | 10 | 11 | ###################################################### 12 | ### TRAINING CONFIGURATION AND HYPERPARAMETERS ### 13 | ###################################################### 14 | 15 | c = { 16 | # GENERAL STUFF 17 | 'suffix': f'{model.name}_hint-4', # identifier for trained models and outputs 18 | 'device': 'cuda', # 'cuda' for GPU, 'cpu' for CPU 19 | 'interactive_visualization': True, # requires visdom package to be installed 20 | 21 | # DATA 22 | 'ndim_x': n_parameters, 23 | 'ndim_y': 0, 24 | 'ndim_z': n_parameters, 25 | 'data_model': model(), 26 | 'vis_y_target': None, 27 | 28 | # MODEL ARCHITECTURE 29 | 'n_blocks': 4, 30 | 'hidden_layer_sizes': 184, # 500k 31 | 'init_scale': 0.005, 32 | 33 | # TRAINING HYPERPARAMETERS 34 | 'n_epochs': 50, # total number of epochs 35 | 'max_batches_per_epoch': 1000, # iterations per epoch (or if training data exhausted) 36 | 'batch_size': 853, 37 | 38 | 'lr_init': 0.01, # initial learning rate 39 | 'pre_low_lr': 3, # number of epochs at the start with very low learning rate 40 | 'final_decay': 0.01, # fraction of the learning rate to reach at final epoch 41 | 'l2_weight_reg': 1.86e-05, # strength of the weight regularization 42 | 'adam_betas': (0.9, 0.95), 43 | } 44 | 45 | # DATA LOADERS 46 | train_loader, test_loader = prepare_uci_loaders(c['data_model'].name, c['batch_size']) 47 | c['train_loader'] = train_loader 48 | c['test_loader'] = test_loader 49 | 50 | 51 | ############################## 52 | ### MODEL ARCHITECTURE ### 53 | ############################## 54 | 55 | x_lane = [InputNode(c['ndim_x'], name='x')] 56 | 57 | for i in range(c['n_blocks']): 58 | if i > 0: 59 | x_lane.append(Node(x_lane[-1], 60 | HouseholderPerm, 61 | {'fixed': True, 'n_reflections': c['ndim_x']}, 62 | name=f'perm_{i}')) 63 | 64 | x_lane.append(Node(x_lane[-1], 65 | HierarchicalAffineCouplingBlock, 66 | {'c_internal': [c['hidden_layer_sizes'], c['hidden_layer_sizes']//2, c['hidden_layer_sizes']//4, c['hidden_layer_sizes']//8]}, 67 | name=f'hac_{i+1}')) 68 | 69 | x_lane.append(OutputNode(x_lane[-1], name='z')) 70 | 71 | model = ReversibleGraphNet(x_lane, verbose=False) 72 | model.to(c['device']) 73 | model.params_trainable = list(filter(lambda p: p.requires_grad, model.parameters())) 74 | 75 | 76 | def model_inverse(test_z): 77 | return model(test_z, rev=True) 78 | 79 | 80 | c['model'] = model 81 | c['model_inverse'] = model_inverse 82 | 83 | # create namedtuple from config dictionary 84 | c = namedtuple("Configuration",c.keys())(*c.values()) 85 | -------------------------------------------------------------------------------- /configs/uci_data/gas_hint_8.py: -------------------------------------------------------------------------------- 1 | from collections import namedtuple 2 | import torch 3 | 4 | from FrEIA.framework import * 5 | from FrEIA.modules import * 6 | 7 | from data import prepare_uci_loaders, Gas as model 8 | n_parameters = model.n_parameters 9 | 10 | 11 | ###################################################### 12 | ### TRAINING CONFIGURATION AND HYPERPARAMETERS ### 13 | ###################################################### 14 | 15 | c = { 16 | # GENERAL STUFF 17 | 'suffix': f'{model.name}_hint-8', # identifier for trained models and outputs 18 | 'device': 'cuda', # 'cuda' for GPU, 'cpu' for CPU 19 | 'interactive_visualization': True, # requires visdom package to be installed 20 | 21 | # DATA 22 | 'ndim_x': n_parameters, 23 | 'ndim_y': 0, 24 | 'ndim_z': n_parameters, 25 | 'data_model': model(), 26 | 'vis_y_target': None, 27 | 28 | # MODEL ARCHITECTURE 29 | 'n_blocks': 8, 30 | 'hidden_layer_sizes': 128, # 500k 31 | 'init_scale': 0.005, 32 | 33 | # TRAINING HYPERPARAMETERS 34 | 'n_epochs': 50, # total number of epochs 35 | 'max_batches_per_epoch': 1000, # iterations per epoch (or if training data exhausted) 36 | 'batch_size': 853, 37 | 38 | 'lr_init': 0.01, # initial learning rate 39 | 'pre_low_lr': 3, # number of epochs at the start with very low learning rate 40 | 'final_decay': 0.01, # fraction of the learning rate to reach at final epoch 41 | 'l2_weight_reg': 1.86e-05, # strength of the weight regularization 42 | 'adam_betas': (0.9, 0.95), 43 | } 44 | 45 | # DATA LOADERS 46 | train_loader, test_loader = prepare_uci_loaders(c['data_model'].name, c['batch_size']) 47 | c['train_loader'] = train_loader 48 | c['test_loader'] = test_loader 49 | 50 | 51 | ############################## 52 | ### MODEL ARCHITECTURE ### 53 | ############################## 54 | 55 | x_lane = [InputNode(c['ndim_x'], name='x')] 56 | 57 | for i in range(c['n_blocks']): 58 | if i > 0: 59 | x_lane.append(Node(x_lane[-1], 60 | HouseholderPerm, 61 | {'fixed': True, 'n_reflections': c['ndim_x']}, 62 | name=f'perm_{i}')) 63 | 64 | x_lane.append(Node(x_lane[-1], 65 | HierarchicalAffineCouplingBlock, 66 | {'c_internal': [c['hidden_layer_sizes'], c['hidden_layer_sizes']//2, c['hidden_layer_sizes']//4, c['hidden_layer_sizes']//8]}, 67 | name=f'hac_{i+1}')) 68 | 69 | x_lane.append(OutputNode(x_lane[-1], name='z')) 70 | 71 | model = ReversibleGraphNet(x_lane, verbose=False) 72 | model.to(c['device']) 73 | model.params_trainable = list(filter(lambda p: p.requires_grad, model.parameters())) 74 | 75 | 76 | def model_inverse(test_z): 77 | return model(test_z, rev=True) 78 | 79 | 80 | c['model'] = model 81 | c['model_inverse'] = model_inverse 82 | 83 | # create namedtuple from config dictionary 84 | c = namedtuple("Configuration",c.keys())(*c.values()) 85 | -------------------------------------------------------------------------------- /configs/uci_data/gas_inn_4.py: -------------------------------------------------------------------------------- 1 | from collections import namedtuple 2 | import torch 3 | 4 | from FrEIA.framework import * 5 | from FrEIA.modules import * 6 | 7 | from data import prepare_uci_loaders, Gas as model 8 | n_parameters = model.n_parameters 9 | 10 | 11 | ###################################################### 12 | ### TRAINING CONFIGURATION AND HYPERPARAMETERS ### 13 | ###################################################### 14 | 15 | c = { 16 | # GENERAL STUFF 17 | 'suffix': f'{model.name}_inn-4', # identifier for trained models and outputs 18 | 'device': 'cuda', # 'cuda' for GPU, 'cpu' for CPU 19 | 'interactive_visualization': True, # requires visdom package to be installed 20 | 21 | # DATA 22 | 'ndim_x': n_parameters, 23 | 'ndim_y': 0, 24 | 'ndim_z': n_parameters, 25 | 'data_model': model(), 26 | 'vis_y_target': None, 27 | 28 | # MODEL ARCHITECTURE 29 | 'n_blocks': 4, 30 | 'hidden_layer_sizes': 174, # 500k 31 | 'init_scale': 0.005, 32 | 33 | # TRAINING HYPERPARAMETERS 34 | 'n_epochs': 50, # total number of epochs 35 | 'max_batches_per_epoch': 1000, # iterations per epoch (or if training data exhausted) 36 | 'batch_size': 853, 37 | 38 | 'lr_init': 0.01, # initial learning rate 39 | 'pre_low_lr': 3, # number of epochs at the start with very low learning rate 40 | 'final_decay': 0.01, # fraction of the learning rate to reach at final epoch 41 | 'l2_weight_reg': 1.86e-05, # strength of the weight regularization 42 | 'adam_betas': (0.9, 0.95), 43 | } 44 | 45 | # DATA LOADERS 46 | train_loader, test_loader = prepare_uci_loaders(c['data_model'].name, c['batch_size']) 47 | c['train_loader'] = train_loader 48 | c['test_loader'] = test_loader 49 | 50 | 51 | ############################## 52 | ### MODEL ARCHITECTURE ### 53 | ############################## 54 | 55 | nodes = [InputNode(c['ndim_x'], name='x')] 56 | 57 | for i in range(c['n_blocks']): 58 | nodes.append(Node(nodes[-1], 59 | HouseholderPerm, 60 | {'fixed': True, 'n_reflections': c['ndim_x']}, 61 | name=f'perm_{i+1}')) 62 | 63 | nodes.append(Node(nodes[-1], 64 | AffineCoupling, 65 | {'F_class': F_fully_connected, 66 | 'F_args': {'internal_size': c['hidden_layer_sizes']}}, 67 | name=f'ac_{i+1}')) 68 | 69 | nodes.append(OutputNode(nodes[-1], name='z')) 70 | 71 | model = ReversibleGraphNet(nodes, verbose=False) 72 | model.to(c['device']) 73 | model.params_trainable = list(filter(lambda p: p.requires_grad, model.parameters())) 74 | 75 | 76 | def model_inverse(test_z): 77 | return model(test_z, rev=True) 78 | 79 | 80 | c['model'] = model 81 | c['model_inverse'] = model_inverse 82 | 83 | # create namedtuple from config dictionary 84 | c = namedtuple("Configuration",c.keys())(*c.values()) 85 | -------------------------------------------------------------------------------- /configs/uci_data/gas_inn_8.py: -------------------------------------------------------------------------------- 1 | from collections import namedtuple 2 | import torch 3 | 4 | from FrEIA.framework import * 5 | from FrEIA.modules import * 6 | 7 | from data import prepare_uci_loaders, Gas as model 8 | n_parameters = model.n_parameters 9 | 10 | 11 | ###################################################### 12 | ### TRAINING CONFIGURATION AND HYPERPARAMETERS ### 13 | ###################################################### 14 | 15 | c = { 16 | # GENERAL STUFF 17 | 'suffix': f'{model.name}_inn-8', # identifier for trained models and outputs 18 | 'device': 'cuda', # 'cuda' for GPU, 'cpu' for CPU 19 | 'interactive_visualization': True, # requires visdom package to be installed 20 | 21 | # DATA 22 | 'ndim_x': n_parameters, 23 | 'ndim_y': 0, 24 | 'ndim_z': n_parameters, 25 | 'data_model': model(), 26 | 'vis_y_target': None, 27 | 28 | # MODEL ARCHITECTURE 29 | 'n_blocks': 8, 30 | 'hidden_layer_sizes': 122, # 500k 31 | 'init_scale': 0.005, 32 | 33 | # TRAINING HYPERPARAMETERS 34 | 'n_epochs': 50, # total number of epochs 35 | 'max_batches_per_epoch': 1000, # iterations per epoch (or if training data exhausted) 36 | 'batch_size': 853, 37 | 38 | 'lr_init': 0.01, # initial learning rate 39 | 'pre_low_lr': 3, # number of epochs at the start with very low learning rate 40 | 'final_decay': 0.01, # fraction of the learning rate to reach at final epoch 41 | 'l2_weight_reg': 1.86e-05, # strength of the weight regularization 42 | 'adam_betas': (0.9, 0.95), 43 | } 44 | 45 | # DATA LOADERS 46 | train_loader, test_loader = prepare_uci_loaders(c['data_model'].name, c['batch_size']) 47 | c['train_loader'] = train_loader 48 | c['test_loader'] = test_loader 49 | 50 | 51 | ############################## 52 | ### MODEL ARCHITECTURE ### 53 | ############################## 54 | 55 | nodes = [InputNode(c['ndim_x'], name='x')] 56 | 57 | for i in range(c['n_blocks']): 58 | nodes.append(Node(nodes[-1], 59 | HouseholderPerm, 60 | {'fixed': True, 'n_reflections': c['ndim_x']}, 61 | name=f'perm_{i+1}')) 62 | 63 | nodes.append(Node(nodes[-1], 64 | AffineCoupling, 65 | {'F_class': F_fully_connected, 66 | 'F_args': {'internal_size': c['hidden_layer_sizes']}}, 67 | name=f'ac_{i+1}')) 68 | 69 | nodes.append(OutputNode(nodes[-1], name='z')) 70 | 71 | model = ReversibleGraphNet(nodes, verbose=False) 72 | model.to(c['device']) 73 | model.params_trainable = list(filter(lambda p: p.requires_grad, model.parameters())) 74 | 75 | 76 | def model_inverse(test_z): 77 | return model(test_z, rev=True) 78 | 79 | 80 | c['model'] = model 81 | c['model_inverse'] = model_inverse 82 | 83 | # create namedtuple from config dictionary 84 | c = namedtuple("Configuration",c.keys())(*c.values()) 85 | -------------------------------------------------------------------------------- /configs/uci_data/miniboone_hint_4.py: -------------------------------------------------------------------------------- 1 | from collections import namedtuple 2 | import torch 3 | 4 | from FrEIA.framework import * 5 | from FrEIA.modules import * 6 | 7 | from data import prepare_uci_loaders, Miniboone as model 8 | n_parameters = model.n_parameters 9 | 10 | 11 | ###################################################### 12 | ### TRAINING CONFIGURATION AND HYPERPARAMETERS ### 13 | ###################################################### 14 | 15 | c = { 16 | # GENERAL STUFF 17 | 'suffix': f'{model.name}_hint-4', # identifier for trained models and outputs 18 | 'device': 'cuda', # 'cuda' for GPU, 'cpu' for CPU 19 | 'interactive_visualization': True, # requires visdom package to be installed 20 | 21 | # DATA 22 | 'ndim_x': n_parameters, 23 | 'ndim_y': 0, 24 | 'ndim_z': n_parameters, 25 | 'data_model': model(), 26 | 'vis_y_target': None, 27 | 28 | # MODEL ARCHITECTURE 29 | 'n_blocks': 4, 30 | # 'hidden_layer_sizes': 152, # 500k 31 | 'hidden_layer_sizes': 102, # 250k 32 | 'init_scale': 0.005, 33 | 34 | # TRAINING HYPERPARAMETERS 35 | 'n_epochs': 50, # total number of epochs 36 | 'max_batches_per_epoch': 100, # iterations per epoch (or if training data exhausted) 37 | 'batch_size': 300, 38 | 39 | 'lr_init': 0.01, # initial learning rate 40 | 'pre_low_lr': 3, # number of epochs at the start with very low learning rate 41 | 'final_decay': 0.01, # fraction of the learning rate to reach at final epoch 42 | 'l2_weight_reg': 1.86e-05, # strength of the weight regularization 43 | 'adam_betas': (0.9, 0.95), 44 | } 45 | 46 | # DATA LOADERS 47 | train_loader, test_loader = prepare_uci_loaders(c['data_model'].name, c['batch_size']) 48 | c['train_loader'] = train_loader 49 | c['test_loader'] = test_loader 50 | 51 | 52 | ############################## 53 | ### MODEL ARCHITECTURE ### 54 | ############################## 55 | 56 | x_lane = [InputNode(c['ndim_x'], name='x')] 57 | 58 | for i in range(c['n_blocks']): 59 | if i > 0: 60 | x_lane.append(Node(x_lane[-1], 61 | HouseholderPerm, 62 | {'fixed': True, 'n_reflections': c['ndim_x']}, 63 | name=f'perm_{i}')) 64 | 65 | x_lane.append(Node(x_lane[-1], 66 | HierarchicalAffineCouplingBlock, 67 | {'c_internal': [c['hidden_layer_sizes'], c['hidden_layer_sizes']//2, c['hidden_layer_sizes']//4, c['hidden_layer_sizes']//8]}, 68 | name=f'hac_{i+1}')) 69 | 70 | x_lane.append(OutputNode(x_lane[-1], name='z')) 71 | 72 | model = ReversibleGraphNet(x_lane, verbose=False) 73 | model.to(c['device']) 74 | model.params_trainable = list(filter(lambda p: p.requires_grad, model.parameters())) 75 | 76 | 77 | def model_inverse(test_z): 78 | return model(test_z, rev=True) 79 | 80 | 81 | c['model'] = model 82 | c['model_inverse'] = model_inverse 83 | 84 | # create namedtuple from config dictionary 85 | c = namedtuple("Configuration",c.keys())(*c.values()) 86 | -------------------------------------------------------------------------------- /configs/uci_data/miniboone_hint_8.py: -------------------------------------------------------------------------------- 1 | from collections import namedtuple 2 | import torch 3 | 4 | from FrEIA.framework import * 5 | from FrEIA.modules import * 6 | 7 | from data import prepare_uci_loaders, Miniboone as model 8 | n_parameters = model.n_parameters 9 | 10 | 11 | ###################################################### 12 | ### TRAINING CONFIGURATION AND HYPERPARAMETERS ### 13 | ###################################################### 14 | 15 | c = { 16 | # GENERAL STUFF 17 | 'suffix': f'{model.name}_hint-8', # identifier for trained models and outputs 18 | 'device': 'cuda', # 'cuda' for GPU, 'cpu' for CPU 19 | 'interactive_visualization': True, # requires visdom package to be installed 20 | 21 | # DATA 22 | 'ndim_x': n_parameters, 23 | 'ndim_y': 0, 24 | 'ndim_z': n_parameters, 25 | 'data_model': model(), 26 | 'vis_y_target': None, 27 | 28 | # MODEL ARCHITECTURE 29 | 'n_blocks': 8, 30 | # 'hidden_layer_sizes': 102, # 500k 31 | 'hidden_layer_sizes': 67, # 250k 32 | 'init_scale': 0.005, 33 | 34 | # TRAINING HYPERPARAMETERS 35 | 'n_epochs': 50, # total number of epochs 36 | 'max_batches_per_epoch': 100, # iterations per epoch (or if training data exhausted) 37 | 'batch_size': 300, 38 | 39 | 'lr_init': 0.01, # initial learning rate 40 | 'pre_low_lr': 3, # number of epochs at the start with very low learning rate 41 | 'final_decay': 0.01, # fraction of the learning rate to reach at final epoch 42 | 'l2_weight_reg': 1.86e-05, # strength of the weight regularization 43 | 'adam_betas': (0.9, 0.95), 44 | } 45 | 46 | # DATA LOADERS 47 | train_loader, test_loader = prepare_uci_loaders(c['data_model'].name, c['batch_size']) 48 | c['train_loader'] = train_loader 49 | c['test_loader'] = test_loader 50 | 51 | 52 | ############################## 53 | ### MODEL ARCHITECTURE ### 54 | ############################## 55 | 56 | x_lane = [InputNode(c['ndim_x'], name='x')] 57 | 58 | for i in range(c['n_blocks']): 59 | if i > 0: 60 | x_lane.append(Node(x_lane[-1], 61 | HouseholderPerm, 62 | {'fixed': True, 'n_reflections': c['ndim_x']}, 63 | name=f'perm_{i}')) 64 | 65 | x_lane.append(Node(x_lane[-1], 66 | HierarchicalAffineCouplingBlock, 67 | {'c_internal': [c['hidden_layer_sizes'], c['hidden_layer_sizes']//2, c['hidden_layer_sizes']//4, c['hidden_layer_sizes']//8]}, 68 | name=f'hac_{i+1}')) 69 | 70 | x_lane.append(OutputNode(x_lane[-1], name='z')) 71 | 72 | model = ReversibleGraphNet(x_lane, verbose=False) 73 | model.to(c['device']) 74 | model.params_trainable = list(filter(lambda p: p.requires_grad, model.parameters())) 75 | 76 | 77 | def model_inverse(test_z): 78 | return model(test_z, rev=True) 79 | 80 | 81 | c['model'] = model 82 | c['model_inverse'] = model_inverse 83 | 84 | # create namedtuple from config dictionary 85 | c = namedtuple("Configuration",c.keys())(*c.values()) 86 | -------------------------------------------------------------------------------- /configs/uci_data/miniboone_inn_4.py: -------------------------------------------------------------------------------- 1 | from collections import namedtuple 2 | import torch 3 | 4 | from FrEIA.framework import * 5 | from FrEIA.modules import * 6 | 7 | from data import prepare_uci_loaders, Miniboone as model 8 | n_parameters = model.n_parameters 9 | 10 | 11 | ###################################################### 12 | ### TRAINING CONFIGURATION AND HYPERPARAMETERS ### 13 | ###################################################### 14 | 15 | c = { 16 | # GENERAL STUFF 17 | 'suffix': f'{model.name}_inn-4', # identifier for trained models and outputs 18 | 'device': 'cuda', # 'cuda' for GPU, 'cpu' for CPU 19 | 'interactive_visualization': True, # requires visdom package to be installed 20 | 21 | # DATA 22 | 'ndim_x': n_parameters, 23 | 'ndim_y': 0, 24 | 'ndim_z': n_parameters, 25 | 'data_model': model(), 26 | 'vis_y_target': None, 27 | 28 | # MODEL ARCHITECTURE 29 | 'n_blocks': 4, 30 | 'hidden_layer_sizes': 114, # 250k 31 | 'init_scale': 0.005, 32 | 33 | # TRAINING HYPERPARAMETERS 34 | 'n_epochs': 150, # total number of epochs 35 | 'max_batches_per_epoch': 100, # iterations per epoch (or if training data exhausted) 36 | 'batch_size': 300, 37 | 38 | 'lr_init': 0.01, # initial learning rate 39 | 'pre_low_lr': 3, # number of epochs at the start with very low learning rate 40 | 'final_decay': 0.01, # fraction of the learning rate to reach at final epoch 41 | 'l2_weight_reg': 1.86e-05, # strength of the weight regularization 42 | 'adam_betas': (0.9, 0.95), 43 | } 44 | 45 | # DATA LOADERS 46 | train_loader, test_loader = prepare_uci_loaders(c['data_model'].name, c['batch_size']) 47 | c['train_loader'] = train_loader 48 | c['test_loader'] = test_loader 49 | 50 | 51 | ############################## 52 | ### MODEL ARCHITECTURE ### 53 | ############################## 54 | 55 | nodes = [InputNode(c['ndim_x'], name='x')] 56 | 57 | for i in range(c['n_blocks']): 58 | nodes.append(Node(nodes[-1], 59 | HouseholderPerm, 60 | {'fixed': True, 'n_reflections': c['ndim_x']}, 61 | name=f'perm_{i+1}')) 62 | 63 | nodes.append(Node(nodes[-1], 64 | AffineCoupling, 65 | {'F_class': F_fully_connected, 66 | 'F_args': {'internal_size': c['hidden_layer_sizes']}}, 67 | name=f'ac_{i+1}')) 68 | 69 | nodes.append(OutputNode(nodes[-1], name='z')) 70 | 71 | model = ReversibleGraphNet(nodes, verbose=False) 72 | model.to(c['device']) 73 | model.params_trainable = list(filter(lambda p: p.requires_grad, model.parameters())) 74 | 75 | 76 | def model_inverse(test_z): 77 | return model(test_z, rev=True) 78 | 79 | 80 | c['model'] = model 81 | c['model_inverse'] = model_inverse 82 | 83 | # create namedtuple from config dictionary 84 | c = namedtuple("Configuration",c.keys())(*c.values()) 85 | -------------------------------------------------------------------------------- /configs/uci_data/miniboone_inn_8.py: -------------------------------------------------------------------------------- 1 | from collections import namedtuple 2 | import torch 3 | 4 | from FrEIA.framework import * 5 | from FrEIA.modules import * 6 | 7 | from data import prepare_uci_loaders, Miniboone as model 8 | n_parameters = model.n_parameters 9 | 10 | 11 | ###################################################### 12 | ### TRAINING CONFIGURATION AND HYPERPARAMETERS ### 13 | ###################################################### 14 | 15 | c = { 16 | # GENERAL STUFF 17 | 'suffix': f'{model.name}_inn-8', # identifier for trained models and outputs 18 | 'device': 'cuda', # 'cuda' for GPU, 'cpu' for CPU 19 | 'interactive_visualization': True, # requires visdom package to be installed 20 | 21 | # DATA 22 | 'ndim_x': n_parameters, 23 | 'ndim_y': 0, 24 | 'ndim_z': n_parameters, 25 | 'data_model': model(), 26 | 'vis_y_target': None, 27 | 28 | # MODEL ARCHITECTURE 29 | 'n_blocks': 8, 30 | 'hidden_layer_sizes': 78, # 250k 31 | 'init_scale': 0.005, 32 | 33 | # TRAINING HYPERPARAMETERS 34 | 'n_epochs': 50, # total number of epochs 35 | 'max_batches_per_epoch': 100, # iterations per epoch (or if training data exhausted) 36 | 'batch_size': 300, 37 | 38 | 'lr_init': 0.01, # initial learning rate 39 | 'pre_low_lr': 3, # number of epochs at the start with very low learning rate 40 | 'final_decay': 0.01, # fraction of the learning rate to reach at final epoch 41 | 'l2_weight_reg': 1.86e-05, # strength of the weight regularization 42 | 'adam_betas': (0.9, 0.95), 43 | } 44 | 45 | # DATA LOADERS 46 | train_loader, test_loader = prepare_uci_loaders(c['data_model'].name, c['batch_size']) 47 | c['train_loader'] = train_loader 48 | c['test_loader'] = test_loader 49 | 50 | 51 | ############################## 52 | ### MODEL ARCHITECTURE ### 53 | ############################## 54 | 55 | nodes = [InputNode(c['ndim_x'], name='x')] 56 | 57 | for i in range(c['n_blocks']): 58 | nodes.append(Node(nodes[-1], 59 | HouseholderPerm, 60 | {'fixed': True, 'n_reflections': c['ndim_x']}, 61 | name=f'perm_{i+1}')) 62 | 63 | nodes.append(Node(nodes[-1], 64 | AffineCoupling, 65 | {'F_class': F_fully_connected, 66 | 'F_args': {'internal_size': c['hidden_layer_sizes']}}, 67 | name=f'ac_{i+1}')) 68 | 69 | nodes.append(OutputNode(nodes[-1], name='z')) 70 | 71 | model = ReversibleGraphNet(nodes, verbose=False) 72 | model.to(c['device']) 73 | model.params_trainable = list(filter(lambda p: p.requires_grad, model.parameters())) 74 | 75 | 76 | def model_inverse(test_z): 77 | return model(test_z, rev=True) 78 | 79 | 80 | c['model'] = model 81 | c['model_inverse'] = model_inverse 82 | 83 | # create namedtuple from config dictionary 84 | c = namedtuple("Configuration",c.keys())(*c.values()) 85 | -------------------------------------------------------------------------------- /configs/uci_data/power_hint_4.py: -------------------------------------------------------------------------------- 1 | from collections import namedtuple 2 | import torch 3 | 4 | from FrEIA.framework import * 5 | from FrEIA.modules import * 6 | 7 | from data import prepare_uci_loaders, Power as model 8 | n_parameters = model.n_parameters 9 | 10 | 11 | ###################################################### 12 | ### TRAINING CONFIGURATION AND HYPERPARAMETERS ### 13 | ###################################################### 14 | 15 | c = { 16 | # GENERAL STUFF 17 | 'suffix': f'{model.name}_hint-4', # identifier for trained models and outputs 18 | 'device': 'cuda', # 'cuda' for GPU, 'cpu' for CPU 19 | 'interactive_visualization': True, # requires visdom package to be installed 20 | 21 | # DATA 22 | 'ndim_x': n_parameters, 23 | 'ndim_y': 0, 24 | 'ndim_z': n_parameters, 25 | 'data_model': model(), 26 | 'vis_y_target': None, 27 | 28 | # MODEL ARCHITECTURE 29 | 'n_blocks': 4, 30 | 'hidden_layer_sizes': 200, # 500k 31 | 'init_scale': 0.005, 32 | 33 | # TRAINING HYPERPARAMETERS 34 | 'n_epochs': 50, # total number of epochs 35 | 'max_batches_per_epoch': 1000, # iterations per epoch (or if training data exhausted) 36 | 'batch_size': 1660, 37 | 38 | 'lr_init': 0.01, # initial learning rate 39 | 'pre_low_lr': 3, # number of epochs at the start with very low learning rate 40 | 'final_decay': 0.01, # fraction of the learning rate to reach at final epoch 41 | 'l2_weight_reg': 1.86e-05, # strength of the weight regularization 42 | 'adam_betas': (0.9, 0.95), 43 | } 44 | 45 | # DATA LOADERS 46 | train_loader, test_loader = prepare_uci_loaders(c['data_model'].name, c['batch_size']) 47 | c['train_loader'] = train_loader 48 | c['test_loader'] = test_loader 49 | 50 | 51 | ############################## 52 | ### MODEL ARCHITECTURE ### 53 | ############################## 54 | 55 | x_lane = [InputNode(c['ndim_x'], name='x')] 56 | 57 | for i in range(c['n_blocks']): 58 | if i > 0: 59 | x_lane.append(Node(x_lane[-1], 60 | HouseholderPerm, 61 | {'fixed': True, 'n_reflections': c['ndim_x']}, 62 | name=f'perm_{i}')) 63 | 64 | x_lane.append(Node(x_lane[-1], 65 | HierarchicalAffineCouplingBlock, 66 | {'c_internal': [c['hidden_layer_sizes'], c['hidden_layer_sizes']//2, c['hidden_layer_sizes']//4, c['hidden_layer_sizes']//8]}, 67 | name=f'hac_{i+1}')) 68 | 69 | x_lane.append(OutputNode(x_lane[-1], name='z')) 70 | 71 | model = ReversibleGraphNet(x_lane, verbose=False) 72 | model.to(c['device']) 73 | model.params_trainable = list(filter(lambda p: p.requires_grad, model.parameters())) 74 | 75 | 76 | def model_inverse(test_z): 77 | return model(test_z, rev=True) 78 | 79 | 80 | c['model'] = model 81 | c['model_inverse'] = model_inverse 82 | 83 | # create namedtuple from config dictionary 84 | c = namedtuple("Configuration",c.keys())(*c.values()) 85 | -------------------------------------------------------------------------------- /configs/uci_data/power_hint_8.py: -------------------------------------------------------------------------------- 1 | from collections import namedtuple 2 | import torch 3 | 4 | from FrEIA.framework import * 5 | from FrEIA.modules import * 6 | 7 | from data import prepare_uci_loaders, Power as model 8 | n_parameters = model.n_parameters 9 | 10 | 11 | ###################################################### 12 | ### TRAINING CONFIGURATION AND HYPERPARAMETERS ### 13 | ###################################################### 14 | 15 | c = { 16 | # GENERAL STUFF 17 | 'suffix': f'{model.name}_hint-8', # identifier for trained models and outputs 18 | 'device': 'cuda', # 'cuda' for GPU, 'cpu' for CPU 19 | 'interactive_visualization': True, # requires visdom package to be installed 20 | 21 | # DATA 22 | 'ndim_x': n_parameters, 23 | 'ndim_y': 0, 24 | 'ndim_z': n_parameters, 25 | 'data_model': model(), 26 | 'vis_y_target': None, 27 | 28 | # MODEL ARCHITECTURE 29 | 'n_blocks': 8, 30 | 'hidden_layer_sizes': 140, # 500k 31 | 'init_scale': 0.005, 32 | 33 | # TRAINING HYPERPARAMETERS 34 | 'n_epochs': 50, # total number of epochs 35 | 'max_batches_per_epoch': 1000, # iterations per epoch (or if training data exhausted) 36 | 'batch_size': 1660, 37 | 38 | 'lr_init': 0.01, # initial learning rate 39 | 'pre_low_lr': 3, # number of epochs at the start with very low learning rate 40 | 'final_decay': 0.01, # fraction of the learning rate to reach at final epoch 41 | 'l2_weight_reg': 1.86e-05, # strength of the weight regularization 42 | 'adam_betas': (0.9, 0.95), 43 | } 44 | 45 | # DATA LOADERS 46 | train_loader, test_loader = prepare_uci_loaders(c['data_model'].name, c['batch_size']) 47 | c['train_loader'] = train_loader 48 | c['test_loader'] = test_loader 49 | 50 | 51 | ############################## 52 | ### MODEL ARCHITECTURE ### 53 | ############################## 54 | 55 | x_lane = [InputNode(c['ndim_x'], name='x')] 56 | 57 | for i in range(c['n_blocks']): 58 | if i > 0: 59 | x_lane.append(Node(x_lane[-1], 60 | HouseholderPerm, 61 | {'fixed': True, 'n_reflections': c['ndim_x']}, 62 | name=f'perm_{i}')) 63 | 64 | x_lane.append(Node(x_lane[-1], 65 | HierarchicalAffineCouplingBlock, 66 | {'c_internal': [c['hidden_layer_sizes'], c['hidden_layer_sizes']//2, c['hidden_layer_sizes']//4, c['hidden_layer_sizes']//8]}, 67 | name=f'hac_{i+1}')) 68 | 69 | x_lane.append(OutputNode(x_lane[-1], name='z')) 70 | 71 | model = ReversibleGraphNet(x_lane, verbose=False) 72 | model.to(c['device']) 73 | model.params_trainable = list(filter(lambda p: p.requires_grad, model.parameters())) 74 | 75 | 76 | def model_inverse(test_z): 77 | return model(test_z, rev=True) 78 | 79 | 80 | c['model'] = model 81 | c['model_inverse'] = model_inverse 82 | 83 | # create namedtuple from config dictionary 84 | c = namedtuple("Configuration",c.keys())(*c.values()) 85 | -------------------------------------------------------------------------------- /configs/uci_data/power_inn_4.py: -------------------------------------------------------------------------------- 1 | from collections import namedtuple 2 | import torch 3 | 4 | from FrEIA.framework import * 5 | from FrEIA.modules import * 6 | 7 | from data import prepare_uci_loaders, Power as model 8 | n_parameters = model.n_parameters 9 | 10 | 11 | ###################################################### 12 | ### TRAINING CONFIGURATION AND HYPERPARAMETERS ### 13 | ###################################################### 14 | 15 | c = { 16 | # GENERAL STUFF 17 | 'suffix': f'{model.name}_inn-4', # identifier for trained models and outputs 18 | 'device': 'cuda', # 'cuda' for GPU, 'cpu' for CPU 19 | 'interactive_visualization': True, # requires visdom package to be installed 20 | 21 | # DATA 22 | 'ndim_x': n_parameters, 23 | 'ndim_y': 0, 24 | 'ndim_z': n_parameters, 25 | 'data_model': model(), 26 | 'vis_y_target': None, 27 | 28 | # MODEL ARCHITECTURE 29 | 'n_blocks': 4, 30 | 'hidden_layer_sizes': 175, # 500k 31 | 'init_scale': 0.005, 32 | 33 | # TRAINING HYPERPARAMETERS 34 | 'n_epochs': 50, # total number of epochs 35 | 'max_batches_per_epoch': 1000, # iterations per epoch (or if training data exhausted) 36 | 'batch_size': 1660, 37 | 38 | 'lr_init': 0.01, # initial learning rate 39 | 'pre_low_lr': 3, # number of epochs at the start with very low learning rate 40 | 'final_decay': 0.01, # fraction of the learning rate to reach at final epoch 41 | 'l2_weight_reg': 1.86e-05, # strength of the weight regularization 42 | 'adam_betas': (0.9, 0.95), 43 | } 44 | 45 | # DATA LOADERS 46 | train_loader, test_loader = prepare_uci_loaders(c['data_model'].name, c['batch_size']) 47 | c['train_loader'] = train_loader 48 | c['test_loader'] = test_loader 49 | 50 | 51 | ############################## 52 | ### MODEL ARCHITECTURE ### 53 | ############################## 54 | 55 | nodes = [InputNode(c['ndim_x'], name='x')] 56 | 57 | for i in range(c['n_blocks']): 58 | nodes.append(Node(nodes[-1], 59 | HouseholderPerm, 60 | {'fixed': True, 'n_reflections': c['ndim_x']}, 61 | name=f'perm_{i+1}')) 62 | 63 | nodes.append(Node(nodes[-1], 64 | AffineCoupling, 65 | {'F_class': F_fully_connected, 66 | 'F_args': {'internal_size': c['hidden_layer_sizes']}}, 67 | name=f'ac_{i+1}')) 68 | 69 | nodes.append(OutputNode(nodes[-1], name='z')) 70 | 71 | model = ReversibleGraphNet(nodes, verbose=False) 72 | model.to(c['device']) 73 | model.params_trainable = list(filter(lambda p: p.requires_grad, model.parameters())) 74 | 75 | 76 | def model_inverse(test_z): 77 | return model(test_z, rev=True) 78 | 79 | 80 | c['model'] = model 81 | c['model_inverse'] = model_inverse 82 | 83 | # create namedtuple from config dictionary 84 | c = namedtuple("Configuration",c.keys())(*c.values()) 85 | -------------------------------------------------------------------------------- /configs/uci_data/power_inn_8.py: -------------------------------------------------------------------------------- 1 | from collections import namedtuple 2 | import torch 3 | 4 | from FrEIA.framework import * 5 | from FrEIA.modules import * 6 | 7 | from data import prepare_uci_loaders, Power as model 8 | n_parameters = model.n_parameters 9 | 10 | 11 | ###################################################### 12 | ### TRAINING CONFIGURATION AND HYPERPARAMETERS ### 13 | ###################################################### 14 | 15 | c = { 16 | # GENERAL STUFF 17 | 'suffix': f'{model.name}_inn-8', # identifier for trained models and outputs 18 | 'device': 'cuda', # 'cuda' for GPU, 'cpu' for CPU 19 | 'interactive_visualization': True, # requires visdom package to be installed 20 | 21 | # DATA 22 | 'ndim_x': n_parameters, 23 | 'ndim_y': 0, 24 | 'ndim_z': n_parameters, 25 | 'data_model': model(), 26 | 'vis_y_target': None, 27 | 28 | # MODEL ARCHITECTURE 29 | 'n_blocks': 8, 30 | 'hidden_layer_sizes': 123, # 500k 31 | 'init_scale': 0.005, 32 | 33 | # TRAINING HYPERPARAMETERS 34 | 'n_epochs': 50, # total number of epochs 35 | 'max_batches_per_epoch': 1000, # iterations per epoch (or if training data exhausted) 36 | 'batch_size': 1660, 37 | 38 | 'lr_init': 0.01, # initial learning rate 39 | 'pre_low_lr': 3, # number of epochs at the start with very low learning rate 40 | 'final_decay': 0.01, # fraction of the learning rate to reach at final epoch 41 | 'l2_weight_reg': 1.86e-05, # strength of the weight regularization 42 | 'adam_betas': (0.9, 0.95), 43 | } 44 | 45 | # DATA LOADERS 46 | train_loader, test_loader = prepare_uci_loaders(c['data_model'].name, c['batch_size']) 47 | c['train_loader'] = train_loader 48 | c['test_loader'] = test_loader 49 | 50 | 51 | ############################## 52 | ### MODEL ARCHITECTURE ### 53 | ############################## 54 | 55 | nodes = [InputNode(c['ndim_x'], name='x')] 56 | 57 | for i in range(c['n_blocks']): 58 | nodes.append(Node(nodes[-1], 59 | HouseholderPerm, 60 | {'fixed': True, 'n_reflections': c['ndim_x']}, 61 | name=f'perm_{i+1}')) 62 | 63 | nodes.append(Node(nodes[-1], 64 | AffineCoupling, 65 | {'F_class': F_fully_connected, 66 | 'F_args': {'internal_size': c['hidden_layer_sizes']}}, 67 | name=f'ac_{i+1}')) 68 | 69 | nodes.append(OutputNode(nodes[-1], name='z')) 70 | 71 | model = ReversibleGraphNet(nodes, verbose=False) 72 | model.to(c['device']) 73 | model.params_trainable = list(filter(lambda p: p.requires_grad, model.parameters())) 74 | 75 | 76 | def model_inverse(test_z): 77 | return model(test_z, rev=True) 78 | 79 | 80 | c['model'] = model 81 | c['model_inverse'] = model_inverse 82 | 83 | # create namedtuple from config dictionary 84 | c = namedtuple("Configuration",c.keys())(*c.values()) 85 | -------------------------------------------------------------------------------- /dense_jacobian.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/vislearn/HINT/b26c5026bd486bb392a9416430cb1dfebe5aa06f/dense_jacobian.png -------------------------------------------------------------------------------- /fourier_curve.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/vislearn/HINT/b26c5026bd486bb392a9416430cb1dfebe5aa06f/fourier_curve.png -------------------------------------------------------------------------------- /recursive_block.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/vislearn/HINT/b26c5026bd486bb392a9416430cb1dfebe5aa06f/recursive_block.png -------------------------------------------------------------------------------- /run_uci_experiments.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import numpy as np 3 | import traceback 4 | from best_shape_fit import * 5 | from tqdm import tqdm 6 | 7 | 8 | 9 | n_runs = 3 10 | 11 | configs = [ 12 | 'uci_data.power_hint_4', 13 | 'uci_data.power_hint_8', 14 | 'uci_data.power_inn_4', 15 | 'uci_data.power_inn_8', 16 | 17 | 'uci_data.gas_hint_4', 18 | 'uci_data.gas_hint_8', 19 | 'uci_data.gas_inn_4', 20 | 'uci_data.gas_inn_8', 21 | 22 | 'uci_data.miniboone_hint_4', 23 | 'uci_data.miniboone_hint_8', 24 | 'uci_data.miniboone_inn_4', 25 | 'uci_data.miniboone_inn_8', 26 | ] 27 | 28 | 29 | def train_and_eval(): 30 | for config in configs: 31 | test_losses = [] 32 | for i in range(n_runs): 33 | try: 34 | # Import config 35 | exec(f'from configs.{config} import c', globals()) 36 | 37 | # n_model_params = sum([p.numel() for p in c.model.params_trainable]) 38 | # print(f'Model {c.suffix} has {n_model_params:,} trainable parameters.') 39 | 40 | from train_unconditional import main, save, load, evaluate 41 | test_losses.append(main(c)) 42 | save(c, f'results/{config.replace(".", "-")}_{i}.pt') 43 | # load(c, f'results/{config.replace(".", "-")}_{i}.pt') 44 | # test_losses.append(evaluate(c)) 45 | 46 | except Exception as e: 47 | print(f'ERROR with config "{config}"', i) 48 | # print(e) 49 | # traceback.print_exc() 50 | 51 | print(config) 52 | print(test_losses) 53 | np.save(f'results/{config.replace(".", "-")}', np.array(test_losses)) 54 | 55 | 56 | 57 | def collect_results(): 58 | # from data import Power, Gas, Miniboone 59 | 60 | for config in configs: 61 | if 'power' in config: 62 | n_dims = Power.n_parameters 63 | # mean, std = Power.mean_and_std() 64 | if 'gas' in config: 65 | n_dims = Gas.n_parameters 66 | # mean, std = Gas.mean_and_std() 67 | if 'miniboone' in config: 68 | n_dims = Miniboone.n_parameters 69 | # mean, std = Miniboone.mean_and_std() 70 | 71 | test_losses = -np.load(f'results/{config.replace(".", "-")}.npy') 72 | test_losses -= np.log(2*np.pi) * (n_dims/2) 73 | 74 | print(config) 75 | print(f'{test_losses.mean():.3f} \pm {test_losses.std():.3f}') 76 | # print('mean:', test_losses.mean()) 77 | # print('std: ', test_losses.std()) 78 | print() 79 | 80 | 81 | 82 | if __name__ == '__main__': 83 | pass 84 | 85 | # train_and_eval() 86 | collect_results() 87 | 88 | # print(Power.mean_and_std(), '\n') 89 | # print(Gas.mean_and_std(), '\n') 90 | # print(Miniboone.mean_and_std(), '\n') 91 | --------------------------------------------------------------------------------