├── LICENSE ├── README.md ├── poetry.lock ├── prefetch_module └── prefetch.cc ├── pyproject.toml ├── setup.py └── tensorflow_huge_model_support ├── __init__.py ├── hms.py ├── tf_keras.py ├── tf_sess.py └── utils.py /LICENSE: -------------------------------------------------------------------------------- 1 | Attribution-NonCommercial-ShareAlike 4.0 International 2 | 3 | ======================================================================= 4 | 5 | Creative Commons Corporation ("Creative Commons") is not a law firm and 6 | does not provide legal services or legal advice. Distribution of 7 | Creative Commons public licenses does not create a lawyer-client or 8 | other relationship. Creative Commons makes its licenses and related 9 | information available on an "as-is" basis. Creative Commons gives no 10 | warranties regarding its licenses, any material licensed under their 11 | terms and conditions, or any related information. Creative Commons 12 | disclaims all liability for damages resulting from their use to the 13 | fullest extent possible. 14 | 15 | Using Creative Commons Public Licenses 16 | 17 | Creative Commons public licenses provide a standard set of terms and 18 | conditions that creators and other rights holders may use to share 19 | original works of authorship and other material subject to copyright 20 | and certain other rights specified in the public license below. The 21 | following considerations are for informational purposes only, are not 22 | exhaustive, and do not form part of our licenses. 23 | 24 | Considerations for licensors: Our public licenses are 25 | intended for use by those authorized to give the public 26 | permission to use material in ways otherwise restricted by 27 | copyright and certain other rights. Our licenses are 28 | irrevocable. Licensors should read and understand the terms 29 | and conditions of the license they choose before applying it. 30 | Licensors should also secure all rights necessary before 31 | applying our licenses so that the public can reuse the 32 | material as expected. Licensors should clearly mark any 33 | material not subject to the license. This includes other CC- 34 | licensed material, or material used under an exception or 35 | limitation to copyright. More considerations for licensors: 36 | wiki.creativecommons.org/Considerations_for_licensors 37 | 38 | Considerations for the public: By using one of our public 39 | licenses, a licensor grants the public permission to use the 40 | licensed material under specified terms and conditions. If 41 | the licensor's permission is not necessary for any reason--for 42 | example, because of any applicable exception or limitation to 43 | copyright--then that use is not regulated by the license. Our 44 | licenses grant only permissions under copyright and certain 45 | other rights that a licensor has authority to grant. Use of 46 | the licensed material may still be restricted for other 47 | reasons, including because others have copyright or other 48 | rights in the material. A licensor may make special requests, 49 | such as asking that all changes be marked or described. 50 | Although not required by our licenses, you are encouraged to 51 | respect those requests where reasonable. More_considerations 52 | for the public: 53 | wiki.creativecommons.org/Considerations_for_licensees 54 | 55 | ======================================================================= 56 | 57 | Creative Commons Attribution-NonCommercial-ShareAlike 4.0 International 58 | Public License 59 | 60 | By exercising the Licensed Rights (defined below), You accept and agree 61 | to be bound by the terms and conditions of this Creative Commons 62 | Attribution-NonCommercial-ShareAlike 4.0 International Public License 63 | ("Public License"). To the extent this Public License may be 64 | interpreted as a contract, You are granted the Licensed Rights in 65 | consideration of Your acceptance of these terms and conditions, and the 66 | Licensor grants You such rights in consideration of benefits the 67 | Licensor receives from making the Licensed Material available under 68 | these terms and conditions. 69 | 70 | 71 | Section 1 -- Definitions. 72 | 73 | a. Adapted Material means material subject to Copyright and Similar 74 | Rights that is derived from or based upon the Licensed Material 75 | and in which the Licensed Material is translated, altered, 76 | arranged, transformed, or otherwise modified in a manner requiring 77 | permission under the Copyright and Similar Rights held by the 78 | Licensor. For purposes of this Public License, where the Licensed 79 | Material is a musical work, performance, or sound recording, 80 | Adapted Material is always produced where the Licensed Material is 81 | synched in timed relation with a moving image. 82 | 83 | b. Adapter's License means the license You apply to Your Copyright 84 | and Similar Rights in Your contributions to Adapted Material in 85 | accordance with the terms and conditions of this Public License. 86 | 87 | c. BY-NC-SA Compatible License means a license listed at 88 | creativecommons.org/compatiblelicenses, approved by Creative 89 | Commons as essentially the equivalent of this Public License. 90 | 91 | d. Copyright and Similar Rights means copyright and/or similar rights 92 | closely related to copyright including, without limitation, 93 | performance, broadcast, sound recording, and Sui Generis Database 94 | Rights, without regard to how the rights are labeled or 95 | categorized. For purposes of this Public License, the rights 96 | specified in Section 2(b)(1)-(2) are not Copyright and Similar 97 | Rights. 98 | 99 | e. Effective Technological Measures means those measures that, in the 100 | absence of proper authority, may not be circumvented under laws 101 | fulfilling obligations under Article 11 of the WIPO Copyright 102 | Treaty adopted on December 20, 1996, and/or similar international 103 | agreements. 104 | 105 | f. Exceptions and Limitations means fair use, fair dealing, and/or 106 | any other exception or limitation to Copyright and Similar Rights 107 | that applies to Your use of the Licensed Material. 108 | 109 | g. License Elements means the license attributes listed in the name 110 | of a Creative Commons Public License. The License Elements of this 111 | Public License are Attribution, NonCommercial, and ShareAlike. 112 | 113 | h. Licensed Material means the artistic or literary work, database, 114 | or other material to which the Licensor applied this Public 115 | License. 116 | 117 | i. Licensed Rights means the rights granted to You subject to the 118 | terms and conditions of this Public License, which are limited to 119 | all Copyright and Similar Rights that apply to Your use of the 120 | Licensed Material and that the Licensor has authority to license. 121 | 122 | j. Licensor means the individual(s) or entity(ies) granting rights 123 | under this Public License. 124 | 125 | k. NonCommercial means not primarily intended for or directed towards 126 | commercial advantage or monetary compensation. For purposes of 127 | this Public License, the exchange of the Licensed Material for 128 | other material subject to Copyright and Similar Rights by digital 129 | file-sharing or similar means is NonCommercial provided there is 130 | no payment of monetary compensation in connection with the 131 | exchange. 132 | 133 | l. Share means to provide material to the public by any means or 134 | process that requires permission under the Licensed Rights, such 135 | as reproduction, public display, public performance, distribution, 136 | dissemination, communication, or importation, and to make material 137 | available to the public including in ways that members of the 138 | public may access the material from a place and at a time 139 | individually chosen by them. 140 | 141 | m. Sui Generis Database Rights means rights other than copyright 142 | resulting from Directive 96/9/EC of the European Parliament and of 143 | the Council of 11 March 1996 on the legal protection of databases, 144 | as amended and/or succeeded, as well as other essentially 145 | equivalent rights anywhere in the world. 146 | 147 | n. You means the individual or entity exercising the Licensed Rights 148 | under this Public License. Your has a corresponding meaning. 149 | 150 | 151 | Section 2 -- Scope. 152 | 153 | a. License grant. 154 | 155 | 1. Subject to the terms and conditions of this Public License, 156 | the Licensor hereby grants You a worldwide, royalty-free, 157 | non-sublicensable, non-exclusive, irrevocable license to 158 | exercise the Licensed Rights in the Licensed Material to: 159 | 160 | a. reproduce and Share the Licensed Material, in whole or 161 | in part, for NonCommercial purposes only; and 162 | 163 | b. produce, reproduce, and Share Adapted Material for 164 | NonCommercial purposes only. 165 | 166 | 2. Exceptions and Limitations. For the avoidance of doubt, where 167 | Exceptions and Limitations apply to Your use, this Public 168 | License does not apply, and You do not need to comply with 169 | its terms and conditions. 170 | 171 | 3. Term. The term of this Public License is specified in Section 172 | 6(a). 173 | 174 | 4. Media and formats; technical modifications allowed. The 175 | Licensor authorizes You to exercise the Licensed Rights in 176 | all media and formats whether now known or hereafter created, 177 | and to make technical modifications necessary to do so. The 178 | Licensor waives and/or agrees not to assert any right or 179 | authority to forbid You from making technical modifications 180 | necessary to exercise the Licensed Rights, including 181 | technical modifications necessary to circumvent Effective 182 | Technological Measures. For purposes of this Public License, 183 | simply making modifications authorized by this Section 2(a) 184 | (4) never produces Adapted Material. 185 | 186 | 5. Downstream recipients. 187 | 188 | a. Offer from the Licensor -- Licensed Material. Every 189 | recipient of the Licensed Material automatically 190 | receives an offer from the Licensor to exercise the 191 | Licensed Rights under the terms and conditions of this 192 | Public License. 193 | 194 | b. Additional offer from the Licensor -- Adapted Material. 195 | Every recipient of Adapted Material from You 196 | automatically receives an offer from the Licensor to 197 | exercise the Licensed Rights in the Adapted Material 198 | under the conditions of the Adapter's License You apply. 199 | 200 | c. No downstream restrictions. You may not offer or impose 201 | any additional or different terms or conditions on, or 202 | apply any Effective Technological Measures to, the 203 | Licensed Material if doing so restricts exercise of the 204 | Licensed Rights by any recipient of the Licensed 205 | Material. 206 | 207 | 6. No endorsement. Nothing in this Public License constitutes or 208 | may be construed as permission to assert or imply that You 209 | are, or that Your use of the Licensed Material is, connected 210 | with, or sponsored, endorsed, or granted official status by, 211 | the Licensor or others designated to receive attribution as 212 | provided in Section 3(a)(1)(A)(i). 213 | 214 | b. Other rights. 215 | 216 | 1. Moral rights, such as the right of integrity, are not 217 | licensed under this Public License, nor are publicity, 218 | privacy, and/or other similar personality rights; however, to 219 | the extent possible, the Licensor waives and/or agrees not to 220 | assert any such rights held by the Licensor to the limited 221 | extent necessary to allow You to exercise the Licensed 222 | Rights, but not otherwise. 223 | 224 | 2. Patent and trademark rights are not licensed under this 225 | Public License. 226 | 227 | 3. To the extent possible, the Licensor waives any right to 228 | collect royalties from You for the exercise of the Licensed 229 | Rights, whether directly or through a collecting society 230 | under any voluntary or waivable statutory or compulsory 231 | licensing scheme. In all other cases the Licensor expressly 232 | reserves any right to collect such royalties, including when 233 | the Licensed Material is used other than for NonCommercial 234 | purposes. 235 | 236 | 237 | Section 3 -- License Conditions. 238 | 239 | Your exercise of the Licensed Rights is expressly made subject to the 240 | following conditions. 241 | 242 | a. Attribution. 243 | 244 | 1. If You Share the Licensed Material (including in modified 245 | form), You must: 246 | 247 | a. retain the following if it is supplied by the Licensor 248 | with the Licensed Material: 249 | 250 | i. identification of the creator(s) of the Licensed 251 | Material and any others designated to receive 252 | attribution, in any reasonable manner requested by 253 | the Licensor (including by pseudonym if 254 | designated); 255 | 256 | ii. a copyright notice; 257 | 258 | iii. a notice that refers to this Public License; 259 | 260 | iv. a notice that refers to the disclaimer of 261 | warranties; 262 | 263 | v. a URI or hyperlink to the Licensed Material to the 264 | extent reasonably practicable; 265 | 266 | b. indicate if You modified the Licensed Material and 267 | retain an indication of any previous modifications; and 268 | 269 | c. indicate the Licensed Material is licensed under this 270 | Public License, and include the text of, or the URI or 271 | hyperlink to, this Public License. 272 | 273 | 2. You may satisfy the conditions in Section 3(a)(1) in any 274 | reasonable manner based on the medium, means, and context in 275 | which You Share the Licensed Material. For example, it may be 276 | reasonable to satisfy the conditions by providing a URI or 277 | hyperlink to a resource that includes the required 278 | information. 279 | 3. If requested by the Licensor, You must remove any of the 280 | information required by Section 3(a)(1)(A) to the extent 281 | reasonably practicable. 282 | 283 | b. ShareAlike. 284 | 285 | In addition to the conditions in Section 3(a), if You Share 286 | Adapted Material You produce, the following conditions also apply. 287 | 288 | 1. The Adapter's License You apply must be a Creative Commons 289 | license with the same License Elements, this version or 290 | later, or a BY-NC-SA Compatible License. 291 | 292 | 2. You must include the text of, or the URI or hyperlink to, the 293 | Adapter's License You apply. You may satisfy this condition 294 | in any reasonable manner based on the medium, means, and 295 | context in which You Share Adapted Material. 296 | 297 | 3. You may not offer or impose any additional or different terms 298 | or conditions on, or apply any Effective Technological 299 | Measures to, Adapted Material that restrict exercise of the 300 | rights granted under the Adapter's License You apply. 301 | 302 | 303 | Section 4 -- Sui Generis Database Rights. 304 | 305 | Where the Licensed Rights include Sui Generis Database Rights that 306 | apply to Your use of the Licensed Material: 307 | 308 | a. for the avoidance of doubt, Section 2(a)(1) grants You the right 309 | to extract, reuse, reproduce, and Share all or a substantial 310 | portion of the contents of the database for NonCommercial purposes 311 | only; 312 | 313 | b. if You include all or a substantial portion of the database 314 | contents in a database in which You have Sui Generis Database 315 | Rights, then the database in which You have Sui Generis Database 316 | Rights (but not its individual contents) is Adapted Material, 317 | including for purposes of Section 3(b); and 318 | 319 | c. You must comply with the conditions in Section 3(a) if You Share 320 | all or a substantial portion of the contents of the database. 321 | 322 | For the avoidance of doubt, this Section 4 supplements and does not 323 | replace Your obligations under this Public License where the Licensed 324 | Rights include other Copyright and Similar Rights. 325 | 326 | 327 | Section 5 -- Disclaimer of Warranties and Limitation of Liability. 328 | 329 | a. UNLESS OTHERWISE SEPARATELY UNDERTAKEN BY THE LICENSOR, TO THE 330 | EXTENT POSSIBLE, THE LICENSOR OFFERS THE LICENSED MATERIAL AS-IS 331 | AND AS-AVAILABLE, AND MAKES NO REPRESENTATIONS OR WARRANTIES OF 332 | ANY KIND CONCERNING THE LICENSED MATERIAL, WHETHER EXPRESS, 333 | IMPLIED, STATUTORY, OR OTHER. THIS INCLUDES, WITHOUT LIMITATION, 334 | WARRANTIES OF TITLE, MERCHANTABILITY, FITNESS FOR A PARTICULAR 335 | PURPOSE, NON-INFRINGEMENT, ABSENCE OF LATENT OR OTHER DEFECTS, 336 | ACCURACY, OR THE PRESENCE OR ABSENCE OF ERRORS, WHETHER OR NOT 337 | KNOWN OR DISCOVERABLE. WHERE DISCLAIMERS OF WARRANTIES ARE NOT 338 | ALLOWED IN FULL OR IN PART, THIS DISCLAIMER MAY NOT APPLY TO YOU. 339 | 340 | b. TO THE EXTENT POSSIBLE, IN NO EVENT WILL THE LICENSOR BE LIABLE 341 | TO YOU ON ANY LEGAL THEORY (INCLUDING, WITHOUT LIMITATION, 342 | NEGLIGENCE) OR OTHERWISE FOR ANY DIRECT, SPECIAL, INDIRECT, 343 | INCIDENTAL, CONSEQUENTIAL, PUNITIVE, EXEMPLARY, OR OTHER LOSSES, 344 | COSTS, EXPENSES, OR DAMAGES ARISING OUT OF THIS PUBLIC LICENSE OR 345 | USE OF THE LICENSED MATERIAL, EVEN IF THE LICENSOR HAS BEEN 346 | ADVISED OF THE POSSIBILITY OF SUCH LOSSES, COSTS, EXPENSES, OR 347 | DAMAGES. WHERE A LIMITATION OF LIABILITY IS NOT ALLOWED IN FULL OR 348 | IN PART, THIS LIMITATION MAY NOT APPLY TO YOU. 349 | 350 | c. The disclaimer of warranties and limitation of liability provided 351 | above shall be interpreted in a manner that, to the extent 352 | possible, most closely approximates an absolute disclaimer and 353 | waiver of all liability. 354 | 355 | 356 | Section 6 -- Term and Termination. 357 | 358 | a. This Public License applies for the term of the Copyright and 359 | Similar Rights licensed here. However, if You fail to comply with 360 | this Public License, then Your rights under this Public License 361 | terminate automatically. 362 | 363 | b. Where Your right to use the Licensed Material has terminated under 364 | Section 6(a), it reinstates: 365 | 366 | 1. automatically as of the date the violation is cured, provided 367 | it is cured within 30 days of Your discovery of the 368 | violation; or 369 | 370 | 2. upon express reinstatement by the Licensor. 371 | 372 | For the avoidance of doubt, this Section 6(b) does not affect any 373 | right the Licensor may have to seek remedies for Your violations 374 | of this Public License. 375 | 376 | c. For the avoidance of doubt, the Licensor may also offer the 377 | Licensed Material under separate terms or conditions or stop 378 | distributing the Licensed Material at any time; however, doing so 379 | will not terminate this Public License. 380 | 381 | d. Sections 1, 5, 6, 7, and 8 survive termination of this Public 382 | License. 383 | 384 | 385 | Section 7 -- Other Terms and Conditions. 386 | 387 | a. The Licensor shall not be bound by any additional or different 388 | terms or conditions communicated by You unless expressly agreed. 389 | 390 | b. Any arrangements, understandings, or agreements regarding the 391 | Licensed Material not stated herein are separate from and 392 | independent of the terms and conditions of this Public License. 393 | 394 | 395 | Section 8 -- Interpretation. 396 | 397 | a. For the avoidance of doubt, this Public License does not, and 398 | shall not be interpreted to, reduce, limit, restrict, or impose 399 | conditions on any use of the Licensed Material that could lawfully 400 | be made without permission under this Public License. 401 | 402 | b. To the extent possible, if any provision of this Public License is 403 | deemed unenforceable, it shall be automatically reformed to the 404 | minimum extent necessary to make it enforceable. If the provision 405 | cannot be reformed, it shall be severed from this Public License 406 | without affecting the enforceability of the remaining terms and 407 | conditions. 408 | 409 | c. No term or condition of this Public License will be waived and no 410 | failure to comply consented to unless expressly agreed to by the 411 | Licensor. 412 | 413 | d. Nothing in this Public License constitutes or may be interpreted 414 | as a limitation upon, or waiver of, any privileges and immunities 415 | that apply to the Licensor or You, including from the legal 416 | processes of any jurisdiction or authority. 417 | 418 | ======================================================================= 419 | 420 | Creative Commons is not a party to its public 421 | licenses. Notwithstanding, Creative Commons may elect to apply one of 422 | its public licenses to material it publishes and in those instances 423 | will be considered the “Licensor.” The text of the Creative Commons 424 | public licenses is dedicated to the public domain under the CC0 Public 425 | Domain Dedication. Except for the limited purpose of indicating that 426 | material is shared under a Creative Commons public license or as 427 | otherwise permitted by the Creative Commons policies published at 428 | creativecommons.org/policies, Creative Commons does not authorize the 429 | use of the trademark "Creative Commons" or any other trademark or logo 430 | of Creative Commons without its prior written consent including, 431 | without limitation, in connection with any unauthorized modifications 432 | to any of its public licenses or any other arrangements, 433 | understandings, or agreements concerning use of licensed material. For 434 | the avoidance of doubt, this paragraph does not form part of the 435 | public licenses. 436 | 437 | Creative Commons may be contacted at creativecommons.org. 438 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Tensorflow Huge Model Support (HMS) 2 | 3 | This library is designed to speed up huge model training on unified memory. 4 | It takes a computation graph built by the user, conducts analysis, implements group execution and prefetch by editing the graph. 5 | A callback hook is provided to easily apply HMS on a tf.keras model. 6 | 7 | ## Publications 8 | 9 | Chen, CL., Chen, CC., Yu, WH. *et al.* An annotation-free whole-slide training approach to pathological classification of lung cancer types using deep learning. *Nat Commun* **12,** 1193 (2021). https://doi.org/10.1038/s41467-021-21467-y 10 | 11 | Chuang, WY., Chen, CC., Yu, WH. *et al.* Identification of nodal micrometastasis in colorectal cancer using deep learning on annotation-free whole-slide images. *Mod Pathol* (2021). https://doi.org/10.1038/s41379-021-00838-2 12 | 13 | ## License 14 | 15 | Copyright (C) 2021 aetherAI Co., Ltd. 16 | All rights reserved. 17 | Licensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode). 18 | 19 | ## Requirements 20 | 21 | - Tensorflow v1 (tensorflow-gpu==1.15.3) 22 | - GCC >= 7 23 | 24 | ## Installation 25 | 26 | To install HMS, simply run the following commands: 27 | ``` 28 | [CUDA_PATH=YOU_CUDA_PATH] pip install . 29 | ``` 30 | , where CUDA_PATH is /usr/local/cuda by default. 31 | 32 | ## Usage 33 | 34 | HMS can be simply applied on tf.keras model by a callback function, as described below. 35 | 36 | 1. Import HMS tf_keras module. 37 | ```python 38 | from tensorflow_huge_model_support.tf_keras import init, HMSTFKerasCallback 39 | ``` 40 | 41 | 2. Call init before model building(, and after horovod initializes). 42 | 43 | Without horovod: 44 | ```python 45 | init() 46 | ``` 47 | 48 | With horovod: 49 | ```python 50 | import horovod.tensorflow.keras as hvd 51 | hvd.init() 52 | init(hvd=hvd) 53 | ``` 54 | 55 | 3. Define a HMSKerasCallback. 56 | ```python 57 | hms_callback = HMSTFKerasCallback( 58 | hvd=hvd, 59 | default_batch_size=DEFAULT_BATCH_SIZE 60 | ) 61 | ``` 62 | , where `hvd` can be skipped if not using Horovod. 63 | 64 | 4. Pass the callback to the Keras fit or fit_generator function. 65 | ```python 66 | model.fit_generator(..., callbacks=[hms_callback] + OTHER_CALLBACKS, ...) 67 | ``` 68 | 69 | Note: Don't forget to add hvd.callbacks.BroadcastGlobalVariablesCallback(0) in the callback list if using Horovod. 70 | -------------------------------------------------------------------------------- /poetry.lock: -------------------------------------------------------------------------------- 1 | [[package]] 2 | name = "absl-py" 3 | version = "1.0.0" 4 | description = "Abseil Python Common Libraries, see https://github.com/abseil/abseil-py." 5 | category = "main" 6 | optional = false 7 | python-versions = ">=3.6" 8 | 9 | [package.dependencies] 10 | six = "*" 11 | 12 | [[package]] 13 | name = "astor" 14 | version = "0.8.1" 15 | description = "Read/rewrite/write Python ASTs" 16 | category = "main" 17 | optional = false 18 | python-versions = "!=3.0.*,!=3.1.*,!=3.2.*,!=3.3.*,>=2.7" 19 | 20 | [[package]] 21 | name = "cython" 22 | version = "0.29.24" 23 | description = "The Cython compiler for writing C extensions for the Python language." 24 | category = "main" 25 | optional = false 26 | python-versions = ">=2.6, !=3.0.*, !=3.1.*, !=3.2.*" 27 | 28 | [[package]] 29 | name = "decorator" 30 | version = "5.1.0" 31 | description = "Decorators for Humans" 32 | category = "main" 33 | optional = false 34 | python-versions = ">=3.5" 35 | 36 | [[package]] 37 | name = "gast" 38 | version = "0.2.2" 39 | description = "Python AST that abstracts the underlying Python version" 40 | category = "main" 41 | optional = false 42 | python-versions = ">=2.7, !=3.0.*, !=3.1.*, !=3.2.*, !=3.3.*" 43 | 44 | [[package]] 45 | name = "google-pasta" 46 | version = "0.2.0" 47 | description = "pasta is an AST-based Python refactoring library" 48 | category = "main" 49 | optional = false 50 | python-versions = "*" 51 | 52 | [package.dependencies] 53 | six = "*" 54 | 55 | [[package]] 56 | name = "gputil" 57 | version = "1.4.0" 58 | description = "GPUtil is a Python module for getting the GPU status from NVIDA GPUs using nvidia-smi." 59 | category = "main" 60 | optional = false 61 | python-versions = "*" 62 | 63 | [[package]] 64 | name = "grpcio" 65 | version = "1.41.1" 66 | description = "HTTP/2-based RPC framework" 67 | category = "main" 68 | optional = false 69 | python-versions = "*" 70 | 71 | [package.dependencies] 72 | six = ">=1.5.2" 73 | 74 | [package.extras] 75 | protobuf = ["grpcio-tools (>=1.41.1)"] 76 | 77 | [[package]] 78 | name = "h5py" 79 | version = "2.10.0" 80 | description = "Read and write HDF5 files from Python" 81 | category = "main" 82 | optional = false 83 | python-versions = "*" 84 | 85 | [package.dependencies] 86 | numpy = ">=1.7" 87 | six = "*" 88 | 89 | [[package]] 90 | name = "importlib-metadata" 91 | version = "4.8.2" 92 | description = "Read metadata from Python packages" 93 | category = "main" 94 | optional = false 95 | python-versions = ">=3.6" 96 | 97 | [package.dependencies] 98 | typing-extensions = {version = ">=3.6.4", markers = "python_version < \"3.8\""} 99 | zipp = ">=0.5" 100 | 101 | [package.extras] 102 | docs = ["sphinx", "jaraco.packaging (>=8.2)", "rst.linker (>=1.9)"] 103 | perf = ["ipython"] 104 | testing = ["pytest (>=6)", "pytest-checkdocs (>=2.4)", "pytest-flake8", "pytest-cov", "pytest-enabler (>=1.0.1)", "packaging", "pep517", "pyfakefs", "flufl.flake8", "pytest-perf (>=0.9.2)", "pytest-black (>=0.3.7)", "pytest-mypy", "importlib-resources (>=1.3)"] 105 | 106 | [[package]] 107 | name = "keras-applications" 108 | version = "1.0.8" 109 | description = "Reference implementations of popular deep learning models" 110 | category = "main" 111 | optional = false 112 | python-versions = "*" 113 | 114 | [package.dependencies] 115 | h5py = "*" 116 | numpy = ">=1.9.1" 117 | 118 | [package.extras] 119 | tests = ["pytest", "pytest-pep8", "pytest-xdist", "pytest-cov"] 120 | 121 | [[package]] 122 | name = "keras-preprocessing" 123 | version = "1.1.2" 124 | description = "Easy data preprocessing and data augmentation for deep learning models" 125 | category = "main" 126 | optional = false 127 | python-versions = "*" 128 | 129 | [package.dependencies] 130 | numpy = ">=1.9.1" 131 | six = ">=1.9.0" 132 | 133 | [package.extras] 134 | image = ["scipy (>=0.14)", "Pillow (>=5.2.0)"] 135 | pep8 = ["flake8"] 136 | tests = ["pandas", "pillow", "tensorflow", "keras", "pytest", "pytest-xdist", "pytest-cov"] 137 | 138 | [[package]] 139 | name = "markdown" 140 | version = "3.3.4" 141 | description = "Python implementation of Markdown." 142 | category = "main" 143 | optional = false 144 | python-versions = ">=3.6" 145 | 146 | [package.dependencies] 147 | importlib-metadata = {version = "*", markers = "python_version < \"3.8\""} 148 | 149 | [package.extras] 150 | testing = ["coverage", "pyyaml"] 151 | 152 | [[package]] 153 | name = "networkx" 154 | version = "2.2" 155 | description = "Python package for creating and manipulating graphs and networks" 156 | category = "main" 157 | optional = false 158 | python-versions = "*" 159 | 160 | [package.dependencies] 161 | decorator = ">=4.3.0" 162 | 163 | [package.extras] 164 | all = ["numpy", "scipy", "pandas", "matplotlib", "pygraphviz", "pydot", "pyyaml", "gdal", "lxml", "nose"] 165 | gdal = ["gdal"] 166 | lxml = ["lxml"] 167 | matplotlib = ["matplotlib"] 168 | nose = ["nose"] 169 | numpy = ["numpy"] 170 | pandas = ["pandas"] 171 | pydot = ["pydot"] 172 | pygraphviz = ["pygraphviz"] 173 | pyyaml = ["pyyaml"] 174 | scipy = ["scipy"] 175 | 176 | [[package]] 177 | name = "numpy" 178 | version = "1.18.5" 179 | description = "NumPy is the fundamental package for array computing with Python." 180 | category = "main" 181 | optional = false 182 | python-versions = ">=3.5" 183 | 184 | [[package]] 185 | name = "opt-einsum" 186 | version = "3.3.0" 187 | description = "Optimizing numpys einsum function" 188 | category = "main" 189 | optional = false 190 | python-versions = ">=3.5" 191 | 192 | [package.dependencies] 193 | numpy = ">=1.7" 194 | 195 | [package.extras] 196 | docs = ["sphinx (==1.2.3)", "sphinxcontrib-napoleon", "sphinx-rtd-theme", "numpydoc"] 197 | tests = ["pytest", "pytest-cov", "pytest-pep8"] 198 | 199 | [[package]] 200 | name = "protobuf" 201 | version = "3.19.1" 202 | description = "Protocol Buffers" 203 | category = "main" 204 | optional = false 205 | python-versions = ">=3.5" 206 | 207 | [[package]] 208 | name = "psutil" 209 | version = "5.8.0" 210 | description = "Cross-platform lib for process and system monitoring in Python." 211 | category = "main" 212 | optional = false 213 | python-versions = ">=2.6, !=3.0.*, !=3.1.*, !=3.2.*, !=3.3.*" 214 | 215 | [package.extras] 216 | test = ["ipaddress", "mock", "unittest2", "enum34", "pywin32", "wmi"] 217 | 218 | [[package]] 219 | name = "pycrypto" 220 | version = "2.6.1" 221 | description = "Cryptographic modules for Python." 222 | category = "main" 223 | optional = false 224 | python-versions = "*" 225 | 226 | [[package]] 227 | name = "six" 228 | version = "1.16.0" 229 | description = "Python 2 and 3 compatibility utilities" 230 | category = "main" 231 | optional = false 232 | python-versions = ">=2.7, !=3.0.*, !=3.1.*, !=3.2.*" 233 | 234 | [[package]] 235 | name = "tensorboard" 236 | version = "1.15.0" 237 | description = "TensorBoard lets you watch Tensors Flow" 238 | category = "main" 239 | optional = false 240 | python-versions = ">= 2.7, != 3.0.*, != 3.1.*" 241 | 242 | [package.dependencies] 243 | absl-py = ">=0.4" 244 | grpcio = ">=1.6.3" 245 | markdown = ">=2.6.8" 246 | numpy = ">=1.12.0" 247 | protobuf = ">=3.6.0" 248 | six = ">=1.10.0" 249 | werkzeug = ">=0.11.15" 250 | 251 | [[package]] 252 | name = "tensorflow-estimator" 253 | version = "1.15.1" 254 | description = "TensorFlow Estimator." 255 | category = "main" 256 | optional = false 257 | python-versions = "*" 258 | 259 | [[package]] 260 | name = "tensorflow-gpu" 261 | version = "1.15.5" 262 | description = "TensorFlow is an open source machine learning framework for everyone." 263 | category = "main" 264 | optional = false 265 | python-versions = "*" 266 | 267 | [package.dependencies] 268 | absl-py = ">=0.7.0" 269 | astor = ">=0.6.0" 270 | gast = "0.2.2" 271 | google-pasta = ">=0.1.6" 272 | grpcio = ">=1.8.6" 273 | h5py = "<=2.10.0" 274 | keras-applications = ">=1.0.8" 275 | keras-preprocessing = ">=1.0.5" 276 | numpy = ">=1.16.0,<1.19.0" 277 | opt-einsum = ">=2.3.2" 278 | protobuf = ">=3.6.1" 279 | six = ">=1.10.0" 280 | tensorboard = ">=1.15.0,<1.16.0" 281 | tensorflow-estimator = "1.15.1" 282 | termcolor = ">=1.1.0" 283 | wrapt = ">=1.11.1" 284 | 285 | [[package]] 286 | name = "termcolor" 287 | version = "1.1.0" 288 | description = "ANSII Color formatting for output in terminal." 289 | category = "main" 290 | optional = false 291 | python-versions = "*" 292 | 293 | [[package]] 294 | name = "typing-extensions" 295 | version = "4.0.0" 296 | description = "Backported and Experimental Type Hints for Python 3.6+" 297 | category = "main" 298 | optional = false 299 | python-versions = ">=3.6" 300 | 301 | [[package]] 302 | name = "werkzeug" 303 | version = "2.0.2" 304 | description = "The comprehensive WSGI web application library." 305 | category = "main" 306 | optional = false 307 | python-versions = ">=3.6" 308 | 309 | [package.extras] 310 | watchdog = ["watchdog"] 311 | 312 | [[package]] 313 | name = "wrapt" 314 | version = "1.13.3" 315 | description = "Module for decorators, wrappers and monkey patching." 316 | category = "main" 317 | optional = false 318 | python-versions = "!=3.0.*,!=3.1.*,!=3.2.*,!=3.3.*,!=3.4.*,>=2.7" 319 | 320 | [[package]] 321 | name = "zipp" 322 | version = "3.6.0" 323 | description = "Backport of pathlib-compatible object wrapper for zip files" 324 | category = "main" 325 | optional = false 326 | python-versions = ">=3.6" 327 | 328 | [package.extras] 329 | docs = ["sphinx", "jaraco.packaging (>=8.2)", "rst.linker (>=1.9)"] 330 | testing = ["pytest (>=4.6)", "pytest-checkdocs (>=2.4)", "pytest-flake8", "pytest-cov", "pytest-enabler (>=1.0.1)", "jaraco.itertools", "func-timeout", "pytest-black (>=0.3.7)", "pytest-mypy"] 331 | 332 | [metadata] 333 | lock-version = "1.1" 334 | python-versions = "^3.7" 335 | content-hash = "ab96671102f4bccde89a309c90a361650e2401f2443aca7a543fa0d5a734ed50" 336 | 337 | [metadata.files] 338 | absl-py = [ 339 | {file = "absl-py-1.0.0.tar.gz", hash = "sha256:ac511215c01ee9ae47b19716599e8ccfa746f2e18de72bdf641b79b22afa27ea"}, 340 | {file = "absl_py-1.0.0-py3-none-any.whl", hash = "sha256:84e6dcdc69c947d0c13e5457d056bd43cade4c2393dce00d684aedea77ddc2a3"}, 341 | ] 342 | astor = [ 343 | {file = "astor-0.8.1-py2.py3-none-any.whl", hash = "sha256:070a54e890cefb5b3739d19f30f5a5ec840ffc9c50ffa7d23cc9fc1a38ebbfc5"}, 344 | {file = "astor-0.8.1.tar.gz", hash = "sha256:6a6effda93f4e1ce9f618779b2dd1d9d84f1e32812c23a29b3fff6fd7f63fa5e"}, 345 | ] 346 | cython = [ 347 | {file = "Cython-0.29.24-cp27-cp27m-macosx_10_9_x86_64.whl", hash = "sha256:6a2cf2ccccc25413864928dfd730c29db6f63eaf98206c1e600003a445ca7f58"}, 348 | {file = "Cython-0.29.24-cp27-cp27m-manylinux1_i686.whl", hash = "sha256:b28f92e617f540d3f21f8fd479a9c6491be920ffff672a4c61b7fc4d7f749f39"}, 349 | {file = "Cython-0.29.24-cp27-cp27m-manylinux1_x86_64.whl", hash = "sha256:37bcfa5df2a3009f49624695d917c3804fccbdfcdc5eda6378754a879711a4d5"}, 350 | {file = "Cython-0.29.24-cp27-cp27m-win32.whl", hash = "sha256:9164aeef1af6f837e4fc20402a31d256188ba4d535e262c6cb78caf57ad744f8"}, 351 | {file = "Cython-0.29.24-cp27-cp27m-win_amd64.whl", hash = "sha256:73ac33a4379056a02031baa4def255717fadb9181b5ac2b244792d53eae1c925"}, 352 | {file = "Cython-0.29.24-cp27-cp27mu-manylinux1_i686.whl", hash = "sha256:09ac3087ac7a3d489ebcb3fb8402e00c13d1a3a1c6bc73fd3b0d756a3e341e79"}, 353 | {file = "Cython-0.29.24-cp27-cp27mu-manylinux1_x86_64.whl", hash = "sha256:774cb8fd931ee1ba52c472bc1c19077cd6895c1b24014ae07bb27df59aed5ebe"}, 354 | {file = "Cython-0.29.24-cp34-cp34m-win32.whl", hash = "sha256:5dd56d0be50073f0e54825a8bc3393852de0eed126339ecbca0ae149dba55cfc"}, 355 | {file = "Cython-0.29.24-cp34-cp34m-win_amd64.whl", hash = "sha256:88dc3c250dec280b0489a83950b15809762e27232f4799b1b8d0bad503f5ab84"}, 356 | {file = "Cython-0.29.24-cp35-cp35m-manylinux1_i686.whl", hash = "sha256:5fa12ebafc2f688ea6d26ab6d1d2e634a9872509ba7135b902bb0d8b368fb04b"}, 357 | {file = "Cython-0.29.24-cp35-cp35m-manylinux1_x86_64.whl", hash = "sha256:60c958bcab0ff315b4036a949bed1c65334e1f6a69e17e9966d742febb59043a"}, 358 | {file = "Cython-0.29.24-cp35-cp35m-win32.whl", hash = "sha256:166f9f29cd0058ce1a14a7b3a2458b849ed34b1ec5fd4108af3fdd2c24afcbb0"}, 359 | {file = "Cython-0.29.24-cp35-cp35m-win_amd64.whl", hash = "sha256:76cbca0188d278e93d12ebdaf5990678e6e436485fdfad49dbe9b07717d41a3c"}, 360 | {file = "Cython-0.29.24-cp36-cp36m-macosx_10_9_x86_64.whl", hash = "sha256:f2e9381497b12e8f622af620bde0d1d094035d79b899abb2ddd3a7891f535083"}, 361 | {file = "Cython-0.29.24-cp36-cp36m-manylinux1_i686.whl", hash = "sha256:d8d1a087f35e39384303f5e6b75d465d6f29d746d7138eae9d3b6e8e6f769eae"}, 362 | {file = "Cython-0.29.24-cp36-cp36m-manylinux1_x86_64.whl", hash = "sha256:112efa54a58293a4fb0acf0dd8e5b3736e95b595eee24dd88615648e445abe41"}, 363 | {file = "Cython-0.29.24-cp36-cp36m-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:4cf4452f0e4d50e11701bca38f3857fe6fa16593e7fd6a4d5f7be66f611b7da2"}, 364 | {file = "Cython-0.29.24-cp36-cp36m-win32.whl", hash = "sha256:854fe2193d3ad4c8b61932ff54d6dbe10c5fa8749eb8958d72cc0ab28243f833"}, 365 | {file = "Cython-0.29.24-cp36-cp36m-win_amd64.whl", hash = "sha256:84826ec1c11cda56261a252ddecac0c7d6b02e47e81b94f40b27b4c23c29c17c"}, 366 | {file = "Cython-0.29.24-cp37-cp37m-macosx_10_9_x86_64.whl", hash = "sha256:6ade74eece909fd3a437d9a5084829180751d7ade118e281e9824dd75eafaff2"}, 367 | {file = "Cython-0.29.24-cp37-cp37m-manylinux1_i686.whl", hash = "sha256:0a142c6b862e6ed6b02209d543062c038c110585b5e32d1ad7c9717af4f07e41"}, 368 | {file = "Cython-0.29.24-cp37-cp37m-manylinux1_x86_64.whl", hash = "sha256:10cb3def9774fa99e4583617a5616874aed3255dc241fd1f4a3c2978c78e1c53"}, 369 | {file = "Cython-0.29.24-cp37-cp37m-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:2f41ef7edd76dd23315925e003f0c58c8585f3ab24be6885c4b3f60e77c82746"}, 370 | {file = "Cython-0.29.24-cp37-cp37m-win32.whl", hash = "sha256:821c2d416ad7d006b069657ee1034c0e0cb45bdbe9ab6ab631e8c495dfcfa4ac"}, 371 | {file = "Cython-0.29.24-cp37-cp37m-win_amd64.whl", hash = "sha256:2d9e61ed1056a3b6a4b9156b62297ad18b357a7948e57a2f49b061217696567e"}, 372 | {file = "Cython-0.29.24-cp38-cp38-macosx_10_9_x86_64.whl", hash = "sha256:55b0ee28c2c8118bfb3ad9b25cf7a6cbd724e442ea96956e32ccd908d5e3e043"}, 373 | {file = "Cython-0.29.24-cp38-cp38-manylinux1_i686.whl", hash = "sha256:eb2843f8cc01c645725e6fc690a84e99cdb266ce8ebe427cf3a680ff09f876aa"}, 374 | {file = "Cython-0.29.24-cp38-cp38-manylinux1_x86_64.whl", hash = "sha256:661dbdea519d9cfb288867252b75fef73ffa8e8bb674cec27acf70646afb369b"}, 375 | {file = "Cython-0.29.24-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:bc05de569f811be1fcfde6756c9048ae518f0c4b6d9f8f024752c5365d934cac"}, 376 | {file = "Cython-0.29.24-cp38-cp38-win32.whl", hash = "sha256:a102cfa795c6b3b81a29bdb9dbec545367cd7f353c03e6f30a056fdfefd92854"}, 377 | {file = "Cython-0.29.24-cp38-cp38-win_amd64.whl", hash = "sha256:416046a98255eff97ec02077d20ebeaae52682dfca1c35aadf31260442b92514"}, 378 | {file = "Cython-0.29.24-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:ad43e684ade673565f6f9d6638015112f6c7f11aa2a632167b79014f613f0f5f"}, 379 | {file = "Cython-0.29.24-cp39-cp39-manylinux1_i686.whl", hash = "sha256:afb521523cb46ddaa8d269b421f88ea2731fee05e65b952b96d4db760f5a2a1c"}, 380 | {file = "Cython-0.29.24-cp39-cp39-manylinux1_x86_64.whl", hash = "sha256:0d414458cb22f8a90d64260da6dace5d5fcebde43f31be52ca51f818c46db8cb"}, 381 | {file = "Cython-0.29.24-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:8cb87777e82d1996aef6c146560a19270684271c9c669ba62ac6803b3cd2ff82"}, 382 | {file = "Cython-0.29.24-cp39-cp39-win32.whl", hash = "sha256:91339ee4b465924a3ea4b2a9cec7f7227bc4cadf673ce859d24c2b9ef60b1214"}, 383 | {file = "Cython-0.29.24-cp39-cp39-win_amd64.whl", hash = "sha256:5fb977945a2111f6b64501fdf7ed0ec162cc502b84457fd648d6a558ea8de0d6"}, 384 | {file = "Cython-0.29.24-py2.py3-none-any.whl", hash = "sha256:f96411f0120b5cae483923aaacd2872af8709be4b46522daedc32f051d778385"}, 385 | {file = "Cython-0.29.24.tar.gz", hash = "sha256:cdf04d07c3600860e8c2ebaad4e8f52ac3feb212453c1764a49ac08c827e8443"}, 386 | ] 387 | decorator = [ 388 | {file = "decorator-5.1.0-py3-none-any.whl", hash = "sha256:7b12e7c3c6ab203a29e157335e9122cb03de9ab7264b137594103fd4a683b374"}, 389 | {file = "decorator-5.1.0.tar.gz", hash = "sha256:e59913af105b9860aa2c8d3272d9de5a56a4e608db9a2f167a8480b323d529a7"}, 390 | ] 391 | gast = [ 392 | {file = "gast-0.2.2.tar.gz", hash = "sha256:fe939df4583692f0512161ec1c880e0a10e71e6a232da045ab8edd3756fbadf0"}, 393 | ] 394 | google-pasta = [ 395 | {file = "google-pasta-0.2.0.tar.gz", hash = "sha256:c9f2c8dfc8f96d0d5808299920721be30c9eec37f2389f28904f454565c8a16e"}, 396 | {file = "google_pasta-0.2.0-py2-none-any.whl", hash = "sha256:4612951da876b1a10fe3960d7226f0c7682cf901e16ac06e473b267a5afa8954"}, 397 | {file = "google_pasta-0.2.0-py3-none-any.whl", hash = "sha256:b32482794a366b5366a32c92a9a9201b107821889935a02b3e51f6b432ea84ed"}, 398 | ] 399 | gputil = [ 400 | {file = "GPUtil-1.4.0.tar.gz", hash = "sha256:099e52c65e512cdfa8c8763fca67f5a5c2afb63469602d5dcb4d296b3661efb9"}, 401 | ] 402 | grpcio = [ 403 | {file = "grpcio-1.41.1-cp310-cp310-linux_armv7l.whl", hash = "sha256:ead9885b53777bed4b0694ff0baea9d2c519ff774b17b177bde43d73e2b4aa38"}, 404 | {file = "grpcio-1.41.1-cp310-cp310-macosx_10_10_universal2.whl", hash = "sha256:3b4b7c1ab18283eb64af5648d20eabef9237a2aec09e30a805f18adc9497258d"}, 405 | {file = "grpcio-1.41.1-cp310-cp310-manylinux_2_17_aarch64.whl", hash = "sha256:f68aa98f5970eccb6c94456f3447a99916c42fbddae1971256bc4e7c40a6593b"}, 406 | {file = "grpcio-1.41.1-cp310-cp310-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:71d9ed5a732a54b9c87764609f2fd2bc4ae72fa85e271038eb132ea723222209"}, 407 | {file = "grpcio-1.41.1-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:0aa1af3e1480b6dd3092ee67c4b67b1ea88d638fcdc4d1a611ae11e311800b34"}, 408 | {file = "grpcio-1.41.1-cp310-cp310-win32.whl", hash = "sha256:766f1b943abc3e27842b72fba6e28fb9f57c9b84029fd7e91146e4c37034d937"}, 409 | {file = "grpcio-1.41.1-cp310-cp310-win_amd64.whl", hash = "sha256:3713e3918da6ae10812a64e75620a172f01af2ff0a1c99d6481c910e1d4a9053"}, 410 | {file = "grpcio-1.41.1-cp36-cp36m-linux_armv7l.whl", hash = "sha256:3f0b70cf8632028714a8341b841b011a47900b1c163bf5fababb4ab3888c9b6c"}, 411 | {file = "grpcio-1.41.1-cp36-cp36m-macosx_10_10_x86_64.whl", hash = "sha256:8824b36e6b0e45fefe0b4eac5ad460830e0cbc856a0c794f711289b4b8933d53"}, 412 | {file = "grpcio-1.41.1-cp36-cp36m-manylinux2010_i686.whl", hash = "sha256:788154b32bf712e9711d001df024af5f7b2522117876c129bb27b9ad6e5461fb"}, 413 | {file = "grpcio-1.41.1-cp36-cp36m-manylinux2010_x86_64.whl", hash = "sha256:c32c470e077b34a52e87e7de26644ad0f9e9ff89a785ff7e6466870869659e05"}, 414 | {file = "grpcio-1.41.1-cp36-cp36m-manylinux_2_17_aarch64.whl", hash = "sha256:a3bb4302389b23f2006ecaaea5eb4a39cc80ea98d1964159e59c1c20ef39a483"}, 415 | {file = "grpcio-1.41.1-cp36-cp36m-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:e156ea12adb7a7ca8d8280c9df850c15510b790c785fc26c9a3fb928cd221fd4"}, 416 | {file = "grpcio-1.41.1-cp36-cp36m-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:2eb8180a6d9e47fc865a4e92a2678f3202145021ef2c1bccf165fa5744f6ec95"}, 417 | {file = "grpcio-1.41.1-cp36-cp36m-win32.whl", hash = "sha256:888d8519709652dd39415de5f79abd50257201b345dd4f40151feffc3dad3232"}, 418 | {file = "grpcio-1.41.1-cp36-cp36m-win_amd64.whl", hash = "sha256:734690b3f35468f8ed4003ec7622d2d47567f1881f5fcdca34f1e52551c2ef55"}, 419 | {file = "grpcio-1.41.1-cp37-cp37m-linux_armv7l.whl", hash = "sha256:133fb9a3cf4519543e4e41eb18b5dac0da26941aeabca8122dbcf3decbad2d21"}, 420 | {file = "grpcio-1.41.1-cp37-cp37m-macosx_10_10_x86_64.whl", hash = "sha256:a5ac91db3c588296366554b2d91116fc3a9f05bae516cafae07220e1f05bfef7"}, 421 | {file = "grpcio-1.41.1-cp37-cp37m-manylinux2010_i686.whl", hash = "sha256:b8dd1b6456c6fb3681affe0f81dff4b3bc46f825fc05e086d64216545da9ad92"}, 422 | {file = "grpcio-1.41.1-cp37-cp37m-manylinux2010_x86_64.whl", hash = "sha256:9e403d07d77ed4495ad3c18994191525b11274693e72e464241c9139e2f9cd7c"}, 423 | {file = "grpcio-1.41.1-cp37-cp37m-manylinux_2_17_aarch64.whl", hash = "sha256:9d1be99f216b18f8a9dbdfbdbcc9a6caee504d0d27295fdbb5c8da35f5254a69"}, 424 | {file = "grpcio-1.41.1-cp37-cp37m-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:ebbe9582ef06559a2358827a588ab4b92a2639517de8fe428288772820ab03b5"}, 425 | {file = "grpcio-1.41.1-cp37-cp37m-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:fd570720871dc84d2adc8430ce287319c9238d1e2f70c140f9bc54c690fabd1b"}, 426 | {file = "grpcio-1.41.1-cp37-cp37m-win32.whl", hash = "sha256:0c075616d5e86fb65fd4784d5a87d6e5a1882d277dce5c33d9b67cfc71d79899"}, 427 | {file = "grpcio-1.41.1-cp37-cp37m-win_amd64.whl", hash = "sha256:9170b5d2082fc00c057c6ccd6b893033c1ade05717fcec1d63557c3bc7afdb1b"}, 428 | {file = "grpcio-1.41.1-cp38-cp38-linux_armv7l.whl", hash = "sha256:61aa02f4505c5bbbaeba80fef1bd6871d1aef05a8778a707ab91303ee0865ad0"}, 429 | {file = "grpcio-1.41.1-cp38-cp38-macosx_10_10_x86_64.whl", hash = "sha256:d1461672b2eaef9affb60a71014ebd2f789deea7c9acb1d4bd163de92dd8e044"}, 430 | {file = "grpcio-1.41.1-cp38-cp38-manylinux2010_i686.whl", hash = "sha256:c35b847bc6bd3c3a118a13277d91a772e7dd163ce7dd2791239f9941b6eaafe3"}, 431 | {file = "grpcio-1.41.1-cp38-cp38-manylinux2010_x86_64.whl", hash = "sha256:8487fb0649ebebc9c5dca1a6dc4eb7fddf701183426b3eefeb3584639d223d43"}, 432 | {file = "grpcio-1.41.1-cp38-cp38-manylinux_2_17_aarch64.whl", hash = "sha256:6ca497ccecaa8727f14c4ccc9ffb70a19c6413fe1d4650500c90a7febd662860"}, 433 | {file = "grpcio-1.41.1-cp38-cp38-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:b1232c5efc8a9e4b7a13db235c51135412beb9e62e618a2a89dd0463edb3d929"}, 434 | {file = "grpcio-1.41.1-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:31a47af7356fb5ed3120636dd75c5efb571ecf15737484119e31286687f0e52a"}, 435 | {file = "grpcio-1.41.1-cp38-cp38-win32.whl", hash = "sha256:7a22a7378ea59ad1e6f2e79f9da6862eb9e1f6586253aee784d419a49e3f4bd9"}, 436 | {file = "grpcio-1.41.1-cp38-cp38-win_amd64.whl", hash = "sha256:32b7ca83f1a6929217098aaaac89fc49879ae714c95501d40df41a0e7506164c"}, 437 | {file = "grpcio-1.41.1-cp39-cp39-linux_armv7l.whl", hash = "sha256:3213dfe3abfc3fda7f30e86aa5967dce0c2eb4cc90a0504f95434091bf6b219a"}, 438 | {file = "grpcio-1.41.1-cp39-cp39-macosx_10_10_x86_64.whl", hash = "sha256:fd11995e3402af0f838844194707da8b3235f1719bcac961493f0138f1325893"}, 439 | {file = "grpcio-1.41.1-cp39-cp39-manylinux2010_i686.whl", hash = "sha256:23a3f03e1d9ac429ff78d23d2ab07756d3728ee1a68b5f244d8435006608b6aa"}, 440 | {file = "grpcio-1.41.1-cp39-cp39-manylinux2010_x86_64.whl", hash = "sha256:fc2eadfb5ec956c556c138fab0dfc1d2395c57ae0bfea047edae1976a26b250c"}, 441 | {file = "grpcio-1.41.1-cp39-cp39-manylinux_2_17_aarch64.whl", hash = "sha256:2a34c8979de10b04a44d2cad07d41d83643e65e49f84a05b1adf150aeb41c95f"}, 442 | {file = "grpcio-1.41.1-cp39-cp39-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:72d0bdc3605dc8f4187b302e1180643963896e3f2917a52becb51afb54448e3e"}, 443 | {file = "grpcio-1.41.1-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:740f5b21a7108a8c08bf522434752dc1d306274d47ca8b4d51af5588a16b6113"}, 444 | {file = "grpcio-1.41.1-cp39-cp39-win32.whl", hash = "sha256:2f2ee78a6ae88d668ceda56fa4a18d8a38b34c2f2e1332083dd1da1a92870703"}, 445 | {file = "grpcio-1.41.1-cp39-cp39-win_amd64.whl", hash = "sha256:c3a446b6a1f8077cc03d0d496fc1cecdd3d0b66860c0c5b65cc92d0549117840"}, 446 | {file = "grpcio-1.41.1.tar.gz", hash = "sha256:9b751271b029432a526a4970dc9b70d93eb6f0963b6a841b574f780b72651969"}, 447 | ] 448 | h5py = [ 449 | {file = "h5py-2.10.0-cp27-cp27m-macosx_10_6_intel.whl", hash = "sha256:ecf4d0b56ee394a0984de15bceeb97cbe1fe485f1ac205121293fc44dcf3f31f"}, 450 | {file = "h5py-2.10.0-cp27-cp27m-manylinux1_i686.whl", hash = "sha256:86868dc07b9cc8cb7627372a2e6636cdc7a53b7e2854ad020c9e9d8a4d3fd0f5"}, 451 | {file = "h5py-2.10.0-cp27-cp27m-manylinux1_x86_64.whl", hash = "sha256:aac4b57097ac29089f179bbc2a6e14102dd210618e94d77ee4831c65f82f17c0"}, 452 | {file = "h5py-2.10.0-cp27-cp27m-win32.whl", hash = "sha256:7be5754a159236e95bd196419485343e2b5875e806fe68919e087b6351f40a70"}, 453 | {file = "h5py-2.10.0-cp27-cp27m-win_amd64.whl", hash = "sha256:13c87efa24768a5e24e360a40e0bc4c49bcb7ce1bb13a3a7f9902cec302ccd36"}, 454 | {file = "h5py-2.10.0-cp27-cp27mu-manylinux1_i686.whl", hash = "sha256:79b23f47c6524d61f899254f5cd5e486e19868f1823298bc0c29d345c2447172"}, 455 | {file = "h5py-2.10.0-cp27-cp27mu-manylinux1_x86_64.whl", hash = "sha256:cbf28ae4b5af0f05aa6e7551cee304f1d317dbed1eb7ac1d827cee2f1ef97a99"}, 456 | {file = "h5py-2.10.0-cp34-cp34m-manylinux1_i686.whl", hash = "sha256:c0d4b04bbf96c47b6d360cd06939e72def512b20a18a8547fa4af810258355d5"}, 457 | {file = "h5py-2.10.0-cp34-cp34m-manylinux1_x86_64.whl", hash = "sha256:549ad124df27c056b2e255ea1c44d30fb7a17d17676d03096ad5cd85edb32dc1"}, 458 | {file = "h5py-2.10.0-cp35-cp35m-macosx_10_6_intel.whl", hash = "sha256:a5f82cd4938ff8761d9760af3274acf55afc3c91c649c50ab18fcff5510a14a5"}, 459 | {file = "h5py-2.10.0-cp35-cp35m-manylinux1_i686.whl", hash = "sha256:3dad1730b6470fad853ef56d755d06bb916ee68a3d8272b3bab0c1ddf83bb99e"}, 460 | {file = "h5py-2.10.0-cp35-cp35m-manylinux1_x86_64.whl", hash = "sha256:063947eaed5f271679ed4ffa36bb96f57bc14f44dd4336a827d9a02702e6ce6b"}, 461 | {file = "h5py-2.10.0-cp35-cp35m-win32.whl", hash = "sha256:c54a2c0dd4957776ace7f95879d81582298c5daf89e77fb8bee7378f132951de"}, 462 | {file = "h5py-2.10.0-cp35-cp35m-win_amd64.whl", hash = "sha256:6998be619c695910cb0effe5eb15d3a511d3d1a5d217d4bd0bebad1151ec2262"}, 463 | {file = "h5py-2.10.0-cp36-cp36m-macosx_10_6_intel.whl", hash = "sha256:ff7d241f866b718e4584fa95f520cb19405220c501bd3a53ee11871ba5166ea2"}, 464 | {file = "h5py-2.10.0-cp36-cp36m-manylinux1_i686.whl", hash = "sha256:54817b696e87eb9e403e42643305f142cd8b940fe9b3b490bbf98c3b8a894cf4"}, 465 | {file = "h5py-2.10.0-cp36-cp36m-manylinux1_x86_64.whl", hash = "sha256:d3c59549f90a891691991c17f8e58c8544060fdf3ccdea267100fa5f561ff62f"}, 466 | {file = "h5py-2.10.0-cp36-cp36m-win32.whl", hash = "sha256:d7ae7a0576b06cb8e8a1c265a8bc4b73d05fdee6429bffc9a26a6eb531e79d72"}, 467 | {file = "h5py-2.10.0-cp36-cp36m-win_amd64.whl", hash = "sha256:bffbc48331b4a801d2f4b7dac8a72609f0b10e6e516e5c480a3e3241e091c878"}, 468 | {file = "h5py-2.10.0-cp37-cp37m-macosx_10_6_intel.whl", hash = "sha256:51ae56894c6c93159086ffa2c94b5b3388c0400548ab26555c143e7cfa05b8e5"}, 469 | {file = "h5py-2.10.0-cp37-cp37m-manylinux1_i686.whl", hash = "sha256:16ead3c57141101e3296ebeed79c9c143c32bdd0e82a61a2fc67e8e6d493e9d1"}, 470 | {file = "h5py-2.10.0-cp37-cp37m-manylinux1_x86_64.whl", hash = "sha256:f0e25bb91e7a02efccb50aba6591d3fe2c725479e34769802fcdd4076abfa917"}, 471 | {file = "h5py-2.10.0-cp37-cp37m-win32.whl", hash = "sha256:f23951a53d18398ef1344c186fb04b26163ca6ce449ebd23404b153fd111ded9"}, 472 | {file = "h5py-2.10.0-cp37-cp37m-win_amd64.whl", hash = "sha256:8bb1d2de101f39743f91512a9750fb6c351c032e5cd3204b4487383e34da7f75"}, 473 | {file = "h5py-2.10.0-cp38-cp38-macosx_10_9_x86_64.whl", hash = "sha256:64f74da4a1dd0d2042e7d04cf8294e04ddad686f8eba9bb79e517ae582f6668d"}, 474 | {file = "h5py-2.10.0-cp38-cp38-manylinux1_x86_64.whl", hash = "sha256:d35f7a3a6cefec82bfdad2785e78359a0e6a5fbb3f605dd5623ce88082ccd681"}, 475 | {file = "h5py-2.10.0-cp38-cp38-win32.whl", hash = "sha256:6ef7ab1089e3ef53ca099038f3c0a94d03e3560e6aff0e9d6c64c55fb13fc681"}, 476 | {file = "h5py-2.10.0-cp38-cp38-win_amd64.whl", hash = "sha256:769e141512b54dee14ec76ed354fcacfc7d97fea5a7646b709f7400cf1838630"}, 477 | {file = "h5py-2.10.0.tar.gz", hash = "sha256:84412798925dc870ffd7107f045d7659e60f5d46d1c70c700375248bf6bf512d"}, 478 | ] 479 | importlib-metadata = [ 480 | {file = "importlib_metadata-4.8.2-py3-none-any.whl", hash = "sha256:53ccfd5c134223e497627b9815d5030edf77d2ed573922f7a0b8f8bb81a1c100"}, 481 | {file = "importlib_metadata-4.8.2.tar.gz", hash = "sha256:75bdec14c397f528724c1bfd9709d660b33a4d2e77387a3358f20b848bb5e5fb"}, 482 | ] 483 | keras-applications = [ 484 | {file = "Keras_Applications-1.0.8-py3-none-any.whl", hash = "sha256:df4323692b8c1174af821bf906f1e442e63fa7589bf0f1230a0b6bdc5a810c95"}, 485 | {file = "Keras_Applications-1.0.8.tar.gz", hash = "sha256:5579f9a12bcde9748f4a12233925a59b93b73ae6947409ff34aa2ba258189fe5"}, 486 | ] 487 | keras-preprocessing = [ 488 | {file = "Keras_Preprocessing-1.1.2-py2.py3-none-any.whl", hash = "sha256:7b82029b130ff61cc99b55f3bd27427df4838576838c5b2f65940e4fcec99a7b"}, 489 | {file = "Keras_Preprocessing-1.1.2.tar.gz", hash = "sha256:add82567c50c8bc648c14195bf544a5ce7c1f76761536956c3d2978970179ef3"}, 490 | ] 491 | markdown = [ 492 | {file = "Markdown-3.3.4-py3-none-any.whl", hash = "sha256:96c3ba1261de2f7547b46a00ea8463832c921d3f9d6aba3f255a6f71386db20c"}, 493 | {file = "Markdown-3.3.4.tar.gz", hash = "sha256:31b5b491868dcc87d6c24b7e3d19a0d730d59d3e46f4eea6430a321bed387a49"}, 494 | ] 495 | networkx = [ 496 | {file = "networkx-2.2.zip", hash = "sha256:45e56f7ab6fe81652fb4bc9f44faddb0e9025f469f602df14e3b2551c2ea5c8b"}, 497 | ] 498 | numpy = [ 499 | {file = "numpy-1.18.5-cp35-cp35m-macosx_10_9_intel.whl", hash = "sha256:e91d31b34fc7c2c8f756b4e902f901f856ae53a93399368d9a0dc7be17ed2ca0"}, 500 | {file = "numpy-1.18.5-cp35-cp35m-manylinux1_i686.whl", hash = "sha256:7d42ab8cedd175b5ebcb39b5208b25ba104842489ed59fbb29356f671ac93583"}, 501 | {file = "numpy-1.18.5-cp35-cp35m-manylinux1_x86_64.whl", hash = "sha256:a78e438db8ec26d5d9d0e584b27ef25c7afa5a182d1bf4d05e313d2d6d515271"}, 502 | {file = "numpy-1.18.5-cp35-cp35m-win32.whl", hash = "sha256:a87f59508c2b7ceb8631c20630118cc546f1f815e034193dc72390db038a5cb3"}, 503 | {file = "numpy-1.18.5-cp35-cp35m-win_amd64.whl", hash = "sha256:965df25449305092b23d5145b9bdaeb0149b6e41a77a7d728b1644b3c99277c1"}, 504 | {file = "numpy-1.18.5-cp36-cp36m-macosx_10_9_x86_64.whl", hash = "sha256:ac792b385d81151bae2a5a8adb2b88261ceb4976dbfaaad9ce3a200e036753dc"}, 505 | {file = "numpy-1.18.5-cp36-cp36m-manylinux1_i686.whl", hash = "sha256:ef627986941b5edd1ed74ba89ca43196ed197f1a206a3f18cc9faf2fb84fd675"}, 506 | {file = "numpy-1.18.5-cp36-cp36m-manylinux1_x86_64.whl", hash = "sha256:f718a7949d1c4f622ff548c572e0c03440b49b9531ff00e4ed5738b459f011e8"}, 507 | {file = "numpy-1.18.5-cp36-cp36m-win32.whl", hash = "sha256:4064f53d4cce69e9ac613256dc2162e56f20a4e2d2086b1956dd2fcf77b7fac5"}, 508 | {file = "numpy-1.18.5-cp36-cp36m-win_amd64.whl", hash = "sha256:b03b2c0badeb606d1232e5f78852c102c0a7989d3a534b3129e7856a52f3d161"}, 509 | {file = "numpy-1.18.5-cp37-cp37m-macosx_10_9_x86_64.whl", hash = "sha256:a7acefddf994af1aeba05bbbafe4ba983a187079f125146dc5859e6d817df824"}, 510 | {file = "numpy-1.18.5-cp37-cp37m-manylinux1_i686.whl", hash = "sha256:cd49930af1d1e49a812d987c2620ee63965b619257bd76eaaa95870ca08837cf"}, 511 | {file = "numpy-1.18.5-cp37-cp37m-manylinux1_x86_64.whl", hash = "sha256:b39321f1a74d1f9183bf1638a745b4fd6fe80efbb1f6b32b932a588b4bc7695f"}, 512 | {file = "numpy-1.18.5-cp37-cp37m-win32.whl", hash = "sha256:cae14a01a159b1ed91a324722d746523ec757357260c6804d11d6147a9e53e3f"}, 513 | {file = "numpy-1.18.5-cp37-cp37m-win_amd64.whl", hash = "sha256:0172304e7d8d40e9e49553901903dc5f5a49a703363ed756796f5808a06fc233"}, 514 | {file = "numpy-1.18.5-cp38-cp38-macosx_10_9_x86_64.whl", hash = "sha256:e15b382603c58f24265c9c931c9a45eebf44fe2e6b4eaedbb0d025ab3255228b"}, 515 | {file = "numpy-1.18.5-cp38-cp38-manylinux1_i686.whl", hash = "sha256:3676abe3d621fc467c4c1469ee11e395c82b2d6b5463a9454e37fe9da07cd0d7"}, 516 | {file = "numpy-1.18.5-cp38-cp38-manylinux1_x86_64.whl", hash = "sha256:4674f7d27a6c1c52a4d1aa5f0881f1eff840d2206989bae6acb1c7668c02ebfb"}, 517 | {file = "numpy-1.18.5-cp38-cp38-win32.whl", hash = "sha256:9c9d6531bc1886454f44aa8f809268bc481295cf9740827254f53c30104f074a"}, 518 | {file = "numpy-1.18.5-cp38-cp38-win_amd64.whl", hash = "sha256:3dd6823d3e04b5f223e3e265b4a1eae15f104f4366edd409e5a5e413a98f911f"}, 519 | {file = "numpy-1.18.5.zip", hash = "sha256:34e96e9dae65c4839bd80012023aadd6ee2ccb73ce7fdf3074c62f301e63120b"}, 520 | ] 521 | opt-einsum = [ 522 | {file = "opt_einsum-3.3.0-py3-none-any.whl", hash = "sha256:2455e59e3947d3c275477df7f5205b30635e266fe6dc300e3d9f9646bfcea147"}, 523 | {file = "opt_einsum-3.3.0.tar.gz", hash = "sha256:59f6475f77bbc37dcf7cd748519c0ec60722e91e63ca114e68821c0c54a46549"}, 524 | ] 525 | protobuf = [ 526 | {file = "protobuf-3.19.1-cp310-cp310-macosx_10_9_universal2.whl", hash = "sha256:d80f80eb175bf5f1169139c2e0c5ada98b1c098e2b3c3736667f28cbbea39fc8"}, 527 | {file = "protobuf-3.19.1-cp310-cp310-manylinux2014_aarch64.whl", hash = "sha256:a529e7df52204565bcd33738a7a5f288f3d2d37d86caa5d78c458fa5fabbd54d"}, 528 | {file = "protobuf-3.19.1-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:28ccea56d4dc38d35cd70c43c2da2f40ac0be0a355ef882242e8586c6d66666f"}, 529 | {file = "protobuf-3.19.1-cp36-cp36m-macosx_10_9_x86_64.whl", hash = "sha256:8b30a7de128c46b5ecb343917d9fa737612a6e8280f440874e5cc2ba0d79b8f6"}, 530 | {file = "protobuf-3.19.1-cp36-cp36m-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:5935c8ce02e3d89c7900140a8a42b35bc037ec07a6aeb61cc108be8d3c9438a6"}, 531 | {file = "protobuf-3.19.1-cp36-cp36m-win32.whl", hash = "sha256:74f33edeb4f3b7ed13d567881da8e5a92a72b36495d57d696c2ea1ae0cfee80c"}, 532 | {file = "protobuf-3.19.1-cp36-cp36m-win_amd64.whl", hash = "sha256:038daf4fa38a7e818dd61f51f22588d61755160a98db087a046f80d66b855942"}, 533 | {file = "protobuf-3.19.1-cp37-cp37m-macosx_10_9_x86_64.whl", hash = "sha256:8e51561d72efd5bd5c91490af1f13e32bcba8dab4643761eb7de3ce18e64a853"}, 534 | {file = "protobuf-3.19.1-cp37-cp37m-manylinux2014_aarch64.whl", hash = "sha256:6e8ea9173403219239cdfd8d946ed101f2ab6ecc025b0fda0c6c713c35c9981d"}, 535 | {file = "protobuf-3.19.1-cp37-cp37m-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:db3532d9f7a6ebbe2392041350437953b6d7a792de10e629c1e4f5a6b1fe1ac6"}, 536 | {file = "protobuf-3.19.1-cp37-cp37m-win32.whl", hash = "sha256:615b426a177780ce381ecd212edc1e0f70db8557ed72560b82096bd36b01bc04"}, 537 | {file = "protobuf-3.19.1-cp37-cp37m-win_amd64.whl", hash = "sha256:d8919368410110633717c406ab5c97e8df5ce93020cfcf3012834f28b1fab1ea"}, 538 | {file = "protobuf-3.19.1-cp38-cp38-macosx_10_9_x86_64.whl", hash = "sha256:71b0250b0cfb738442d60cab68abc166de43411f2a4f791d31378590bfb71bd7"}, 539 | {file = "protobuf-3.19.1-cp38-cp38-manylinux2014_aarch64.whl", hash = "sha256:3cd0458870ea7d1c58e948ac8078f6ba8a7ecc44a57e03032ed066c5bb318089"}, 540 | {file = "protobuf-3.19.1-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:655264ed0d0efe47a523e2255fc1106a22f6faab7cc46cfe99b5bae085c2a13e"}, 541 | {file = "protobuf-3.19.1-cp38-cp38-win32.whl", hash = "sha256:b691d996c6d0984947c4cf8b7ae2fe372d99b32821d0584f0b90277aa36982d3"}, 542 | {file = "protobuf-3.19.1-cp38-cp38-win_amd64.whl", hash = "sha256:e7e8d2c20921f8da0dea277dfefc6abac05903ceac8e72839b2da519db69206b"}, 543 | {file = "protobuf-3.19.1-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:fd390367fc211cc0ffcf3a9e149dfeca78fecc62adb911371db0cec5c8b7472d"}, 544 | {file = "protobuf-3.19.1-cp39-cp39-manylinux2014_aarch64.whl", hash = "sha256:d83e1ef8cb74009bebee3e61cc84b1c9cd04935b72bca0cbc83217d140424995"}, 545 | {file = "protobuf-3.19.1-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:36d90676d6f426718463fe382ec6274909337ca6319d375eebd2044e6c6ac560"}, 546 | {file = "protobuf-3.19.1-cp39-cp39-win32.whl", hash = "sha256:e7b24c11df36ee8e0c085e5b0dc560289e4b58804746fb487287dda51410f1e2"}, 547 | {file = "protobuf-3.19.1-cp39-cp39-win_amd64.whl", hash = "sha256:77d2fadcf369b3f22859ab25bd12bb8e98fb11e05d9ff9b7cd45b711c719c002"}, 548 | {file = "protobuf-3.19.1-py2.py3-none-any.whl", hash = "sha256:e813b1c9006b6399308e917ac5d298f345d95bb31f46f02b60cd92970a9afa17"}, 549 | {file = "protobuf-3.19.1.tar.gz", hash = "sha256:62a8e4baa9cb9e064eb62d1002eca820857ab2138440cb4b3ea4243830f94ca7"}, 550 | ] 551 | psutil = [ 552 | {file = "psutil-5.8.0-cp27-cp27m-macosx_10_9_x86_64.whl", hash = "sha256:0066a82f7b1b37d334e68697faba68e5ad5e858279fd6351c8ca6024e8d6ba64"}, 553 | {file = "psutil-5.8.0-cp27-cp27m-manylinux2010_i686.whl", hash = "sha256:0ae6f386d8d297177fd288be6e8d1afc05966878704dad9847719650e44fc49c"}, 554 | {file = "psutil-5.8.0-cp27-cp27m-manylinux2010_x86_64.whl", hash = "sha256:12d844996d6c2b1d3881cfa6fa201fd635971869a9da945cf6756105af73d2df"}, 555 | {file = "psutil-5.8.0-cp27-cp27mu-manylinux2010_i686.whl", hash = "sha256:02b8292609b1f7fcb34173b25e48d0da8667bc85f81d7476584d889c6e0f2131"}, 556 | {file = "psutil-5.8.0-cp27-cp27mu-manylinux2010_x86_64.whl", hash = "sha256:6ffe81843131ee0ffa02c317186ed1e759a145267d54fdef1bc4ea5f5931ab60"}, 557 | {file = "psutil-5.8.0-cp27-none-win32.whl", hash = "sha256:ea313bb02e5e25224e518e4352af4bf5e062755160f77e4b1767dd5ccb65f876"}, 558 | {file = "psutil-5.8.0-cp27-none-win_amd64.whl", hash = "sha256:5da29e394bdedd9144c7331192e20c1f79283fb03b06e6abd3a8ae45ffecee65"}, 559 | {file = "psutil-5.8.0-cp36-cp36m-macosx_10_9_x86_64.whl", hash = "sha256:74fb2557d1430fff18ff0d72613c5ca30c45cdbfcddd6a5773e9fc1fe9364be8"}, 560 | {file = "psutil-5.8.0-cp36-cp36m-manylinux2010_i686.whl", hash = "sha256:74f2d0be88db96ada78756cb3a3e1b107ce8ab79f65aa885f76d7664e56928f6"}, 561 | {file = "psutil-5.8.0-cp36-cp36m-manylinux2010_x86_64.whl", hash = "sha256:99de3e8739258b3c3e8669cb9757c9a861b2a25ad0955f8e53ac662d66de61ac"}, 562 | {file = "psutil-5.8.0-cp36-cp36m-win32.whl", hash = "sha256:36b3b6c9e2a34b7d7fbae330a85bf72c30b1c827a4366a07443fc4b6270449e2"}, 563 | {file = "psutil-5.8.0-cp36-cp36m-win_amd64.whl", hash = "sha256:52de075468cd394ac98c66f9ca33b2f54ae1d9bff1ef6b67a212ee8f639ec06d"}, 564 | {file = "psutil-5.8.0-cp37-cp37m-macosx_10_9_x86_64.whl", hash = "sha256:c6a5fd10ce6b6344e616cf01cc5b849fa8103fbb5ba507b6b2dee4c11e84c935"}, 565 | {file = "psutil-5.8.0-cp37-cp37m-manylinux2010_i686.whl", hash = "sha256:61f05864b42fedc0771d6d8e49c35f07efd209ade09a5afe6a5059e7bb7bf83d"}, 566 | {file = "psutil-5.8.0-cp37-cp37m-manylinux2010_x86_64.whl", hash = "sha256:0dd4465a039d343925cdc29023bb6960ccf4e74a65ad53e768403746a9207023"}, 567 | {file = "psutil-5.8.0-cp37-cp37m-win32.whl", hash = "sha256:1bff0d07e76114ec24ee32e7f7f8d0c4b0514b3fae93e3d2aaafd65d22502394"}, 568 | {file = "psutil-5.8.0-cp37-cp37m-win_amd64.whl", hash = "sha256:fcc01e900c1d7bee2a37e5d6e4f9194760a93597c97fee89c4ae51701de03563"}, 569 | {file = "psutil-5.8.0-cp38-cp38-macosx_10_9_x86_64.whl", hash = "sha256:6223d07a1ae93f86451d0198a0c361032c4c93ebd4bf6d25e2fb3edfad9571ef"}, 570 | {file = "psutil-5.8.0-cp38-cp38-manylinux2010_i686.whl", hash = "sha256:d225cd8319aa1d3c85bf195c4e07d17d3cd68636b8fc97e6cf198f782f99af28"}, 571 | {file = "psutil-5.8.0-cp38-cp38-manylinux2010_x86_64.whl", hash = "sha256:28ff7c95293ae74bf1ca1a79e8805fcde005c18a122ca983abf676ea3466362b"}, 572 | {file = "psutil-5.8.0-cp38-cp38-win32.whl", hash = "sha256:ce8b867423291cb65cfc6d9c4955ee9bfc1e21fe03bb50e177f2b957f1c2469d"}, 573 | {file = "psutil-5.8.0-cp38-cp38-win_amd64.whl", hash = "sha256:90f31c34d25b1b3ed6c40cdd34ff122b1887a825297c017e4cbd6796dd8b672d"}, 574 | {file = "psutil-5.8.0-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:6323d5d845c2785efb20aded4726636546b26d3b577aded22492908f7c1bdda7"}, 575 | {file = "psutil-5.8.0-cp39-cp39-manylinux2010_i686.whl", hash = "sha256:245b5509968ac0bd179287d91210cd3f37add77dad385ef238b275bad35fa1c4"}, 576 | {file = "psutil-5.8.0-cp39-cp39-manylinux2010_x86_64.whl", hash = "sha256:90d4091c2d30ddd0a03e0b97e6a33a48628469b99585e2ad6bf21f17423b112b"}, 577 | {file = "psutil-5.8.0-cp39-cp39-win32.whl", hash = "sha256:ea372bcc129394485824ae3e3ddabe67dc0b118d262c568b4d2602a7070afdb0"}, 578 | {file = "psutil-5.8.0-cp39-cp39-win_amd64.whl", hash = "sha256:f4634b033faf0d968bb9220dd1c793b897ab7f1189956e1aa9eae752527127d3"}, 579 | {file = "psutil-5.8.0.tar.gz", hash = "sha256:0c9ccb99ab76025f2f0bbecf341d4656e9c1351db8cc8a03ccd62e318ab4b5c6"}, 580 | ] 581 | pycrypto = [ 582 | {file = "pycrypto-2.6.1.tar.gz", hash = "sha256:f2ce1e989b272cfcb677616763e0a2e7ec659effa67a88aa92b3a65528f60a3c"}, 583 | ] 584 | six = [ 585 | {file = "six-1.16.0-py2.py3-none-any.whl", hash = "sha256:8abb2f1d86890a2dfb989f9a77cfcfd3e47c2a354b01111771326f8aa26e0254"}, 586 | {file = "six-1.16.0.tar.gz", hash = "sha256:1e61c37477a1626458e36f7b1d82aa5c9b094fa4802892072e49de9c60c4c926"}, 587 | ] 588 | tensorboard = [ 589 | {file = "tensorboard-1.15.0-py2-none-any.whl", hash = "sha256:612b789386aa1b2c4804e1961273b37f8e4dd97613f98bc90ff0402d24627f50"}, 590 | {file = "tensorboard-1.15.0-py3-none-any.whl", hash = "sha256:4cad2c65f6011e51609b463014c014fd7c6ddd9c1263af1d4f18dd97ed88c2bc"}, 591 | ] 592 | tensorflow-estimator = [ 593 | {file = "tensorflow_estimator-1.15.1-py2.py3-none-any.whl", hash = "sha256:8853bfb7c3c96fbdc80b3d66c37a10af6ccbcd235dc87474764270c02a0f86b9"}, 594 | ] 595 | tensorflow-gpu = [ 596 | {file = "tensorflow_gpu-1.15.5-cp36-cp36m-manylinux2010_x86_64.whl", hash = "sha256:d8243cbacefe8cd77564532e48dbfa8ed98bda6f5cc05d371040ac48f0f80a7b"}, 597 | {file = "tensorflow_gpu-1.15.5-cp36-cp36m-win_amd64.whl", hash = "sha256:d909f84d863d18e7f6247ebddc44d61ad56de4c7466301565cf135f5511f0d76"}, 598 | {file = "tensorflow_gpu-1.15.5-cp37-cp37m-manylinux2010_x86_64.whl", hash = "sha256:d0f9afbdc4d3d66870e626462733ac0ea92757b0c2c92279305d13b12a2f84e8"}, 599 | {file = "tensorflow_gpu-1.15.5-cp37-cp37m-win_amd64.whl", hash = "sha256:b0ac2858a5a0e799bdae1e5b86df46d6b33df79bd6cd8e36791fb225c672a39f"}, 600 | ] 601 | termcolor = [ 602 | {file = "termcolor-1.1.0.tar.gz", hash = "sha256:1d6d69ce66211143803fbc56652b41d73b4a400a2891d7bf7a1cdf4c02de613b"}, 603 | ] 604 | typing-extensions = [ 605 | {file = "typing_extensions-4.0.0-py3-none-any.whl", hash = "sha256:829704698b22e13ec9eaf959122315eabb370b0884400e9818334d8b677023d9"}, 606 | ] 607 | werkzeug = [ 608 | {file = "Werkzeug-2.0.2-py3-none-any.whl", hash = "sha256:63d3dc1cf60e7b7e35e97fa9861f7397283b75d765afcaefd993d6046899de8f"}, 609 | {file = "Werkzeug-2.0.2.tar.gz", hash = "sha256:aa2bb6fc8dee8d6c504c0ac1e7f5f7dc5810a9903e793b6f715a9f015bdadb9a"}, 610 | ] 611 | wrapt = [ 612 | {file = "wrapt-1.13.3-cp27-cp27m-macosx_10_9_x86_64.whl", hash = "sha256:e05e60ff3b2b0342153be4d1b597bbcfd8330890056b9619f4ad6b8d5c96a81a"}, 613 | {file = "wrapt-1.13.3-cp27-cp27m-manylinux1_i686.whl", hash = "sha256:85148f4225287b6a0665eef08a178c15097366d46b210574a658c1ff5b377489"}, 614 | {file = "wrapt-1.13.3-cp27-cp27m-manylinux1_x86_64.whl", hash = "sha256:2dded5496e8f1592ec27079b28b6ad2a1ef0b9296d270f77b8e4a3a796cf6909"}, 615 | {file = "wrapt-1.13.3-cp27-cp27m-manylinux2010_i686.whl", hash = "sha256:e94b7d9deaa4cc7bac9198a58a7240aaf87fe56c6277ee25fa5b3aa1edebd229"}, 616 | {file = "wrapt-1.13.3-cp27-cp27m-manylinux2010_x86_64.whl", hash = "sha256:498e6217523111d07cd67e87a791f5e9ee769f9241fcf8a379696e25806965af"}, 617 | {file = "wrapt-1.13.3-cp27-cp27mu-manylinux1_i686.whl", hash = "sha256:ec7e20258ecc5174029a0f391e1b948bf2906cd64c198a9b8b281b811cbc04de"}, 618 | {file = "wrapt-1.13.3-cp27-cp27mu-manylinux1_x86_64.whl", hash = "sha256:87883690cae293541e08ba2da22cacaae0a092e0ed56bbba8d018cc486fbafbb"}, 619 | {file = "wrapt-1.13.3-cp27-cp27mu-manylinux2010_i686.whl", hash = "sha256:f99c0489258086308aad4ae57da9e8ecf9e1f3f30fa35d5e170b4d4896554d80"}, 620 | {file = "wrapt-1.13.3-cp27-cp27mu-manylinux2010_x86_64.whl", hash = "sha256:6a03d9917aee887690aa3f1747ce634e610f6db6f6b332b35c2dd89412912bca"}, 621 | {file = "wrapt-1.13.3-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:936503cb0a6ed28dbfa87e8fcd0a56458822144e9d11a49ccee6d9a8adb2ac44"}, 622 | {file = "wrapt-1.13.3-cp310-cp310-manylinux_2_5_i686.manylinux1_i686.manylinux_2_12_i686.manylinux2010_i686.whl", hash = "sha256:f9c51d9af9abb899bd34ace878fbec8bf357b3194a10c4e8e0a25512826ef056"}, 623 | {file = "wrapt-1.13.3-cp310-cp310-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_12_x86_64.manylinux2010_x86_64.whl", hash = "sha256:220a869982ea9023e163ba915077816ca439489de6d2c09089b219f4e11b6785"}, 624 | {file = "wrapt-1.13.3-cp310-cp310-musllinux_1_1_i686.whl", hash = "sha256:0877fe981fd76b183711d767500e6b3111378ed2043c145e21816ee589d91096"}, 625 | {file = "wrapt-1.13.3-cp310-cp310-musllinux_1_1_x86_64.whl", hash = "sha256:43e69ffe47e3609a6aec0fe723001c60c65305784d964f5007d5b4fb1bc6bf33"}, 626 | {file = "wrapt-1.13.3-cp310-cp310-win32.whl", hash = "sha256:78dea98c81915bbf510eb6a3c9c24915e4660302937b9ae05a0947164248020f"}, 627 | {file = "wrapt-1.13.3-cp310-cp310-win_amd64.whl", hash = "sha256:ea3e746e29d4000cd98d572f3ee2a6050a4f784bb536f4ac1f035987fc1ed83e"}, 628 | {file = "wrapt-1.13.3-cp35-cp35m-manylinux1_i686.whl", hash = "sha256:8c73c1a2ec7c98d7eaded149f6d225a692caa1bd7b2401a14125446e9e90410d"}, 629 | {file = "wrapt-1.13.3-cp35-cp35m-manylinux1_x86_64.whl", hash = "sha256:086218a72ec7d986a3eddb7707c8c4526d677c7b35e355875a0fe2918b059179"}, 630 | {file = "wrapt-1.13.3-cp35-cp35m-manylinux2010_i686.whl", hash = "sha256:e92d0d4fa68ea0c02d39f1e2f9cb5bc4b4a71e8c442207433d8db47ee79d7aa3"}, 631 | {file = "wrapt-1.13.3-cp35-cp35m-manylinux2010_x86_64.whl", hash = "sha256:d4a5f6146cfa5c7ba0134249665acd322a70d1ea61732723c7d3e8cc0fa80755"}, 632 | {file = "wrapt-1.13.3-cp35-cp35m-win32.whl", hash = "sha256:8aab36778fa9bba1a8f06a4919556f9f8c7b33102bd71b3ab307bb3fecb21851"}, 633 | {file = "wrapt-1.13.3-cp35-cp35m-win_amd64.whl", hash = "sha256:944b180f61f5e36c0634d3202ba8509b986b5fbaf57db3e94df11abee244ba13"}, 634 | {file = "wrapt-1.13.3-cp36-cp36m-macosx_10_9_x86_64.whl", hash = "sha256:2ebdde19cd3c8cdf8df3fc165bc7827334bc4e353465048b36f7deeae8ee0918"}, 635 | {file = "wrapt-1.13.3-cp36-cp36m-manylinux_2_5_i686.manylinux1_i686.manylinux_2_12_i686.manylinux2010_i686.whl", hash = "sha256:610f5f83dd1e0ad40254c306f4764fcdc846641f120c3cf424ff57a19d5f7ade"}, 636 | {file = "wrapt-1.13.3-cp36-cp36m-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_12_x86_64.manylinux2010_x86_64.whl", hash = "sha256:5601f44a0f38fed36cc07db004f0eedeaadbdcec90e4e90509480e7e6060a5bc"}, 637 | {file = "wrapt-1.13.3-cp36-cp36m-musllinux_1_1_i686.whl", hash = "sha256:e6906d6f48437dfd80464f7d7af1740eadc572b9f7a4301e7dd3d65db285cacf"}, 638 | {file = "wrapt-1.13.3-cp36-cp36m-musllinux_1_1_x86_64.whl", hash = "sha256:766b32c762e07e26f50d8a3468e3b4228b3736c805018e4b0ec8cc01ecd88125"}, 639 | {file = "wrapt-1.13.3-cp36-cp36m-win32.whl", hash = "sha256:5f223101f21cfd41deec8ce3889dc59f88a59b409db028c469c9b20cfeefbe36"}, 640 | {file = "wrapt-1.13.3-cp36-cp36m-win_amd64.whl", hash = "sha256:f122ccd12fdc69628786d0c947bdd9cb2733be8f800d88b5a37c57f1f1d73c10"}, 641 | {file = "wrapt-1.13.3-cp37-cp37m-macosx_10_9_x86_64.whl", hash = "sha256:46f7f3af321a573fc0c3586612db4decb7eb37172af1bc6173d81f5b66c2e068"}, 642 | {file = "wrapt-1.13.3-cp37-cp37m-manylinux_2_5_i686.manylinux1_i686.manylinux_2_12_i686.manylinux2010_i686.whl", hash = "sha256:778fd096ee96890c10ce96187c76b3e99b2da44e08c9e24d5652f356873f6709"}, 643 | {file = "wrapt-1.13.3-cp37-cp37m-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_12_x86_64.manylinux2010_x86_64.whl", hash = "sha256:0cb23d36ed03bf46b894cfec777eec754146d68429c30431c99ef28482b5c1df"}, 644 | {file = "wrapt-1.13.3-cp37-cp37m-musllinux_1_1_i686.whl", hash = "sha256:96b81ae75591a795d8c90edc0bfaab44d3d41ffc1aae4d994c5aa21d9b8e19a2"}, 645 | {file = "wrapt-1.13.3-cp37-cp37m-musllinux_1_1_x86_64.whl", hash = "sha256:7dd215e4e8514004c8d810a73e342c536547038fb130205ec4bba9f5de35d45b"}, 646 | {file = "wrapt-1.13.3-cp37-cp37m-win32.whl", hash = "sha256:47f0a183743e7f71f29e4e21574ad3fa95676136f45b91afcf83f6a050914829"}, 647 | {file = "wrapt-1.13.3-cp37-cp37m-win_amd64.whl", hash = "sha256:fd76c47f20984b43d93de9a82011bb6e5f8325df6c9ed4d8310029a55fa361ea"}, 648 | {file = "wrapt-1.13.3-cp38-cp38-macosx_10_9_x86_64.whl", hash = "sha256:b73d4b78807bd299b38e4598b8e7bd34ed55d480160d2e7fdaabd9931afa65f9"}, 649 | {file = "wrapt-1.13.3-cp38-cp38-manylinux_2_5_i686.manylinux1_i686.manylinux_2_12_i686.manylinux2010_i686.whl", hash = "sha256:ec9465dd69d5657b5d2fa6133b3e1e989ae27d29471a672416fd729b429eb554"}, 650 | {file = "wrapt-1.13.3-cp38-cp38-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_12_x86_64.manylinux2010_x86_64.whl", hash = "sha256:dd91006848eb55af2159375134d724032a2d1d13bcc6f81cd8d3ed9f2b8e846c"}, 651 | {file = "wrapt-1.13.3-cp38-cp38-musllinux_1_1_i686.whl", hash = "sha256:ae9de71eb60940e58207f8e71fe113c639da42adb02fb2bcbcaccc1ccecd092b"}, 652 | {file = "wrapt-1.13.3-cp38-cp38-musllinux_1_1_x86_64.whl", hash = "sha256:51799ca950cfee9396a87f4a1240622ac38973b6df5ef7a41e7f0b98797099ce"}, 653 | {file = "wrapt-1.13.3-cp38-cp38-win32.whl", hash = "sha256:4b9c458732450ec42578b5642ac53e312092acf8c0bfce140ada5ca1ac556f79"}, 654 | {file = "wrapt-1.13.3-cp38-cp38-win_amd64.whl", hash = "sha256:7dde79d007cd6dfa65afe404766057c2409316135cb892be4b1c768e3f3a11cb"}, 655 | {file = "wrapt-1.13.3-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:981da26722bebb9247a0601e2922cedf8bb7a600e89c852d063313102de6f2cb"}, 656 | {file = "wrapt-1.13.3-cp39-cp39-manylinux_2_5_i686.manylinux1_i686.manylinux_2_12_i686.manylinux2010_i686.whl", hash = "sha256:705e2af1f7be4707e49ced9153f8d72131090e52be9278b5dbb1498c749a1e32"}, 657 | {file = "wrapt-1.13.3-cp39-cp39-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_12_x86_64.manylinux2010_x86_64.whl", hash = "sha256:25b1b1d5df495d82be1c9d2fad408f7ce5ca8a38085e2da41bb63c914baadff7"}, 658 | {file = "wrapt-1.13.3-cp39-cp39-musllinux_1_1_i686.whl", hash = "sha256:77416e6b17926d953b5c666a3cb718d5945df63ecf922af0ee576206d7033b5e"}, 659 | {file = "wrapt-1.13.3-cp39-cp39-musllinux_1_1_x86_64.whl", hash = "sha256:865c0b50003616f05858b22174c40ffc27a38e67359fa1495605f96125f76640"}, 660 | {file = "wrapt-1.13.3-cp39-cp39-win32.whl", hash = "sha256:0a017a667d1f7411816e4bf214646d0ad5b1da2c1ea13dec6c162736ff25a374"}, 661 | {file = "wrapt-1.13.3-cp39-cp39-win_amd64.whl", hash = "sha256:81bd7c90d28a4b2e1df135bfbd7c23aee3050078ca6441bead44c42483f9ebfb"}, 662 | {file = "wrapt-1.13.3.tar.gz", hash = "sha256:1fea9cd438686e6682271d36f3481a9f3636195578bab9ca3382e2f5f01fc185"}, 663 | ] 664 | zipp = [ 665 | {file = "zipp-3.6.0-py3-none-any.whl", hash = "sha256:9fe5ea21568a0a70e50f273397638d39b03353731e6cbbb3fd8502a33fec40bc"}, 666 | {file = "zipp-3.6.0.tar.gz", hash = "sha256:71c644c5369f4a6e07636f0aa966270449561fcea2e3d6747b8d23efaa9d7832"}, 667 | ] 668 | -------------------------------------------------------------------------------- /prefetch_module/prefetch.cc: -------------------------------------------------------------------------------- 1 | #include "tensorflow/core/framework/common_shape_fns.h" 2 | #include "tensorflow/core/framework/op.h" 3 | #include "tensorflow/core/framework/op_kernel.h" 4 | #include "tensorflow/core/framework/shape_inference.h" 5 | #include 6 | #include 7 | 8 | #include 9 | #include 10 | #include 11 | #include 12 | 13 | using namespace tensorflow; 14 | 15 | REGISTER_OP("GetTensorAddr") 16 | .Input("tensor: T") 17 | .Output("addr: int64") 18 | .Attr("T: type") 19 | .SetShapeFn([](shape_inference::InferenceContext* c) { 20 | c->set_output(0, c->Vector(2)); 21 | return Status::OK(); 22 | }); 23 | 24 | class GetTensorAddrOp : public OpKernel { 25 | public: 26 | explicit GetTensorAddrOp(OpKernelConstruction* context) 27 | : OpKernel(context) 28 | {} 29 | 30 | void Compute(OpKernelContext* context) override 31 | { 32 | const Tensor& input_tensor = context->input(0); 33 | StringPiece input_tensor_data = input_tensor.tensor_data(); 34 | CUdeviceptr device_ptr = reinterpret_cast(input_tensor_data.data()); 35 | size_t size = input_tensor_data.size(); 36 | 37 | Tensor *addr_tensor = nullptr; 38 | TensorShape shape = TensorShape(); 39 | shape.AddDim(2); 40 | context->allocate_output(0, shape, &addr_tensor); 41 | addr_tensor->vec()(0) = static_cast(device_ptr); 42 | addr_tensor->vec()(1) = static_cast(size); 43 | } 44 | }; 45 | 46 | REGISTER_KERNEL_BUILDER(Name("GetTensorAddr").Device(DEVICE_GPU).HostMemory("addr"), GetTensorAddrOp); 47 | 48 | class StreamSingleton { 49 | private: 50 | static StreamSingleton* instance; 51 | static std::once_flag init_instance_flag; 52 | 53 | std::map, CUstream *> stream_list; 54 | std::map, std::mutex *> stream_mutex_list; 55 | std::mutex list_mutex; 56 | 57 | StreamSingleton() 58 | {} 59 | 60 | public: 61 | static StreamSingleton& getInstance() 62 | { 63 | std::call_once( 64 | init_instance_flag, 65 | [&]() -> void { 66 | instance = new StreamSingleton; 67 | } 68 | ); 69 | return *instance; 70 | } 71 | 72 | void claim(bool is_prefetch, CUdevice device) 73 | { 74 | std::mutex *stream_mutex_ptr; 75 | std::tuple tup(is_prefetch, device); 76 | list_mutex.lock(); 77 | if (stream_mutex_list.count(tup) == 0) { 78 | stream_mutex_list[tup] = new std::mutex; 79 | } 80 | stream_mutex_ptr = stream_mutex_list[tup]; 81 | list_mutex.unlock(); 82 | 83 | stream_mutex_ptr->lock(); 84 | } 85 | 86 | void release(bool is_prefetch, CUdevice device) 87 | { 88 | std::tuple tup(is_prefetch, device); 89 | list_mutex.lock(); 90 | stream_mutex_list[tup]->unlock(); 91 | list_mutex.unlock(); 92 | } 93 | 94 | CUstream& getStream(bool is_prefetch, CUdevice device) 95 | { 96 | CUstream *stream_ptr; 97 | std::tuple tup(is_prefetch, device); 98 | list_mutex.lock(); 99 | if (stream_list.count(tup) == 0) { 100 | stream_list[tup] = new CUstream; 101 | cuStreamCreate(stream_list[tup], CU_STREAM_NON_BLOCKING); 102 | } 103 | stream_ptr = stream_list[tup]; 104 | list_mutex.unlock(); 105 | return *stream_ptr; 106 | } 107 | }; 108 | 109 | StreamSingleton* StreamSingleton::instance = nullptr; 110 | std::once_flag StreamSingleton::init_instance_flag; 111 | 112 | struct MemRange { 113 | static const int64 PAGE_SIZE; 114 | int64 addr; 115 | int64 size; 116 | 117 | MemRange(int64 addr, int64 size) 118 | : addr(addr), size(size) 119 | {} 120 | 121 | bool operator<(const MemRange& rhs) const 122 | { 123 | return this->addr < rhs.addr; 124 | } 125 | 126 | bool isOverlapped(const MemRange& rhs) const 127 | { 128 | return !(rhs.addr > this->addr + this->size) && !(this->addr > rhs.addr + rhs.size); 129 | } 130 | 131 | void mergeWith(const MemRange& rhs) 132 | { 133 | int64 new_addr = std::min(this->addr, rhs.addr); 134 | int64 new_addr_end = std::max(this->addr + this->size, rhs.addr + rhs.size); 135 | int64 new_size = new_addr_end - new_addr; 136 | 137 | this->addr = new_addr; 138 | this->size = new_size; 139 | } 140 | 141 | MemRange &toPageAligned(bool expand_or_shrink) 142 | { 143 | if (expand_or_shrink) { 144 | int64 begin_addr = addr - addr % PAGE_SIZE; 145 | int64 end_addr = addr + size; 146 | end_addr = end_addr - end_addr % PAGE_SIZE + ((end_addr % PAGE_SIZE == 0) ? 0 : PAGE_SIZE); 147 | 148 | addr = begin_addr; 149 | size = end_addr - begin_addr; 150 | } 151 | else { 152 | int64 addr_offset = addr % PAGE_SIZE; 153 | int64 begin_addr = (addr_offset == 0 ? addr : addr - addr_offset + PAGE_SIZE); 154 | int64 end_addr = addr + size; 155 | int64 end_addr_offset = end_addr % PAGE_SIZE; 156 | end_addr = (end_addr_offset == 0 ? end_addr : end_addr - end_addr_offset); 157 | 158 | addr = begin_addr; 159 | size = end_addr - begin_addr; 160 | } 161 | 162 | return *this; 163 | } 164 | }; 165 | 166 | //const int64 MemRange::PAGE_SIZE = sysconf(_SC_PAGESIZE); 167 | const int64 MemRange::PAGE_SIZE = 2LL * 1024LL * 1024LL; 168 | 169 | REGISTER_OP("Prefetch") 170 | .Input("addrs: int64") 171 | .Attr("expand_or_shrink: bool") 172 | .SetShapeFn(shape_inference::NoOutputs); 173 | 174 | class PrefetchOp : public AsyncOpKernel { 175 | private: 176 | bool expand_or_shrink; 177 | 178 | public: 179 | explicit PrefetchOp(OpKernelConstruction* context) 180 | : AsyncOpKernel(context) 181 | { 182 | context->GetAttr("expand_or_shrink", &expand_or_shrink); 183 | } 184 | 185 | void ComputeAsync(OpKernelContext* context, DoneCallback done) override 186 | { 187 | const Tensor& addrs = context->input(0); 188 | 189 | // Merge intervals of requested addresses 190 | int64 num_range = addrs.dim_size(0); 191 | auto addrs_tensor = addrs.tensor(); 192 | 193 | const int TAG_BEGIN = 0; 194 | const int TAG_END = 1; 195 | 196 | struct AddrWithTag { 197 | int64 addr; 198 | int tag; 199 | 200 | AddrWithTag(int64 addr, int tag) 201 | : addr(addr), tag(tag) 202 | {} 203 | 204 | bool operator<(const AddrWithTag &rhs) const 205 | { 206 | return this->addr < rhs.addr; 207 | } 208 | }; 209 | 210 | std::vector addr_list; 211 | for (int64 i = 0; i < num_range; i++) { 212 | if (addrs_tensor(i, 1) <= 0) 213 | continue; 214 | 215 | MemRange mem_range(addrs_tensor(i, 0), addrs_tensor(i, 1)); 216 | mem_range.toPageAligned(expand_or_shrink); 217 | 218 | addr_list.emplace_back(mem_range.addr, TAG_BEGIN); 219 | addr_list.emplace_back(mem_range.addr + mem_range.size, TAG_END); 220 | } 221 | std::sort(addr_list.begin(), addr_list.end()); 222 | 223 | std::vector range_list; 224 | int depth = 0; 225 | int64 addr_begin = 0LL; 226 | for (AddrWithTag &addr_with_tag : addr_list) { 227 | int64 addr = addr_with_tag.addr; 228 | int tag = addr_with_tag.tag; 229 | 230 | switch (tag) { 231 | case TAG_BEGIN: 232 | depth += 1; 233 | if (depth == 1) { 234 | addr_begin = addr; 235 | } 236 | break; 237 | 238 | case TAG_END: 239 | depth -= 1; 240 | if (depth == 0) { 241 | range_list.emplace_back(addr_begin, addr - addr_begin); 242 | } 243 | break; 244 | 245 | default: 246 | assert(false); 247 | } 248 | } 249 | 250 | // Issue prefetches 251 | const tensorflow::DeviceBase::GpuDeviceInfo *device_info = 252 | context->device()->tensorflow_gpu_device_info(); 253 | int device_ordinal = device_info->gpu_id; 254 | CUdevice cu_device = -1; 255 | cuDeviceGet(&cu_device, device_ordinal); 256 | 257 | CUevent event; 258 | cuEventCreate(&event, CU_EVENT_DEFAULT); 259 | StreamSingleton &stream_singleton = StreamSingleton::getInstance(); 260 | stream_singleton.claim(true, cu_device); 261 | CUstream &cu_stream = stream_singleton.getStream(true, cu_device); 262 | 263 | for (MemRange& range : range_list) { 264 | //std::cerr << "Prefetch: " << range.addr << ", " << range.size << std::endl; 265 | cuMemPrefetchAsync(range.addr, range.size, cu_device, cu_stream); 266 | } 267 | cuEventRecord(event, cu_stream); 268 | stream_singleton.release(true, cu_device); 269 | cuEventSynchronize(event); 270 | cuEventDestroy(event); 271 | //std::cerr << "---" << std::endl; 272 | 273 | done(); 274 | } 275 | }; 276 | 277 | REGISTER_KERNEL_BUILDER(Name("Prefetch").Device(DEVICE_GPU).HostMemory("addrs"), PrefetchOp); 278 | 279 | REGISTER_OP("Evict") 280 | .Input("addrs: int64") 281 | .Input("exclude_addrs: int64") 282 | .SetShapeFn(shape_inference::NoOutputs); 283 | 284 | class EvictOp : public AsyncOpKernel { 285 | public: 286 | explicit EvictOp(OpKernelConstruction* context) 287 | : AsyncOpKernel(context) 288 | {} 289 | 290 | void ComputeAsync(OpKernelContext* context, DoneCallback done) override 291 | { 292 | const Tensor& addrs = context->input(0); 293 | const Tensor& exclude_addrs = context->input(1); 294 | 295 | // Extract addresses 296 | const int TAG_EVICT_BEGIN = 0; 297 | const int TAG_EVICT_END = 1; 298 | const int TAG_EXCLUDE_BEGIN = 2; 299 | const int TAG_EXCLUDE_END = 3; 300 | 301 | int64 num_addrs = addrs.dim_size(0); 302 | auto addrs_tensor = addrs.tensor(); 303 | int64 num_exclude_addrs = exclude_addrs.dim_size(0); 304 | auto exclude_addrs_tensor = exclude_addrs.tensor(); 305 | 306 | struct AddrWithTag { 307 | int64 addr; 308 | int tag; 309 | 310 | AddrWithTag(int64 addr, int tag) 311 | : addr(addr), tag(tag) 312 | {} 313 | 314 | bool operator<(const AddrWithTag &rhs) const 315 | { 316 | return this->addr < rhs.addr; 317 | } 318 | }; 319 | 320 | std::vector addr_list; 321 | for (int64 i = 0; i < num_addrs; i++) { 322 | if (addrs_tensor(i, 1) <= 0) 323 | // Remove size <= 0 requests 324 | continue; 325 | 326 | MemRange mem_range(addrs_tensor(i, 0), addrs_tensor(i, 1)); 327 | mem_range.toPageAligned(false); 328 | 329 | AddrWithTag addr_begin(mem_range.addr, TAG_EVICT_BEGIN); 330 | addr_list.push_back(addr_begin); 331 | AddrWithTag addr_end(mem_range.addr + mem_range.size, TAG_EVICT_END); 332 | addr_list.push_back(addr_end); 333 | } 334 | for (int64 i = 0; i < num_exclude_addrs; i++) { 335 | if (exclude_addrs_tensor(i, 1) <= 0) 336 | // Remove size <= 0 requests 337 | continue; 338 | 339 | MemRange mem_range(exclude_addrs_tensor(i, 0), exclude_addrs_tensor(i, 1)); 340 | mem_range.toPageAligned(false); 341 | 342 | AddrWithTag addr_begin(mem_range.addr, TAG_EXCLUDE_BEGIN); 343 | addr_list.push_back(addr_begin); 344 | AddrWithTag addr_end(mem_range.addr + mem_range.size, TAG_EXCLUDE_END); 345 | addr_list.push_back(addr_end); 346 | } 347 | 348 | std::sort(addr_list.begin(), addr_list.end()); 349 | 350 | // Deal with exclusions 351 | std::vector range_list; 352 | int evict_depth = 0; 353 | int exclude_depth = 0; 354 | int64 evict_begin = 0LL; 355 | for (int64 i = 0; i < addr_list.size(); i++) { 356 | AddrWithTag &addr_with_tag = addr_list[i]; 357 | int tag = addr_with_tag.tag; 358 | int64 addr = addr_with_tag.addr; 359 | 360 | switch (tag) { 361 | case TAG_EVICT_BEGIN: 362 | evict_depth += 1; 363 | if (evict_depth == 1 && exclude_depth == 0) 364 | evict_begin = addr; 365 | break; 366 | 367 | case TAG_EVICT_END: 368 | evict_depth -= 1; 369 | if (evict_depth == 0 && exclude_depth == 0 && addr - evict_begin > 0) 370 | range_list.emplace_back(evict_begin, addr - evict_begin); 371 | break; 372 | 373 | case TAG_EXCLUDE_BEGIN: 374 | exclude_depth += 1; 375 | if (exclude_depth == 1 && evict_depth > 0 && addr - evict_begin > 0) 376 | range_list.emplace_back(evict_begin, addr - evict_begin); 377 | break; 378 | 379 | case TAG_EXCLUDE_END: 380 | exclude_depth -= 1; 381 | if (exclude_depth == 0 && evict_depth > 0) 382 | evict_begin = addr; 383 | break; 384 | 385 | default: 386 | assert(false); 387 | } 388 | } 389 | 390 | // Issue evictions 391 | const tensorflow::DeviceBase::GpuDeviceInfo *device_info = 392 | context->device()->tensorflow_gpu_device_info(); 393 | int device_ordinal = device_info->gpu_id; 394 | CUdevice cu_device = -1; 395 | cuDeviceGet(&cu_device, device_ordinal); 396 | 397 | CUevent event; 398 | cuEventCreate(&event, CU_EVENT_DEFAULT); 399 | StreamSingleton &stream_singleton = StreamSingleton::getInstance(); 400 | stream_singleton.claim(false, cu_device); 401 | CUstream &cu_stream = stream_singleton.getStream(false, cu_device); 402 | 403 | for (MemRange& range : range_list) { 404 | CUresult res = cuMemPrefetchAsync(range.addr, range.size, CU_DEVICE_CPU, cu_stream); 405 | //std::cerr << res << " " << range.addr << " " << range.size << std::endl; 406 | } 407 | cuEventRecord(event, cu_stream); 408 | stream_singleton.release(false, cu_device); 409 | cuEventSynchronize(event); 410 | cuEventDestroy(event); 411 | 412 | done(); 413 | } 414 | }; 415 | 416 | REGISTER_KERNEL_BUILDER(Name("Evict").Device(DEVICE_GPU).HostMemory("addrs").HostMemory("exclude_addrs"), EvictOp); 417 | -------------------------------------------------------------------------------- /pyproject.toml: -------------------------------------------------------------------------------- 1 | [tool.poetry] 2 | name = "tensorflow-huge-model-support" 3 | version = "1.1.2" 4 | description = "This library is designed to speed up huge model training on unified memory." 5 | authors = ["Chi-Chung Chen "] 6 | license = "CC BY-NC-SA 4.0" 7 | 8 | [tool.poetry.dependencies] 9 | python = "^3.7" 10 | tensorflow-gpu = "1.15.5" 11 | GPUtil = "^1.4.0" 12 | networkx = "2.2" 13 | pycrypto = "^2.6.1" 14 | psutil = "^5.8.0" 15 | Cython = "^0.29.24" 16 | 17 | [tool.poetry.dev-dependencies] 18 | 19 | [build-system] 20 | requires = ["poetry-core>=1.0.0"] 21 | build-backend = "poetry.core.masonry.api" 22 | -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | from setuptools import setup, Extension, find_packages 2 | from setuptools.command.install import install 3 | import subprocess 4 | import tensorflow as tf 5 | import os 6 | 7 | from Cython.Build import cythonize # cythonize must be imported after setuptools 8 | 9 | if "CUDA_PATH" in os.environ: 10 | CUDA_PATH = os.environ["CUDA_PATH"] 11 | else: 12 | CUDA_PATH = '/usr/local/cuda' 13 | 14 | install_requires = [ 15 | 'gputil', 16 | 'networkx==2.2', 17 | 'psutil', 18 | 'pycrypto', 19 | 'Cython' 20 | ] 21 | 22 | setup( 23 | name='tensorflow-huge-model-support', 24 | version='1.1.2', 25 | description='', 26 | author='Chi-Chung Chen', 27 | author_email='chenchc@aetherai.com', 28 | ext_modules=( 29 | cythonize('tensorflow_huge_model_support/*.py', compiler_directives={'language_level' : "3"}) + 30 | [Extension( 31 | 'tensorflow_huge_model_support.prefetch', 32 | ['prefetch_module/prefetch.cc'], 33 | include_dirs=[CUDA_PATH + '/include/'], 34 | library_dirs=[CUDA_PATH + '/lib64/'], 35 | libraries=['cuda'], 36 | extra_compile_args=['-fPIC', '-std=c++11', '-O3'] + tf.sysconfig.get_compile_flags(), 37 | extra_link_args=tf.sysconfig.get_link_flags(), 38 | language='c++', 39 | )] 40 | ), 41 | install_requires=install_requires, 42 | ) 43 | 44 | -------------------------------------------------------------------------------- /tensorflow_huge_model_support/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/aetherAI/tensorflow-huge-model-support/522d4d26706def6bc397ee27f49186573e5db6a0/tensorflow_huge_model_support/__init__.py -------------------------------------------------------------------------------- /tensorflow_huge_model_support/hms.py: -------------------------------------------------------------------------------- 1 | from tensorflow_huge_model_support.utils import USE_TF_2, get_gpu_mem_size 2 | 3 | import networkx as nx 4 | import tensorflow as tf 5 | import os 6 | import sys 7 | import pkgutil 8 | 9 | MERGE_TYPES = ['Merge'] 10 | IDENTITY_LIKE_TYPES = ['Identity', 'Reshape'] 11 | SIZE_STATIC_TYPES = ['Identity', 'Reshape', 'Transpose'] 12 | 13 | def isOpExcluded(op, strict, default_batch_size=None, thres=None): 14 | if 'HMS/' in op.name: 15 | return True 16 | if not strict: 17 | return False 18 | if any([arg.is_ref for arg in op.op_def.output_arg]): 19 | return True 20 | if any([tensor.dtype == tf.resource for tensor in op.outputs]): 21 | return True 22 | 23 | if thres != None: 24 | op_size = 0 25 | for tensor in op.outputs: 26 | tensor_size = _EstimateTensorSizeSingleton()(tensor, default_batch_size) 27 | tensor_size = 0 if tensor_size == None else tensor_size 28 | op_size += tensor_size 29 | if op_size < thres: 30 | return True 31 | 32 | return False 33 | 34 | class _EstimateTensorSizeSingleton(object): 35 | 36 | _instance = None 37 | 38 | def __new__(cls, *args, **kw): 39 | if not cls._instance: 40 | cls._instance = super(_EstimateTensorSizeSingleton, cls).__new__(cls, *args, **kw) 41 | cls._instance.cache = {} 42 | cls._instance.cache_nodtype = {} 43 | 44 | return cls._instance 45 | 46 | def __call__(self, tensor, default_batch_size=None): 47 | if tensor in self.cache: 48 | return self.cache[tensor] 49 | 50 | if tensor.op.type in IDENTITY_LIKE_TYPES: 51 | return 0 52 | 53 | if tensor.dtype == tf.float16 or tensor.dtype == tf.int16: 54 | size = dtype_size = 2 55 | elif tensor.dtype == tf.float32 or tensor.dtype == tf.int32: 56 | size = dtype_size = 4 57 | else: 58 | size = dtype_size = 8 59 | 60 | def get_tensor_size(tensor, default_batch_size=None): 61 | if tensor in self.cache_nodtype: 62 | return self.cache_nodtype[tensor] 63 | 64 | # If the tensor shape is normal 65 | try: 66 | tensor_size = 1 67 | for dim in tensor.shape: 68 | tensor_size *= int(dim) 69 | return tensor_size 70 | except: 71 | pass 72 | 73 | # If the tensor is size-static 74 | if tensor.op.type in SIZE_STATIC_TYPES: 75 | return get_tensor_size(tensor.op.inputs[0]) 76 | 77 | # If the tensor shape is only missing on batch dimension 78 | tensor_size = 1 79 | try: 80 | for idx, dim in enumerate(tensor.shape): 81 | if dim.value == None and idx == 0 and default_batch_size != None: 82 | dim = default_batch_size 83 | tensor_size *= int(dim) 84 | return tensor_size 85 | except: 86 | pass 87 | 88 | # Assume the tensor's size is equal to its largest input 89 | tensor_size = 0 90 | for input_tensor in tensor.op.inputs: 91 | input_tensor_size = get_tensor_size(input_tensor, default_batch_size=default_batch_size) 92 | tensor_size = max(tensor_size, input_tensor_size) if input_tensor_size != None else tensor_size 93 | if tensor_size != 0: 94 | return tensor_size 95 | 96 | # Just give up 97 | return None 98 | 99 | tensor_size = get_tensor_size(tensor, default_batch_size=default_batch_size) 100 | self.cache_nodtype[tensor] = tensor_size 101 | size = None if tensor_size == None else int(size * tensor_size) 102 | 103 | self.cache[tensor] = size 104 | 105 | workspace = 0 106 | if tensor.op.type in ['Conv2D', 'Conv2DBackpropFilter']: 107 | for input_tensor in tensor.op.inputs: 108 | input_tensor_size = get_tensor_size(input_tensor, default_batch_size=default_batch_size) 109 | workspace += (input_tensor_size * dtype_size) if input_tensor_size != None else 0 110 | elif tensor.op.type in ['Conv2DBackpropInput']: 111 | for input_tensor in tensor.op.inputs: 112 | input_tensor_size = get_tensor_size(input_tensor, default_batch_size=default_batch_size) 113 | workspace += (input_tensor_size * dtype_size) if input_tensor_size != None else 0 114 | workspace += size if size != None else 0 115 | size = None if size == None else size + workspace 116 | 117 | return size 118 | 119 | class _DependencyHelper(object): 120 | SCOPE_GROUP = 'group' # the entire group 121 | SCOPE_CRITICAL = 'critical' # the critical path in the group 122 | SCOPE_TARGET_OP = 'target_op' # e.g. train_step, loss 123 | SCOPE_OP = 'op' # A single operation 124 | 125 | def __init__(self, op_graph, target_op, critical_path): 126 | self._op_graph = op_graph 127 | self._target_op = target_op 128 | self._critical_path = critical_path 129 | 130 | def add_dep(self, src, src_scope, dst, dst_scope): 131 | # Check arguments 132 | def check_arguments(target, scope): 133 | if scope in [self.SCOPE_GROUP, self.SCOPE_CRITICAL]: 134 | assert(isinstance(target, _OperationGraph)) 135 | elif scope in [self.SCOPE_TARGET_OP]: 136 | assert(target == None) 137 | elif scope in [self.SCOPE_OP]: 138 | assert(isinstance(target, tf.Operation)) 139 | else: 140 | assert False 141 | 142 | check_arguments(src, src_scope) 143 | check_arguments(dst, dst_scope) 144 | 145 | # Build dependencies 146 | def get_op_list(target, scope, is_src): 147 | if scope == self.SCOPE_GROUP: 148 | if is_src: 149 | return target.get_exit_nodes() 150 | else: 151 | return target.get_entry_nodes() 152 | elif scope == self.SCOPE_CRITICAL: 153 | if is_src: 154 | return target.get_exit_nodes(on_path=self._critical_path) 155 | else: 156 | return target.get_entry_nodes(on_path=self._critical_path) 157 | elif scope == self.SCOPE_TARGET_OP: 158 | return [self._target_op] 159 | elif scope == self.SCOPE_OP: 160 | return [target] 161 | else: 162 | assert False 163 | 164 | aggregation = tf.group(get_op_list(src, src_scope, is_src=True)) 165 | for dst_op in get_op_list(dst, dst_scope, is_src=False): 166 | dst_op._add_control_inputs([aggregation]) 167 | 168 | class _OperationGraph(object): 169 | ''' 170 | Graph of ops and their computational dependencies. 171 | ''' 172 | def __init__( 173 | self, 174 | nx_graph, 175 | seen_ops, 176 | default_batch_size=None, 177 | thres=None 178 | ): 179 | ''' 180 | Create a op graph. 181 | 182 | Args: 183 | nx_graph: An nx.Graph() object. 184 | ''' 185 | self._nx_graph = nx_graph 186 | self._seen_ops = seen_ops 187 | self._default_batch_size = default_batch_size 188 | self._thres = thres 189 | 190 | def __str__(self): 191 | return str(list(nx.algorithms.dag.topological_sort(self._nx_graph))) 192 | 193 | def get_critical_path(self): 194 | ''' 195 | Get the longest path in the graph. This is used for grouping. 196 | ''' 197 | critical_path = nx.algorithms.dag.dag_longest_path(self._nx_graph) 198 | 199 | # Remove operations right behind identity operations 200 | new_critical_path = [] 201 | is_identity = False 202 | for op in reversed(critical_path): 203 | last_is_identity = is_identity 204 | if op.type in IDENTITY_LIKE_TYPES: 205 | is_identity = True 206 | else: 207 | is_identity = False 208 | if not last_is_identity: 209 | new_critical_path.append(op) 210 | new_critical_path = list(reversed(new_critical_path)) 211 | 212 | return new_critical_path 213 | 214 | def get_entry_nodes(self, on_path=None): 215 | nx_graph_wo_seen_ops = self._nx_graph.copy() 216 | nx_graph_wo_seen_ops.remove_nodes_from(self._seen_ops) 217 | if on_path != None: 218 | nx_graph_wo_seen_ops.remove_nodes_from(list( 219 | set(nx_graph_wo_seen_ops.nodes()) - 220 | set(on_path) 221 | )) 222 | 223 | exclude_nodes = set() 224 | for node in nx_graph_wo_seen_ops.nodes(): 225 | if isOpExcluded(node, strict=True, default_batch_size=self._default_batch_size, thres=self._thres): 226 | exclude_nodes.add(node) 227 | 228 | while True: 229 | entry_nodes = set() 230 | for node, in_degree in nx_graph_wo_seen_ops.in_degree: 231 | if in_degree == 0: 232 | entry_nodes.add(node) 233 | remove_nodes = list(entry_nodes & exclude_nodes) 234 | if remove_nodes == []: 235 | break 236 | for node in remove_nodes: 237 | nx_graph_wo_seen_ops.remove_node(node) 238 | 239 | return list(entry_nodes) 240 | 241 | def get_exit_nodes(self, on_path=None): 242 | nx_graph = self._nx_graph.copy() 243 | if on_path != None: 244 | nx_graph.remove_nodes_from(list( 245 | set(nx_graph.nodes()) - 246 | set(on_path) 247 | )) 248 | 249 | exclude_nodes = set() 250 | for node in nx_graph.nodes(): 251 | if isOpExcluded(node, strict=True, default_batch_size=self._default_batch_size, thres=self._thres): 252 | exclude_nodes.add(node) 253 | 254 | while True: 255 | exit_nodes = set() 256 | for node, out_degree in nx_graph.out_degree: 257 | if out_degree == 0: 258 | exit_nodes.add(node) 259 | remove_nodes = list(exit_nodes & exclude_nodes) 260 | if remove_nodes == []: 261 | break 262 | for node in remove_nodes: 263 | nx_graph.remove_node(node) 264 | 265 | return list(exit_nodes) 266 | 267 | def get_ops(self): 268 | return list(self._nx_graph.nodes) 269 | 270 | def get_unseen_ops(self): 271 | return list(set(self._nx_graph.nodes) - set(self._seen_ops)) 272 | 273 | def _check_size(self, tensor, thres): 274 | if thres == 0: 275 | return True 276 | size = _EstimateTensorSizeSingleton()(tensor, self._default_batch_size) 277 | if size == None: 278 | return False 279 | if size >= thres: 280 | return True 281 | else: 282 | return False 283 | 284 | def get_tensors(self, thres=0): 285 | tensor_list = [] 286 | for op in self.get_ops(): 287 | for tensor in op.outputs: 288 | if not self._check_size(tensor, thres): 289 | continue 290 | tensor_list.append(tensor) 291 | 292 | return tensor_list 293 | 294 | def get_seen_tensors(self, thres=0): 295 | tensor_list = [] 296 | for op in self._seen_ops: 297 | for tensor in op.outputs: 298 | if not self._check_size(tensor, thres): 299 | continue 300 | tensor_list.append(tensor) 301 | 302 | return tensor_list 303 | 304 | def get_unseen_tensors(self, thres=0): 305 | return list( 306 | set(self.get_tensors(thres)) - 307 | set(self.get_seen_tensors(thres)) 308 | ) 309 | 310 | def estimate_working_set_size(self): 311 | result = 0 312 | for op in self.get_unseen_ops(): 313 | if 'implied_size' in self._nx_graph.node[op]: 314 | result += self._nx_graph.node[op]['implied_size'] 315 | for tensor in self.get_tensors(): 316 | size = _EstimateTensorSizeSingleton()(tensor, self._default_batch_size) 317 | result += size if size != None else 0 318 | 319 | return result 320 | 321 | def warmup_tensor_size_estimator(self): 322 | for op in nx.algorithms.dag.topological_sort(self._nx_graph): 323 | for tensor in op.outputs: 324 | _EstimateTensorSizeSingleton()(tensor, self._default_batch_size) 325 | 326 | def get_serialized_unseen_ops(self): 327 | # Use DFS to do topological sort. Using DFS minimizes the number of new control dependency edges. 328 | serialized = [] 329 | new_control_deps = [] 330 | stack = self.get_entry_nodes() 331 | unseen_ops = self.get_unseen_ops() 332 | done_ops = list(self._seen_ops) 333 | if len(stack) != 0: 334 | node = stack.pop() 335 | serialized.append(node) 336 | done_ops.append(node) 337 | is_branch_end = True 338 | for next_node in self._nx_graph.successors(node): 339 | if next_node not in unseen_ops: 340 | continue 341 | if next_node in done_ops: 342 | continue 343 | if not all([pred in done_ops for pred in self._nx_graph.predecessors(next_node)]): 344 | continue 345 | stack.append(next_node) 346 | is_branch_end = False 347 | if is_branch_end and len(stack) > 0: 348 | new_control_deps.append((node, stack[-1])) 349 | 350 | return serialized, new_control_deps 351 | 352 | def has_dependency(self, src, dst): 353 | return (dst in self._nx_graph.successors(src) ) 354 | 355 | def __contains__(self, input): 356 | if isinstance(input, tf.Tensor): 357 | return input in self.get_tensors() 358 | elif isinstance(input, tf.Operation): 359 | return input in self.get_ops() 360 | else: 361 | raise 362 | 363 | def summary(self, verbose=False): 364 | def simplify(op_list): 365 | return [op.name for op in op_list] 366 | 367 | if verbose: 368 | string = 'Graph Summary: \n' 369 | string += 'New nodes: {}\n'.format(simplify(self.get_unseen_ops())) 370 | string += 'Seen nodes: {}\n'.format(simplify(set(self.get_ops()) - set(self.get_unseen_ops()))) 371 | string += 'Entry nodes: {}\n'.format(simplify(self.get_entry_nodes())) 372 | string += 'Exit nodes: {}\n'.format(simplify(self.get_exit_nodes())) 373 | else: 374 | string = 'Exit nodes: {}'.format(simplify(self.get_exit_nodes())) 375 | 376 | return string 377 | 378 | @staticmethod 379 | def get_graph_by_target(target_op, seen_ops=[], default_batch_size=None, thres=None): 380 | ''' 381 | Build a op graph associated to computing a target op. 382 | 383 | Args: 384 | target_op: A tf.Operation. 385 | ''' 386 | # Convert the Tensorflow graph to a Networkx graph, 387 | nx_graph = nx.DiGraph() 388 | seen_ops_used = set() 389 | 390 | pending_set = set([target_op]) 391 | visited_set = set([target_op]) 392 | while len(pending_set) != 0: 393 | op = pending_set.pop() 394 | if op not in seen_ops: 395 | for from_tensor in op.inputs: 396 | from_op = from_tensor.op 397 | if isOpExcluded(from_op, strict=False): 398 | continue 399 | nx_graph.add_edge(from_op, op, is_control_dep=False) 400 | if from_op not in visited_set: 401 | pending_set.add(from_op) 402 | visited_set.add(from_op) 403 | for from_op in op.control_inputs: 404 | if isOpExcluded(from_op, strict=False): 405 | continue 406 | if from_op in seen_ops: 407 | continue 408 | nx_graph.add_edge(from_op, op, is_control_dep=True) 409 | if from_op not in visited_set: 410 | pending_set.add(from_op) 411 | visited_set.add(from_op) 412 | else: 413 | seen_ops_used.add(op) 414 | 415 | return _OperationGraph(nx_graph, list(seen_ops_used), default_batch_size, thres) 416 | 417 | @staticmethod 418 | def trim_conditional_branches(op_graph): 419 | op_graph.warmup_tensor_size_estimator() 420 | 421 | nx_graph = op_graph._nx_graph.copy() 422 | seen_ops = set(op_graph._seen_ops) 423 | 424 | # Deal with conditional operations in the graph by bypassing the branching paths. 425 | dealt_merge_ops = set() 426 | while True: 427 | graph_modified = False 428 | for op in reversed(list(nx.algorithms.dag.topological_sort(nx_graph))): 429 | # Skip the op if it's non-Merge op, seen or dealt with. 430 | if op.type not in MERGE_TYPES: 431 | continue 432 | if op in seen_ops: 433 | continue 434 | if op in dealt_merge_ops: 435 | continue 436 | 437 | # Find nondeterministic_ops and last_deterministic_ops 438 | nondeterministic_ops = set() 439 | last_deterministic_ops = set() 440 | nondeterministic_ops_size = 0 441 | 442 | for switch in nx_graph.predecessors(op): 443 | pending_list = [switch] 444 | local_nondeterministic_ops = set() 445 | local_last_deterministic_ops = set() 446 | 447 | while len(pending_list) != 0: 448 | candidate_op = pending_list[0] 449 | pending_list = pending_list[1: ] 450 | if all([succ in local_nondeterministic_ops | set([op]) for succ in nx_graph.successors(candidate_op)]): 451 | local_nondeterministic_ops.add(candidate_op) 452 | pending_list += nx_graph.predecessors(candidate_op) 453 | else: 454 | last_deterministic_ops.add(candidate_op) 455 | 456 | local_nondeterministic_ops_size = 0 457 | for nd_op in local_nondeterministic_ops: 458 | if 'implied_size' in nx_graph[nd_op]: 459 | local_nondeterministic_ops_size += nx_graph[nd_op]['implied_size'] 460 | for tensor in nd_op.outputs: 461 | size = _EstimateTensorSizeSingleton()(tensor, op_graph._default_batch_size) 462 | local_nondeterministic_ops_size += size if size != None else 0 463 | 464 | nondeterministic_ops |= local_nondeterministic_ops 465 | last_deterministic_ops |= local_last_deterministic_ops 466 | nondeterministic_ops_size = max(nondeterministic_ops_size, local_nondeterministic_ops_size) 467 | 468 | # Remove nondeterministic_ops from the graph 469 | nx_graph.remove_nodes_from(nondeterministic_ops) 470 | seen_ops -= nondeterministic_ops 471 | 472 | # Link last_deterministic_ops with the merge op 473 | for last_deterministic_op in last_deterministic_ops: 474 | nx_graph.add_edge(last_deterministic_op, op, is_control_dep=False) # TODO(chenchc): is_control_dep=False OK? 475 | 476 | # Annotate the total size of non-deterministic ops on the merge op 477 | nx_graph.node[op]['implied_size'] = nondeterministic_ops_size 478 | 479 | dealt_merge_ops.add(op) 480 | graph_modified = True 481 | break 482 | 483 | if not graph_modified: 484 | break 485 | 486 | return _OperationGraph(nx_graph, list(seen_ops), op_graph._default_batch_size, op_graph._thres) 487 | 488 | @staticmethod 489 | def get_subgraph_by_target_extended(op_graph, target_op, critical_ops, seen_ops=[]): 490 | nx_graph = nx.DiGraph() 491 | seen_ops_used = set() 492 | seen_ops_set = set(seen_ops + op_graph._seen_ops) 493 | critical_ops_set = set(critical_ops) 494 | 495 | # Ancestor nodes of target_op 496 | pending_set = set([target_op]) 497 | while len(pending_set) != 0: 498 | op = pending_set.pop() 499 | if op not in seen_ops_set: 500 | for in_edge in op_graph._nx_graph.in_edges(op, data=True): 501 | from_op = in_edge[0] 502 | is_control_dep = in_edge[2]['is_control_dep'] 503 | if is_control_dep and from_op in seen_ops_set: 504 | continue 505 | nx_graph.add_edge(from_op, op, is_control_dep=is_control_dep) 506 | pending_set.add(from_op) 507 | else: 508 | seen_ops_used.add(op) 509 | 510 | # Augment more releated nodes - successor nodes of the subnet that is not on critical path 511 | pending_set = set() 512 | for node in nx_graph: 513 | if node in seen_ops_used: 514 | continue 515 | for succ in op_graph._nx_graph.successors(node): 516 | pending_set.add(succ) 517 | 518 | while len(pending_set) != 0: 519 | op = pending_set.pop() 520 | if op in nx_graph: 521 | continue 522 | if op in critical_ops_set: 523 | continue 524 | if any([ 525 | pred not in nx_graph and 526 | pred not in seen_ops_set 527 | for pred in op_graph._nx_graph.predecessors(op) 528 | ]): 529 | continue 530 | for in_edge in op_graph._nx_graph.in_edges(op, data=True): 531 | from_op = in_edge[0] 532 | is_control_dep = in_edge[2]['is_control_dep'] 533 | if is_control_dep: 534 | continue 535 | if from_op not in nx_graph: 536 | assert from_op in seen_ops_set 537 | seen_ops_used.add(from_op) 538 | nx_graph.add_edge(from_op, op, is_control_dep=is_control_dep) 539 | 540 | pending_set.update(set(op_graph._nx_graph.successors(op))) 541 | 542 | return _OperationGraph(nx_graph, list(seen_ops_used), op_graph._default_batch_size, op_graph._thres) 543 | 544 | class _AddrAgent(object): 545 | def __init__(self, pf_module): 546 | self._pf_module = pf_module 547 | 548 | self._tensor_list = [] 549 | self._addr_var_dict = {} 550 | self._addr_var_updated_dict = {} 551 | self._addr_update_dict = {} 552 | 553 | def __contains__(self, tensor): 554 | return tensor in self._tensor_list 555 | 556 | def add_tensor(self, tensor): 557 | name_prefix = 'HMS/AddrAgent/' + tensor.name.replace(':', '_') + '/' 558 | if USE_TF_2: 559 | assign = tf.compat.v1.assign 560 | else: 561 | assign = tf.assign 562 | 563 | addr_latest = self._pf_module.get_tensor_addr(tensor, name=(name_prefix + 'addr_latest')) 564 | with tf.device('/cpu:0'): 565 | addr_var = tf.Variable( 566 | initial_value=tf.zeros(shape=[2], dtype=tf.int64), 567 | name=(name_prefix + 'addr_var'), 568 | trainable=False 569 | ) 570 | addr_var_updated = assign(addr_var, addr_latest, use_locking=False, name=(name_prefix + 'addr_var_updated')) 571 | addr_update = addr_var_updated.op 572 | 573 | self._tensor_list.append(tensor) 574 | self._addr_var_dict[tensor] = addr_var 575 | self._addr_var_updated_dict[tensor] = addr_var_updated 576 | self._addr_update_dict[tensor] = addr_update 577 | 578 | def get_addr_var(self, tensor): 579 | return self._addr_var_dict[tensor] 580 | 581 | def get_addr_var_updated(self, tensor): 582 | return self._addr_var_updated_dict[tensor] 583 | 584 | def get_addr_update(self, tensor): 585 | return self._addr_update_dict[tensor] 586 | 587 | def __add__(self, rhs): 588 | def merge_two_dict(dict1, dict2): 589 | result = dict1.copy() 590 | result.update(dict2) 591 | return result 592 | 593 | result = _AddrAgent(self._pf_module) 594 | result._tensor_list = self._tensor_list + rhs._tensor_list 595 | result._addr_var_dict = merge_two_dict(self._addr_var_dict, rhs._addr_var_dict) 596 | result._addr_var_updated_dict = merge_two_dict(self._addr_var_updated_dict, rhs._addr_var_updated_dict) 597 | result._addr_update_dict = merge_two_dict(self._addr_update_dict, rhs._addr_update_dict) 598 | 599 | return result 600 | 601 | def get_tensors(self): 602 | return self._tensor_list 603 | 604 | class _PrefetchAgent(object): 605 | id_counter = 0 606 | 607 | def __init__(self, addr_agent, pf_module): 608 | self._addr_agent = addr_agent 609 | self._pf_module = pf_module 610 | self._addr_list = [] 611 | self._id = _PrefetchAgent.id_counter 612 | _PrefetchAgent.id_counter += 1 613 | 614 | def add_tensors(self, tensors): 615 | addr_list = [] 616 | for tensor in tensors: 617 | addr_list.append(self._addr_agent.get_addr_var(tensor)) 618 | 619 | self._addr_list += addr_list 620 | 621 | def add_seen_tensors(self, tensors): 622 | addr_list = [] 623 | for tensor in tensors: 624 | addr_list.append(self._addr_agent.get_addr_var_updated(tensor)) 625 | 626 | self._addr_list += addr_list 627 | 628 | def get_prefetch_op(self, expand_or_shrink): 629 | if len(self._addr_list) == 0: 630 | return tf.no_op() 631 | 632 | name_prefix = 'HMS/PrefetchAgent/' + str(self._id) + '/' 633 | with tf.device('/cpu:0'): 634 | addr_list_stacked = tf.stack(self._addr_list, name=(name_prefix + 'addr_list_stacked')) 635 | 636 | prefetch_op = self._pf_module.prefetch( 637 | addr_list_stacked, 638 | expand_or_shrink=expand_or_shrink, 639 | name=(name_prefix + 'prefetch_op') 640 | ) 641 | 642 | return prefetch_op 643 | 644 | class _EvictAgent(object): 645 | WARMUP_STEPS = 5 646 | id_counter = 0 647 | 648 | def __init__(self, addr_agent, pf_module): 649 | self._addr_agent = addr_agent 650 | self._pf_module = pf_module 651 | self._tensor_list = [] 652 | self._addr_list = [] 653 | self._exclude_tensor_list = [] 654 | self._exclude_addr_list = [] 655 | self._id = _EvictAgent.id_counter 656 | _EvictAgent.id_counter += 1 657 | 658 | def _exclude_tensors(self): 659 | for tensor in self._exclude_tensor_list: 660 | while tensor in self._tensor_list: 661 | idx = self._tensor_list.index(tensor) 662 | self._tensor_list.pop(idx) 663 | self._addr_list.pop(idx) 664 | 665 | def add_tensors(self, tensors, use_addr_var_updated=False): 666 | for tensor in tensors: 667 | self._tensor_list.append(tensor) 668 | if use_addr_var_updated: 669 | self._addr_list.append(self._addr_agent.get_addr_var_updated(tensor)) 670 | else: 671 | self._addr_list.append(self._addr_agent.get_addr_var(tensor)) 672 | 673 | self._exclude_tensors() 674 | 675 | def add_exclude_tensors(self, tensors, use_addr_var_updated=False): 676 | for tensor in tensors: 677 | self._exclude_tensor_list.append(tensor) 678 | if use_addr_var_updated: 679 | self._exclude_addr_list.append(self._addr_agent.get_addr_var_updated(tensor)) 680 | else: 681 | self._exclude_addr_list.append(self._addr_agent.get_addr_var(tensor)) 682 | 683 | self._exclude_tensors() 684 | 685 | def get_evict_op(self): 686 | if len(self._addr_list) == 0: 687 | return tf.no_op() 688 | 689 | if USE_TF_2: 690 | assign = tf.compat.v1.assign 691 | else: 692 | assign = tf.assign 693 | 694 | name_prefix = 'HMS/EvictAgent/' + str(self._id) + '/' 695 | with tf.device('/cpu:0'): 696 | warmup_counter = tf.Variable( 697 | initial_value=tf.zeros(shape=[], dtype=tf.int64), 698 | name=(name_prefix + 'warmup_counter'), 699 | trainable=False 700 | ) 701 | updated_warmup_counter = assign(warmup_counter, warmup_counter + 1) 702 | addr_list_stacked = tf.stack(self._addr_list, name=(name_prefix + 'addr_list_stacked')) 703 | if self._exclude_addr_list != []: 704 | exclude_addr_list_stacked = tf.stack(self._exclude_addr_list, name=(name_prefix + 'exclude_addr_list_stacked')) 705 | else: 706 | exclude_addr_list_stacked = tf.zeros(shape=[1, 2], dtype=tf.int64, name=(name_prefix + 'exclude_addr_list_stacked')) 707 | evict_op = tf.cond( 708 | updated_warmup_counter > _EvictAgent.WARMUP_STEPS, 709 | lambda: self._pf_module.evict(addr_list_stacked, exclude_addr_list_stacked, name=(name_prefix + 'evict_op')), 710 | lambda: tf.no_op() 711 | ).op 712 | 713 | return evict_op 714 | 715 | class HMS(object): 716 | ''' 717 | HMS class for Huge Model Support 718 | ''' 719 | def __init__( 720 | self, 721 | predict_op, 722 | train_step_op, 723 | default_batch_size=None, 724 | hms_pf_module_path=None, 725 | graph=None, 726 | fwd_pf_en=True, 727 | bwd_pf_en=True, 728 | fwd_pf_seen_only=False, 729 | bwd_pf_seen_only=False, # TODO: test if this is better to turn on or not. 730 | addr_fetch_depth=1, 731 | pf_thres_size=(1024 * 1024), 732 | pf_depth=2, 733 | fwd_group_exe_en=True, 734 | bwd_group_exe_en=True, 735 | group_exe_depth=1, 736 | fwd_serialize_en=False, 737 | bwd_serialize_en=False, 738 | fwd_evict_en=False, 739 | bwd_evict_en=False, 740 | evict_depth=1, 741 | group_maxsize=None, 742 | gpu_device='/gpu:0', 743 | hvd=None 744 | ): 745 | ''' 746 | Create and initialize an HMS object. 747 | 748 | Args: 749 | predict_op: A tf.Operation. The final op of the forward pass. 750 | train_step_op: A tf.Operation. The op used to update the model. 751 | hms_pf_module_path: A string. The path of the .so file for prefetch. 752 | graph: A tf.Graph. The graph where the model is placed. 753 | fwd_pf_en: A bool. Enabling prefetch in forward pass. 754 | bwd_pf_en: A bool. Enabling prefetch in backward pass. 755 | pf_depth: An int. The number of groups prefetched in advance. 756 | fwd_group_exe_en: A bool. Enabling group execution on forward pass. 757 | bwd_group_exe_en: A bool. Enabling group execution on backward pass. 758 | group_maxsize: An int. The maximal byte count of a group. Set None to let HMS determine by system limit. 759 | Set 0 to disable grouping. 760 | gpu_device: A string. The GPU device name. 761 | ''' 762 | # Validate inputs and set default values 763 | if graph == None: 764 | graph = predict_op.graph 765 | if group_maxsize == None: 766 | group_maxsize = int(get_gpu_mem_size(hvd=hvd) * 0.25) 767 | if hms_pf_module_path == None: 768 | prefetch_package = pkgutil.get_loader('tensorflow_huge_model_support.prefetch') 769 | hms_pf_module_path = prefetch_package.path 770 | 771 | # Declare data members from function inputs 772 | self._predict_op = predict_op 773 | self._train_step_op = train_step_op 774 | self._pf_module = tf.load_op_library(hms_pf_module_path) 775 | self._default_batch_size = default_batch_size 776 | self._graph = graph 777 | self._fwd_pf_en = fwd_pf_en 778 | self._bwd_pf_en = bwd_pf_en 779 | self._fwd_pf_seen_only = fwd_pf_seen_only 780 | self._bwd_pf_seen_only = bwd_pf_seen_only 781 | self._addr_fetch_depth = addr_fetch_depth 782 | self._pf_thres_size = pf_thres_size 783 | self._pf_depth = pf_depth 784 | self._fwd_group_exe_en = fwd_group_exe_en 785 | self._bwd_group_exe_en = bwd_group_exe_en 786 | self._group_exe_depth = group_exe_depth 787 | self._fwd_serialize_en = fwd_serialize_en 788 | self._bwd_serialize_en = bwd_serialize_en 789 | self._fwd_evict_en = fwd_evict_en 790 | self._bwd_evict_en = bwd_evict_en 791 | self._evict_depth = evict_depth 792 | self._group_maxsize = group_maxsize 793 | self._gpu_device = gpu_device 794 | self._hvd = hvd 795 | 796 | self._fwd_addr_agent = None 797 | self._bwd_addr_agent = None 798 | 799 | # Initialization 800 | 801 | def _grouping(self, op_graph, critical_path, group_maxsize): 802 | def log(idx, new_group): 803 | self._log_info( 804 | 'New group {} with {} nodes.'.format(idx, len(new_group.get_ops())) 805 | ) 806 | self._log_info(new_group.summary()) 807 | 808 | seen_ops = set() 809 | group_list = [] 810 | pending_sub_graph = None 811 | i = 0 812 | while i != len(critical_path): 813 | op = critical_path[i] 814 | sub_graph = _OperationGraph.get_subgraph_by_target_extended(op_graph, op, critical_path, list(seen_ops)) 815 | working_set_size = sub_graph.estimate_working_set_size() 816 | if working_set_size < group_maxsize: 817 | pending_sub_graph = sub_graph 818 | i += 1 819 | elif pending_sub_graph == None: 820 | group_list.append(sub_graph) 821 | log(len(group_list) - 1, sub_graph) 822 | seen_ops.update(sub_graph.get_ops()) 823 | i += 1 824 | else: 825 | group_list.append(pending_sub_graph) 826 | log(len(group_list) - 1, pending_sub_graph) 827 | seen_ops.update(pending_sub_graph.get_ops()) 828 | pending_sub_graph = None 829 | 830 | if pending_sub_graph != None: 831 | group_list.append(pending_sub_graph) 832 | log(len(group_list) - 1, pending_sub_graph) 833 | seen_ops.update(pending_sub_graph.get_ops()) 834 | pending_sub_graph = None 835 | 836 | assert seen_ops == set(op_graph.get_ops()) 837 | 838 | return group_list 839 | 840 | def _do_serialize(self, group_list): 841 | for i, group in enumerate(group_list): 842 | serialized_ops, new_control_deps = group.get_serialized_unseen_ops() 843 | for dep in new_control_deps: 844 | dep[1]._add_control_inputs([dep[0]]) 845 | 846 | def _do_group_execution(self, group_list, dep_helper, prev_group_list=None): 847 | for i, group in enumerate(group_list): 848 | if i >= self._group_exe_depth + 1: 849 | dep_helper.add_dep( 850 | src=(group_list[i - self._group_exe_depth - 1]), 851 | src_scope=dep_helper.SCOPE_GROUP, 852 | dst=group, 853 | dst_scope=dep_helper.SCOPE_GROUP 854 | ) 855 | elif prev_group_list != None: 856 | dep_helper.add_dep( 857 | src=(prev_group_list[len(prev_group_list) + i - self._group_exe_depth - 1]), 858 | src_scope=dep_helper.SCOPE_GROUP, 859 | dst=group, 860 | dst_scope=dep_helper.SCOPE_GROUP 861 | ) 862 | 863 | #self._log_info('Group execution built on group {}.'.format(i)) 864 | 865 | def _do_addr_fetch(self, target_op, group_list, addr_agent, critical_path, dep_helper, prev_addr_agent=None): 866 | # Build addr fetching operations and store in addr agent. 867 | for i, group in enumerate(group_list): 868 | aggregate_op = tf.no_op() 869 | for tensor in group.get_tensors(thres=self._pf_thres_size): 870 | if prev_addr_agent != None and tensor in prev_addr_agent: 871 | continue 872 | if tensor in addr_agent: 873 | continue 874 | if tensor.op == target_op: 875 | continue 876 | 877 | addr_agent.add_tensor(tensor) 878 | 879 | dep_helper.add_dep( 880 | src=(addr_agent.get_addr_update(tensor)), 881 | src_scope=dep_helper.SCOPE_OP, 882 | dst=aggregate_op, 883 | dst_scope=dep_helper.SCOPE_OP 884 | ) 885 | 886 | if i + self._addr_fetch_depth + 1 < len(group_list): 887 | dep_helper.add_dep( 888 | src=aggregate_op, 889 | src_scope=dep_helper.SCOPE_OP, 890 | dst=(group_list[i + self._addr_fetch_depth + 1]), 891 | dst_scope=dep_helper.SCOPE_GROUP 892 | ) 893 | else: 894 | dep_helper.add_dep( 895 | src=aggregate_op, 896 | src_scope=dep_helper.SCOPE_OP, 897 | dst=None, 898 | dst_scope=dep_helper.SCOPE_TARGET_OP 899 | ) 900 | 901 | #self._log_info('Addr fetch built on group {}.'.format(i)) 902 | 903 | def _do_prefetch(self, target_op, group_list, addr_agent, dep_helper, circular_pf_op_list, prev_addr_agent=None, prev_group_list=None, seen_only=False): 904 | # Build prefetch operations 905 | if prev_addr_agent == None: 906 | merged_addr_agent = addr_agent 907 | else: 908 | merged_addr_agent = addr_agent + prev_addr_agent 909 | 910 | prev_prefetch_op = None 911 | for i, group in enumerate(group_list): 912 | pf_agent = _PrefetchAgent(merged_addr_agent, self._pf_module) 913 | if not seen_only: 914 | pf_agent.add_tensors(list( 915 | set(group.get_unseen_tensors(thres=self._pf_thres_size)) - 916 | set(target_op.outputs) 917 | )) 918 | for tensor in list( 919 | set(group.get_seen_tensors(thres=self._pf_thres_size)) - 920 | set(target_op.outputs) 921 | ): 922 | recent_access = False 923 | for j in range(1, self._addr_fetch_depth + 1): 924 | if i - j >= 0: 925 | recent_access = recent_access or (tensor in group_list[i - j].get_unseen_tensors()) 926 | elif prev_group_list != None: 927 | recent_access = recent_access or (tensor in prev_group_list[len(prev_group_list) + i - j].get_unseen_tensors()) 928 | if recent_access: 929 | #pf_agent.add_tensors([tensor]) # TODO: Check if this is necessary. 930 | pass 931 | else: 932 | pf_agent.add_seen_tensors([tensor]) 933 | 934 | prefetch_op = pf_agent.get_prefetch_op(True) 935 | 936 | if i >= self._pf_depth + 1: 937 | dep_helper.add_dep( 938 | src=(group_list[i - self._pf_depth - 1]), 939 | src_scope=dep_helper.SCOPE_GROUP, 940 | dst=prefetch_op, 941 | dst_scope=dep_helper.SCOPE_OP 942 | ) 943 | dep_helper.add_dep( 944 | src=prefetch_op, 945 | src_scope=dep_helper.SCOPE_OP, 946 | dst=group, 947 | dst_scope=dep_helper.SCOPE_GROUP 948 | ) 949 | elif prev_group_list != None: 950 | dep_helper.add_dep( 951 | src=(prev_group_list[len(prev_group_list) + i - self._pf_depth - 1]), 952 | src_scope=dep_helper.SCOPE_GROUP, 953 | dst=prefetch_op, 954 | dst_scope=dep_helper.SCOPE_OP 955 | ) 956 | dep_helper.add_dep( 957 | src=prefetch_op, 958 | src_scope=dep_helper.SCOPE_OP, 959 | dst=group, 960 | dst_scope=dep_helper.SCOPE_GROUP 961 | ) 962 | else: 963 | circular_pf_op_list.append(prefetch_op) 964 | 965 | prev_prefetch_op = prefetch_op 966 | #self._log_info('Prefetch built on group {}.'.format(i)) 967 | 968 | if prev_group_list != None: 969 | assert len(circular_pf_op_list) == self._pf_depth + 1 970 | for i in range(self._pf_depth + 1): 971 | dep_helper.add_dep( 972 | src=(group_list[len(group_list) - i - self._pf_depth - 1]), 973 | src_scope=dep_helper.SCOPE_GROUP, 974 | dst=circular_pf_op_list[i], 975 | dst_scope=dep_helper.SCOPE_OP 976 | ) 977 | dep_helper.add_dep( 978 | src=circular_pf_op_list[i], 979 | src_scope=dep_helper.SCOPE_OP, 980 | dst=None, 981 | dst_scope=dep_helper.SCOPE_TARGET_OP 982 | ) 983 | 984 | def _modify_graph( 985 | self, 986 | target_op, 987 | addr_agent, 988 | circular_pf_op_list, 989 | seen_ops=[], 990 | prefetch_en=True, 991 | evict_en=False, 992 | prefetch_seen_only=False, 993 | group_exe_en=True, 994 | serialize_en=False, 995 | prev_addr_agent=None, 996 | prev_group_list=None 997 | ): 998 | ''' 999 | Modify graph on either forward pass or backward pass. 1000 | 1001 | Args: 1002 | target_op: A Operation. predict_op for forward pass and train_step_op for backward pass. 1003 | seen_ops: A list of Operations. Operations that have been already computed. 1004 | prefetch_en: A bool. Enable prefetch or not. 1005 | group_exe_en: A bool. Enable group execution or not. 1006 | ''' 1007 | op_graph = _OperationGraph.get_graph_by_target(target_op, seen_ops=seen_ops, default_batch_size=self._default_batch_size, thres=self._pf_thres_size) 1008 | op_graph = _OperationGraph.trim_conditional_branches(op_graph) 1009 | critical_path = op_graph.get_critical_path() 1010 | dep_helper = _DependencyHelper(op_graph, target_op, critical_path) 1011 | group_list = self._grouping(op_graph, critical_path, self._group_maxsize) 1012 | if serialize_en: 1013 | self._log_info('Serializing.') 1014 | self._do_serialize(group_list) 1015 | elif group_exe_en: 1016 | self._log_info('Building group execution.') 1017 | self._do_group_execution(group_list, dep_helper, prev_group_list=prev_group_list) 1018 | if prefetch_en or evict_en: 1019 | self._log_info('Building address fetch.') 1020 | self._do_addr_fetch(target_op, group_list, addr_agent, critical_path, dep_helper, prev_addr_agent=prev_addr_agent) 1021 | if prefetch_en: 1022 | self._log_info('Building prefetching.') 1023 | self._do_prefetch( 1024 | target_op, 1025 | group_list, 1026 | addr_agent, 1027 | dep_helper, 1028 | circular_pf_op_list, 1029 | prev_addr_agent=prev_addr_agent, 1030 | prev_group_list=prev_group_list, 1031 | seen_only=prefetch_seen_only 1032 | ) 1033 | 1034 | return group_list, dep_helper 1035 | 1036 | def _do_evict(self, target_op, group_list, dep_helper, merged_addr_agent, pf_en, next_group_list=None): 1037 | # Build evict operations 1038 | for i, group in enumerate(group_list): 1039 | if next_group_list == None and i + self._evict_depth + 1 >= len(group_list): 1040 | # One of the very last groups. Skip eviction. 1041 | # TODO: Is circular eviction helpful? 1042 | continue 1043 | 1044 | evict_agent = _EvictAgent(merged_addr_agent, self._pf_module) 1045 | evict_agent.add_tensors( 1046 | group.get_tensors(thres=self._pf_thres_size), 1047 | use_addr_var_updated=(self._evict_depth > self._addr_fetch_depth) 1048 | ) 1049 | 1050 | exclude_depth = self._evict_depth 1051 | if pf_en: 1052 | exclude_depth = max(exclude_depth, self._pf_depth + 1) 1053 | for dist in range(1, exclude_depth + 1): 1054 | if i + dist < len(group_list): 1055 | evict_agent.add_exclude_tensors( 1056 | group_list[i + dist].get_tensors(thres=self._pf_thres_size), 1057 | use_addr_var_updated=(self._evict_depth - dist > self._addr_fetch_depth) 1058 | ) 1059 | elif next_group_list != None and i + dist - len(group_list) < len(next_group_list): 1060 | evict_agent.add_exclude_tensors( 1061 | next_group_list[i + dist - len(group_list)].get_tensors(thres=self._pf_thres_size), 1062 | use_addr_var_updated=(self._evict_depth - dist > self._addr_fetch_depth) 1063 | ) 1064 | 1065 | evict_op = evict_agent.get_evict_op() 1066 | 1067 | if i + self._evict_depth + 1 < len(group_list): 1068 | dep_helper.add_dep( 1069 | src=group, 1070 | src_scope=dep_helper.SCOPE_GROUP, 1071 | dst=evict_op, 1072 | dst_scope=dep_helper.SCOPE_OP 1073 | ) 1074 | dep_helper.add_dep( 1075 | src=evict_op, 1076 | src_scope=dep_helper.SCOPE_OP, 1077 | dst=(group_list[i + self._evict_depth + 1]), 1078 | dst_scope=dep_helper.SCOPE_GROUP 1079 | ) 1080 | elif next_group_list != None: 1081 | dep_helper.add_dep( 1082 | src=group, 1083 | src_scope=dep_helper.SCOPE_GROUP, 1084 | dst=evict_op, 1085 | dst_scope=dep_helper.SCOPE_OP 1086 | ) 1087 | dep_helper.add_dep( 1088 | src=evict_op, 1089 | src_scope=dep_helper.SCOPE_OP, 1090 | dst=(next_group_list[i + self._evict_depth + 1 - len(group_list)]), 1091 | dst_scope=dep_helper.SCOPE_GROUP 1092 | ) 1093 | else: 1094 | raise 1095 | 1096 | def run(self): 1097 | ''' 1098 | Start modifying the graph to invoke HMS. 1099 | ''' 1100 | self._log_info('Start modifying the graph.') 1101 | 1102 | # Modify graph on forward pass 1103 | circular_pf_op_list = list() 1104 | fwd_ops = None 1105 | if self._fwd_group_exe_en or self._fwd_pf_en or self._fwd_evict_en: 1106 | self._log_info('Processing on forward pass.') 1107 | self._fwd_addr_agent = _AddrAgent(self._pf_module) 1108 | fwd_group_list, fwd_dep_helper = self._modify_graph( 1109 | self._predict_op, 1110 | self._fwd_addr_agent, 1111 | circular_pf_op_list, 1112 | seen_ops=[], 1113 | prefetch_en=self._fwd_pf_en, 1114 | evict_en=self._fwd_evict_en, 1115 | prefetch_seen_only=self._fwd_pf_seen_only, 1116 | group_exe_en=self._fwd_group_exe_en, 1117 | serialize_en=self._fwd_serialize_en 1118 | ) 1119 | 1120 | # Modify graph on backward pass 1121 | if self._bwd_group_exe_en or self._bwd_pf_en or self._bwd_evict_en: 1122 | self._log_info('Processing on backward pass.') 1123 | seen_ops = _OperationGraph.get_graph_by_target(self._predict_op, default_batch_size=self._default_batch_size).get_ops() 1124 | self._bwd_addr_agent = _AddrAgent(self._pf_module) 1125 | bwd_group_list, bwd_dep_helper = self._modify_graph( 1126 | self._train_step_op, 1127 | self._bwd_addr_agent, 1128 | circular_pf_op_list, 1129 | seen_ops=seen_ops, 1130 | prefetch_en=self._bwd_pf_en, 1131 | evict_en=self._bwd_evict_en, 1132 | prefetch_seen_only=self._bwd_pf_seen_only, 1133 | group_exe_en=self._bwd_group_exe_en, 1134 | serialize_en=self._bwd_serialize_en, 1135 | prev_addr_agent=self._fwd_addr_agent, 1136 | prev_group_list=fwd_group_list 1137 | ) 1138 | 1139 | # Modify graph on forward pass for eviction 1140 | if self._fwd_evict_en: 1141 | self._log_info('Building eviction on forward pass.') 1142 | self._do_evict( 1143 | self._predict_op, 1144 | fwd_group_list, 1145 | fwd_dep_helper, 1146 | self._fwd_addr_agent + self._bwd_addr_agent, 1147 | self._fwd_pf_en, 1148 | next_group_list=bwd_group_list 1149 | ) 1150 | 1151 | # Modify graph on backward pass for eviction 1152 | if self._bwd_evict_en: 1153 | self._log_info('Building eviction on backward pass.') 1154 | self._do_evict( 1155 | self._train_step_op, 1156 | bwd_group_list, 1157 | bwd_dep_helper, 1158 | self._fwd_addr_agent + self._bwd_addr_agent, 1159 | self._bwd_pf_en, 1160 | ) 1161 | 1162 | def _log_info(self, message): 1163 | ''' 1164 | Log debug message. 1165 | ''' 1166 | if USE_TF_2: 1167 | tf.get_logger().warn( 1168 | '[HMS] {}'.format(message) 1169 | ) 1170 | else: 1171 | tf.logging.log( 1172 | tf.logging.WARN, 1173 | '[HMS] {}'.format(message) 1174 | ) 1175 | 1176 | 1177 | -------------------------------------------------------------------------------- /tensorflow_huge_model_support/tf_keras.py: -------------------------------------------------------------------------------- 1 | import os 2 | from tensorflow_huge_model_support.hms import HMS 3 | from tensorflow_huge_model_support.utils import USE_TF_2, get_gpu_mem_size, get_host_mem_size, create_session 4 | 5 | import tensorflow as tf 6 | 7 | def init(base_config=None, hvd=None, copy=False): 8 | # Let Tensorflow use Unified Memory. 9 | host_mem_size = get_host_mem_size() 10 | gpu_mem_size = get_gpu_mem_size(hvd=hvd) 11 | um_ratio = host_mem_size / gpu_mem_size 12 | sess = create_session( 13 | base_config=base_config, 14 | um_ratio=um_ratio, 15 | visible_devices=(str(hvd.local_rank()) if hvd != None else None) 16 | ) 17 | 18 | if copy: 19 | old_sess = tf.keras.backend.get_session() 20 | var_list = tf.global_variables() 21 | if var_list != []: 22 | var_is_inited_list = [tf.is_variable_initialized(var) for var in var_list] 23 | var_is_inited_val_list = old_sess.run(var_is_inited_list) 24 | var_inited_var_list = [var for i, var in enumerate(var_list) if var_is_inited_val_list[i]] 25 | if var_inited_var_list != []: 26 | var_inited_var_val_list = old_sess.run(var_inited_var_list) 27 | sess.run([tf.assign(var_inited_var_list[i], var_inited_var_val_list[i]) for i in range(len(var_inited_var_list))]) 28 | 29 | tf.keras.backend.set_session(sess) 30 | 31 | return sess 32 | 33 | class HMSTFKerasCallback(tf.keras.callbacks.Callback): 34 | def __init__(self, hvd=None, **kwargs): 35 | self._hvd = hvd 36 | self._kwargs = kwargs 37 | self._hms = None 38 | 39 | def set_model(self, model): 40 | if self._hms != None: 41 | return 42 | #train_function = model._fit_function if hasattr(model, '_fit_function') else model.train_function 43 | model._make_train_function() 44 | train_function = model.train_function 45 | assert train_function != None 46 | 47 | # Invoke HMS 48 | fwd_op = model.total_loss.op 49 | bwd_op = train_function.updates_op 50 | self._hms = hms = HMS(fwd_op, bwd_op, hvd=self._hvd, **self._kwargs) 51 | hms.run() 52 | 53 | # Inialize HMS variables 54 | old_sess = tf.keras.backend.get_session() 55 | old_config = old_sess._config 56 | 57 | sess = init(base_config=old_config, hvd=self._hvd, copy=True) 58 | hms_vars = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope='HMS/') 59 | sess.run(tf.initializers.variables(hms_vars)) 60 | -------------------------------------------------------------------------------- /tensorflow_huge_model_support/tf_sess.py: -------------------------------------------------------------------------------- 1 | import os 2 | from tensorflow_huge_model_support.utils import USE_TF_2, get_gpu_mem_size, get_host_mem_size, create_session 3 | from tensorflow_huge_model_support.hms import HMS 4 | 5 | def HMSSession(config=None, hvd=None, hooks=None): 6 | host_mem_size = get_host_mem_size() 7 | gpu_mem_size = get_gpu_mem_size(hvd=hvd) 8 | um_ratio = host_mem_size / gpu_mem_size 9 | sess = create_session( 10 | base_config=config, 11 | um_ratio=um_ratio, 12 | visible_devices=(str(hvd.local_rank()) if hvd != None else None), 13 | hooks=hooks 14 | ) 15 | 16 | # Increase the limit of max workspace size of cuDNN covolutional layers. 17 | # This setting prevents cuDNN from using slow but memory-conservative algorithms. 18 | if 'TF_CUDNN_WORKSPACE_LIMIT_IN_MB' in os.environ: 19 | print('[HMS] Overriding TF_CUDNN_WORKSPACE_LIMIT_IN_MB...') 20 | os.environ['TF_CUDNN_WORKSPACE_LIMIT_IN_MB'] = str(host_mem_size) 21 | 22 | return sess 23 | -------------------------------------------------------------------------------- /tensorflow_huge_model_support/utils.py: -------------------------------------------------------------------------------- 1 | import copy 2 | import GPUtil 3 | import os 4 | import psutil 5 | import tensorflow as tf 6 | 7 | USE_TF_2 = tf.__version__.startswith('2') 8 | 9 | def get_gpu_mem_size(gpu_id=0, hvd=None): 10 | if hvd: 11 | gpu_id = hvd.local_rank() 12 | 13 | if 'CUDA_VISIBLE_DEVICES' in os.environ: 14 | gpu_list = os.environ['CUDA_VISIBLE_DEVICES'].split(',') 15 | gpu_list = [int(gpu) for gpu in gpu_list] 16 | gpu_mem_size = GPUtil.getGPUs()[gpu_list[gpu_id]].memoryTotal 17 | else: 18 | gpu_mem_size = GPUtil.getGPUs()[gpu_id].memoryTotal 19 | 20 | gpu_mem_size *= 1024 * 1024 21 | 22 | return gpu_mem_size 23 | 24 | def get_host_mem_size(): 25 | return psutil.virtual_memory().total 26 | 27 | def create_session(base_config=None, um_ratio=None, visible_devices=None, hooks=None): 28 | if USE_TF_2: 29 | ConfigProto = tf.compat.v1.ConfigProto 30 | Session = tf.compat.v1.Session 31 | MonitoredTrainingSession = tf.compat.v1.train.MonitoredTrainingSession 32 | else: 33 | ConfigProto = tf.ConfigProto 34 | Session = tf.Session 35 | MonitoredTrainingSession = tf.train.MonitoredTrainingSession 36 | 37 | if base_config: 38 | config = copy.deepcopy(base_config) 39 | else: 40 | config = ConfigProto() 41 | 42 | if um_ratio: 43 | config.gpu_options.per_process_gpu_memory_fraction = um_ratio 44 | if visible_devices: 45 | config.gpu_options.visible_device_list = visible_devices 46 | 47 | config.gpu_options.allow_growth = True 48 | 49 | if hooks: 50 | sess = MonitoredTrainingSession(config=config, hooks=hooks) 51 | else: 52 | sess = Session(config=config) 53 | 54 | return sess 55 | --------------------------------------------------------------------------------