├── .gitignore ├── LICENSE ├── README.md ├── setup.py ├── test └── test_lazyload.py └── torchhacks ├── __init__.py └── lazyload.py /.gitignore: -------------------------------------------------------------------------------- 1 | __pycache__/ 2 | *.pyc 3 | 4 | build/ 5 | dist/ 6 | *.egg-info/ 7 | 8 | .ipynb_checkpoints 9 | *~ -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | Copyright 2023 MathInf GmbH 2 | 3 | Licensed under the Apache License, Version 2.0 (the "License"); 4 | you may not use this files from this repository except in compliance 5 | with the License reproduced below (also at 6 | http://www.apache.org/licenses/LICENSE-2.0). 7 | 8 | Unless required by applicable law or agreed to in writing, software 9 | distributed under the License is distributed on an "AS IS" BASIS, 10 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 11 | See the License for the specific language governing permissions and 12 | limitations under the License. 13 | 14 | 15 | 16 | 17 | Apache License 18 | Version 2.0, January 2004 19 | http://www.apache.org/licenses/ 20 | 21 | TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION 22 | 23 | 1. Definitions. 24 | 25 | "License" shall mean the terms and conditions for use, reproduction, 26 | and distribution as defined by Sections 1 through 9 of this document. 27 | 28 | "Licensor" shall mean the copyright owner or entity authorized by 29 | the copyright owner that is granting the License. 30 | 31 | "Legal Entity" shall mean the union of the acting entity and all 32 | other entities that control, are controlled by, or are under common 33 | control with that entity. For the purposes of this definition, 34 | "control" means (i) the power, direct or indirect, to cause the 35 | direction or management of such entity, whether by contract or 36 | otherwise, or (ii) ownership of fifty percent (50%) or more of the 37 | outstanding shares, or (iii) beneficial ownership of such entity. 38 | 39 | "You" (or "Your") shall mean an individual or Legal Entity 40 | exercising permissions granted by this License. 41 | 42 | "Source" form shall mean the preferred form for making modifications, 43 | including but not limited to software source code, documentation 44 | source, and configuration files. 45 | 46 | "Object" form shall mean any form resulting from mechanical 47 | transformation or translation of a Source form, including but 48 | not limited to compiled object code, generated documentation, 49 | and conversions to other media types. 50 | 51 | "Work" shall mean the work of authorship, whether in Source or 52 | Object form, made available under the License, as indicated by a 53 | copyright notice that is included in or attached to the work 54 | (an example is provided in the Appendix below). 55 | 56 | "Derivative Works" shall mean any work, whether in Source or Object 57 | form, that is based on (or derived from) the Work and for which the 58 | editorial revisions, annotations, elaborations, or other modifications 59 | represent, as a whole, an original work of authorship. For the purposes 60 | of this License, Derivative Works shall not include works that remain 61 | separable from, or merely link (or bind by name) to the interfaces of, 62 | the Work and Derivative Works thereof. 63 | 64 | "Contribution" shall mean any work of authorship, including 65 | the original version of the Work and any modifications or additions 66 | to that Work or Derivative Works thereof, that is intentionally 67 | submitted to Licensor for inclusion in the Work by the copyright owner 68 | or by an individual or Legal Entity authorized to submit on behalf of 69 | the copyright owner. For the purposes of this definition, "submitted" 70 | means any form of electronic, verbal, or written communication sent 71 | to the Licensor or its representatives, including but not limited to 72 | communication on electronic mailing lists, source code control systems, 73 | and issue tracking systems that are managed by, or on behalf of, the 74 | Licensor for the purpose of discussing and improving the Work, but 75 | excluding communication that is conspicuously marked or otherwise 76 | designated in writing by the copyright owner as "Not a Contribution." 77 | 78 | "Contributor" shall mean Licensor and any individual or Legal Entity 79 | on behalf of whom a Contribution has been received by Licensor and 80 | subsequently incorporated within the Work. 81 | 82 | 2. Grant of Copyright License. Subject to the terms and conditions of 83 | this License, each Contributor hereby grants to You a perpetual, 84 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable 85 | copyright license to reproduce, prepare Derivative Works of, 86 | publicly display, publicly perform, sublicense, and distribute the 87 | Work and such Derivative Works in Source or Object form. 88 | 89 | 3. Grant of Patent License. Subject to the terms and conditions of 90 | this License, each Contributor hereby grants to You a perpetual, 91 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable 92 | (except as stated in this section) patent license to make, have made, 93 | use, offer to sell, sell, import, and otherwise transfer the Work, 94 | where such license applies only to those patent claims licensable 95 | by such Contributor that are necessarily infringed by their 96 | Contribution(s) alone or by combination of their Contribution(s) 97 | with the Work to which such Contribution(s) was submitted. If You 98 | institute patent litigation against any entity (including a 99 | cross-claim or counterclaim in a lawsuit) alleging that the Work 100 | or a Contribution incorporated within the Work constitutes direct 101 | or contributory patent infringement, then any patent licenses 102 | granted to You under this License for that Work shall terminate 103 | as of the date such litigation is filed. 104 | 105 | 4. Redistribution. You may reproduce and distribute copies of the 106 | Work or Derivative Works thereof in any medium, with or without 107 | modifications, and in Source or Object form, provided that You 108 | meet the following conditions: 109 | 110 | (a) You must give any other recipients of the Work or 111 | Derivative Works a copy of this License; and 112 | 113 | (b) You must cause any modified files to carry prominent notices 114 | stating that You changed the files; and 115 | 116 | (c) You must retain, in the Source form of any Derivative Works 117 | that You distribute, all copyright, patent, trademark, and 118 | attribution notices from the Source form of the Work, 119 | excluding those notices that do not pertain to any part of 120 | the Derivative Works; and 121 | 122 | (d) If the Work includes a "NOTICE" text file as part of its 123 | distribution, then any Derivative Works that You distribute must 124 | include a readable copy of the attribution notices contained 125 | within such NOTICE file, excluding those notices that do not 126 | pertain to any part of the Derivative Works, in at least one 127 | of the following places: within a NOTICE text file distributed 128 | as part of the Derivative Works; within the Source form or 129 | documentation, if provided along with the Derivative Works; or, 130 | within a display generated by the Derivative Works, if and 131 | wherever such third-party notices normally appear. The contents 132 | of the NOTICE file are for informational purposes only and 133 | do not modify the License. You may add Your own attribution 134 | notices within Derivative Works that You distribute, alongside 135 | or as an addendum to the NOTICE text from the Work, provided 136 | that such additional attribution notices cannot be construed 137 | as modifying the License. 138 | 139 | You may add Your own copyright statement to Your modifications and 140 | may provide additional or different license terms and conditions 141 | for use, reproduction, or distribution of Your modifications, or 142 | for any such Derivative Works as a whole, provided Your use, 143 | reproduction, and distribution of the Work otherwise complies with 144 | the conditions stated in this License. 145 | 146 | 5. Submission of Contributions. Unless You explicitly state otherwise, 147 | any Contribution intentionally submitted for inclusion in the Work 148 | by You to the Licensor shall be under the terms and conditions of 149 | this License, without any additional terms or conditions. 150 | Notwithstanding the above, nothing herein shall supersede or modify 151 | the terms of any separate license agreement you may have executed 152 | with Licensor regarding such Contributions. 153 | 154 | 6. Trademarks. This License does not grant permission to use the trade 155 | names, trademarks, service marks, or product names of the Licensor, 156 | except as required for reasonable and customary use in describing the 157 | origin of the Work and reproducing the content of the NOTICE file. 158 | 159 | 7. Disclaimer of Warranty. Unless required by applicable law or 160 | agreed to in writing, Licensor provides the Work (and each 161 | Contributor provides its Contributions) on an "AS IS" BASIS, 162 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or 163 | implied, including, without limitation, any warranties or conditions 164 | of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A 165 | PARTICULAR PURPOSE. You are solely responsible for determining the 166 | appropriateness of using or redistributing the Work and assume any 167 | risks associated with Your exercise of permissions under this License. 168 | 169 | 8. Limitation of Liability. In no event and under no legal theory, 170 | whether in tort (including negligence), contract, or otherwise, 171 | unless required by applicable law (such as deliberate and grossly 172 | negligent acts) or agreed to in writing, shall any Contributor be 173 | liable to You for damages, including any direct, indirect, special, 174 | incidental, or consequential damages of any character arising as a 175 | result of this License or out of the use or inability to use the 176 | Work (including but not limited to damages for loss of goodwill, 177 | work stoppage, computer failure or malfunction, or any and all 178 | other commercial damages or losses), even if such Contributor 179 | has been advised of the possibility of such damages. 180 | 181 | 9. Accepting Warranty or Additional Liability. While redistributing 182 | the Work or Derivative Works thereof, You may choose to offer, 183 | and charge a fee for, acceptance of support, warranty, indemnity, 184 | or other liability obligations and/or rights consistent with this 185 | License. However, in accepting such obligations, You may act only 186 | on Your own behalf and on Your sole responsibility, not on behalf 187 | of any other Contributor, and only if You agree to indemnify, 188 | defend, and hold each Contributor harmless for any liability 189 | incurred by, or claims asserted against, such Contributor by reason 190 | of your accepting any such warranty or additional liability. 191 | 192 | END OF TERMS AND CONDITIONS 193 | 194 | APPENDIX: How to apply the Apache License to your work. 195 | 196 | To apply the Apache License to your work, attach the following 197 | boilerplate notice, with the fields enclosed by brackets "[]" 198 | replaced with your own identifying information. (Don't include 199 | the brackets!) The text should be enclosed in the appropriate 200 | comment syntax for the file format. We also recommend that a 201 | file or class name and description of purpose be included on the 202 | same "printed page" as the copyright notice for easier 203 | identification within third-party archives. 204 | 205 | Copyright [yyyy] [name of copyright owner] 206 | 207 | Licensed under the Apache License, Version 2.0 (the "License"); 208 | you may not use this file except in compliance with the License. 209 | You may obtain a copy of the License at 210 | 211 | http://www.apache.org/licenses/LICENSE-2.0 212 | 213 | Unless required by applicable law or agreed to in writing, software 214 | distributed under the License is distributed on an "AS IS" BASIS, 215 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 216 | See the License for the specific language governing permissions and 217 | limitations under the License. 218 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # TorchHacks 2 | 3 | Making PyTorch nicer in ways that might not be entirely safe. 4 | -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | import setuptools 2 | 3 | 4 | with open("README.md", "r", encoding="utf-8") as fh: 5 | long_description = fh.read() 6 | 7 | setuptools.setup( 8 | name="torchhacks", 9 | version="0.0.1", 10 | description="Hacks for PyTorch", 11 | long_description=long_description, 12 | long_description_content_type="text/markdown", 13 | author="Thomas Viehmann, MathInf GmbH", 14 | url="https://lernappar.at/torchhacks", 15 | install_requires=["torch"], 16 | packages=setuptools.find_packages(), 17 | ) 18 | -------------------------------------------------------------------------------- /test/test_lazyload.py: -------------------------------------------------------------------------------- 1 | import pytest 2 | import tempfile 3 | import pathlib 4 | import torchhacks 5 | import torch 6 | 7 | 8 | def test_lazy_load_basic(): 9 | with tempfile.TemporaryDirectory() as tmpdirname: 10 | m = torch.nn.Linear(5, 3) 11 | path = pathlib.Path(tmpdirname) 12 | fn = str(path / "test.pt") 13 | torch.save(m.state_dict(), fn) 14 | sd_lazy = torchhacks.lazy_load(fn) 15 | assert "NotYetLoadedTensor" in str(next(iter(sd_lazy.values()))) 16 | m2 = torch.nn.Linear(5, 3) 17 | m2.load_state_dict(sd_lazy) 18 | 19 | x = torch.randn(2, 5) 20 | actual = m2(x) 21 | expected = m(x) 22 | torch.testing.assert_close(actual, expected) 23 | 24 | 25 | if __name__ == "__main__": 26 | pytest.main([__file__]) 27 | -------------------------------------------------------------------------------- /torchhacks/__init__.py: -------------------------------------------------------------------------------- 1 | from .lazyload import lazy_load 2 | -------------------------------------------------------------------------------- /torchhacks/lazyload.py: -------------------------------------------------------------------------------- 1 | import zipfile 2 | import pickle 3 | import functools 4 | import torch 5 | import warnings 6 | import pathlib 7 | 8 | 9 | class NotYetLoadedTensor: 10 | def __init__(self, metatensor, archiveinfo, storageinfo, rebuild_args): 11 | self.metatensor = metatensor 12 | self.archiveinfo = archiveinfo 13 | self.storageinfo = storageinfo 14 | self.rebuild_args = rebuild_args 15 | 16 | @classmethod 17 | def rebuild( 18 | cls, 19 | storage, 20 | storage_offset, 21 | size, 22 | stride, 23 | requires_grad, 24 | backward_hooks, 25 | metadata=None, 26 | archiveinfo=None, 27 | ): 28 | rebuild_args = ( 29 | storage_offset, 30 | size, 31 | stride, 32 | requires_grad, 33 | backward_hooks, 34 | metadata, 35 | ) 36 | metatensor = torch._utils._rebuild_tensor_v2( 37 | storage, 38 | storage_offset, 39 | size, 40 | stride, 41 | requires_grad, 42 | backward_hooks, 43 | metadata, 44 | ) 45 | storageinfo = archiveinfo.cache[storage._cdata] 46 | return NotYetLoadedTensor(metatensor, archiveinfo, storageinfo, rebuild_args) 47 | 48 | def _load_tensor(self): 49 | # we could / should try to lean heavier on PyTorch's reader 50 | name, storage_cls, fn, device, size = self.storageinfo 51 | buffer = self.archiveinfo.zipfile.read( 52 | str(self.archiveinfo.prefix / "data" / fn) 53 | ) 54 | with warnings.catch_warnings(): 55 | warnings.simplefilter("ignore") 56 | storage = storage_cls.from_buffer(buffer, "native") 57 | tensor = torch._utils._rebuild_tensor_v2(storage, *self.rebuild_args) 58 | return tensor 59 | 60 | @classmethod 61 | def __torch_function__(cls, func, types, args=(), kwargs=None): 62 | if kwargs is None: 63 | kwargs = {} 64 | loaded_args = [ 65 | (a._load_tensor() if isinstance(a, NotYetLoadedTensor) else a) for a in args 66 | ] 67 | res = func(*loaded_args, **kwargs) 68 | # gc.collect would be costly here, maybe do it optionally 69 | return res 70 | 71 | def __getattr__(self, name): 72 | # properties 73 | ## TODO: device, is_...?? 74 | ## TODO: mH, mT, H, T, data, imag, real 75 | ## name ??? 76 | if name in { 77 | "dtype", 78 | "grad", 79 | "grad_fn", 80 | "layout", 81 | "names", 82 | "ndim", 83 | "output_nr", 84 | "requires_grad", 85 | "retains_grad", 86 | "shape", 87 | "volatile", 88 | }: 89 | return getattr(self.metatensor, name) 90 | if name in {"size"}: 91 | return getattr(self.metatensor, name) 92 | raise AttributeError(f"{type(self)} does not have {name}") 93 | 94 | def __repr__(self): 95 | return f"NotYetLoadedTensor({repr(self.metatensor)})" 96 | 97 | 98 | class LazyLoadingUnpickler(pickle.Unpickler): 99 | def __init__(self, file, zipfile, prefix): 100 | super().__init__(file) 101 | self.zipfile = zipfile 102 | self.cache = {} 103 | self.prefix = prefix 104 | 105 | def find_class(self, module, name): 106 | if module == "torch._utils" and name == "_rebuild_tensor_v2": 107 | res = super().find_class(module, name) 108 | return functools.partial(NotYetLoadedTensor.rebuild, archiveinfo=self) 109 | return super().find_class(module, name) 110 | 111 | def persistent_load(self, pid): 112 | name, cls, fn, device, size = pid 113 | with warnings.catch_warnings(): 114 | warnings.simplefilter("ignore") 115 | s = torch.storage.TypedStorage(dtype=cls().dtype, device="meta") 116 | self.cache[s._cdata] = pid 117 | return s 118 | 119 | 120 | def lazy_load(fn): 121 | zf = zipfile.ZipFile(fn) 122 | nl = zf.namelist() 123 | prefix = pathlib.Path(pathlib.Path(nl[0]).parts[0]) 124 | with zf.open(str(prefix / "data.pkl"), "r") as pkl: 125 | mup = LazyLoadingUnpickler(pkl, zf, prefix) 126 | sd = mup.load() 127 | return sd 128 | --------------------------------------------------------------------------------