├── LICENSE.txt ├── MODEL_LICENSE.txt ├── README.md ├── environment.yml ├── figs ├── model_framework.png └── results.png ├── guided_diffusion ├── __init__.py ├── __pycache__ │ ├── __init__.cpython-37.pyc │ ├── __init__.cpython-38.pyc │ ├── __init__.cpython-39.pyc │ ├── dist_util.cpython-37.pyc │ ├── dist_util.cpython-38.pyc │ ├── dist_util.cpython-39.pyc │ ├── fp16_util.cpython-37.pyc │ ├── fp16_util.cpython-39.pyc │ ├── gaussian_diffusion.cpython-37.pyc │ ├── gaussian_diffusion.cpython-39.pyc │ ├── image_datasets.cpython-37.pyc │ ├── image_datasets.cpython-39.pyc │ ├── logger.cpython-37.pyc │ ├── logger.cpython-39.pyc │ ├── losses.cpython-37.pyc │ ├── losses.cpython-39.pyc │ ├── nn.cpython-37.pyc │ ├── nn.cpython-39.pyc │ ├── resample.cpython-37.pyc │ ├── resample.cpython-39.pyc │ ├── respace.cpython-37.pyc │ ├── respace.cpython-39.pyc │ ├── script_util.cpython-37.pyc │ ├── script_util.cpython-39.pyc │ ├── train_util.cpython-37.pyc │ ├── train_util.cpython-39.pyc │ ├── unet.cpython-37.pyc │ └── unet.cpython-39.pyc ├── dist_util.py ├── fp16_util.py ├── gaussian_diffusion.py ├── image_datasets.py ├── logger.py ├── losses.py ├── nn.py ├── resample.py ├── respace.py ├── script_util.py ├── train_util.py └── unet.py ├── scripts ├── test.py └── train.py └── test_DDPM_3d_mpi.sh /LICENSE.txt: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2025 Boxiao Yu and Kuang Gong 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /MODEL_LICENSE.txt: -------------------------------------------------------------------------------- 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 |

Robust Whole-body PET Image Denoising Using 3D Diffusion Models: Evaluation Across Various Scanners, Tracers, and Dose Levels

2 |
3 | 4 | Boxiao Yu1, 5 | 6 | Savas Ozdemir2, 7 | 8 | Yafei Dong3, 9 | 10 | Wei Shao4, 11 | 12 | Tinsu Pan5, 13 | 14 | Kuangyu Shi6, 15 | 16 | Kuang Gong1 17 |
18 |
19 | 1J. Crayton Pruitt Family Department of Biomedical Engineering, 20 | University of Florida; 21 | 2Department of Raiology, University of Florida; 22 | 3Yale School of Medicine, Yale University; 23 | 4Department of Medicine, University of Florida; 24 | 5Department of Imaging Physics, University of Texas MD Anderson Cancer Center; 25 | 6Department of Nuclear Medicine, University of Bern 26 |
27 | 28 | ## Purpose 29 | 30 | Whole-body PET imaging plays an essential role in cancer diagnosis and treatment but suffers from low image quality. Traditional deep learning-based denoising methods work well for a specific acquisition but are less effective in handling diverse PET protocols. In this study, we proposed and validated a 3D Denoising Diffusion Probabilistic Model (3D DDPM) as a robust and universal solution for whole-body PET image denoising. 31 | 32 | ## Method 33 | 34 |

35 | 36 |

37 | The proposed 3D DDPM gradually injected noise into the images during the forward diffusion phase, allowing the model to learn to reconstruct the clean data during the reverse diffusion process. A 3D convolutional network was trained using high-quality data from the Biograph Vision Quadra PET/CT scanner to generate the score function, enabling the model to capture accurate PET distribution information extracted from the total-body datasets. The trained 3D DDPM was evaluated on datasets from four scanners, four tracer types, and six dose levels representing a broad spectrum of clinical scenarios. 38 | 39 | ## Result 40 | 41 |

42 | 43 |

44 | The proposed 3D DDPM consistently outperformed 2D DDPM, 3D UNet, and 3D GAN, demonstrating its superior denoising performance across all tested conditions. Additionally, the model’s uncertainty maps exhibited lower variance, reflecting its higher confidence in its outputs. 45 | 46 | ## Installation 47 | ### Step 1: Clone the Repository 48 | 49 | git clone https://github.com/Miche11eU/PET-Image-Denoising-Using-3D-Diffusion-Model.git 50 | cd PET-Image-Denoising-Using-3D-Diffusion-Model 51 | 52 | ### Step 2: Create and activate the conda environment from the environment.yml file: 53 | 54 | conda env create -f environment.yml 55 | conda activate PET-3D-DDPM 56 | 57 | ### Step 3: Download Pre-trained Models 58 | Download the pre-trained model files from this [link](https://www.dropbox.com/scl/fo/nj52fz7p23icnkxo3v5y2/AAsggV-0DAuJjd4ILYAE1m4?rlkey=uivlrx0oi68l7n34fkbmamkdj&st=fztnoohh&dl=0) and place them into the `./checkpoint/` folder. 59 | 60 | **Note**: This model is licensed under CC BY-NC-SA 4.0. **Commercial use is prohibited.** 61 | 62 | 63 | ## Testing 64 | 65 | ### Data Preparation 66 | 67 | Before running the denoising script, modify the `load_data_for_worker` function in `./scripts/test.py` to align with your data format and dimensions. This function is responsible for loading your low-dose PET data into the model. 68 | 69 | ### Running the Denoising Script 70 | 71 | We provide a shell script `test_DDPM_3d_mpi.sh` to facilitate the testing process. 72 | 73 | #### Usage 74 | 75 | - `--base_samples`: Path to the `.npz` files containing your low-dose PET images. 76 | - `--num_samples`: Total number of samples you wish to process. 77 | - `-n`: Number of GPUs to utilize for parallel processing. 78 | - `--save_dir`: Path to the directory where you want to save the denoised images. 79 | 80 | ## License 81 | 82 | - The **code** in this repository is licensed under the **MIT License**. 83 | - The **model weights** are licensed under **CC BY-NC-SA 4.0**, meaning: 84 | - You **can** share and modify the model weights, but **must** use the same license. 85 | - You **cannot** use it for **commercial purposes**. 86 | 87 | For details, check: 88 | - [LICENSE](./LICENSE.txt) (MIT for code) 89 | - [MODEL_LICENSE](./MODEL_LICENSE.txt) (CC BY-NC-SA 4.0 for model weights) 90 | 91 | ## Citation 92 | If you find our work is useful in your research or applications, please consider giving us a star 🌟 and citing it by the following BibTeX entry. 93 | 94 | ```bibtex 95 | @article{yu2025robust, 96 | title={Robust whole-body PET image denoising using 3D diffusion models: evaluation across various scanners, tracers, and dose levels}, 97 | author={Yu, Boxiao and Ozdemir, Savas and Dong, Yafei and Shao, Wei and Pan, Tinsu and Shi, Kuangyu and Gong, Kuang}, 98 | journal={European Journal of Nuclear Medicine and Molecular Imaging}, 99 | pages={1--14}, 100 | year={2025}, 101 | publisher={Springer} 102 | } 103 | ``` 104 | 105 | ## Contact 106 | 107 | For any questions or inquiries, please contact boxiao.yu@ufl.edu. 108 | 109 | -------------------------------------------------------------------------------- /environment.yml: -------------------------------------------------------------------------------- 1 | name: PET-3D-DDPM 2 | channels: 3 | - conda-forge 4 | - anaconda 5 | - defaults 6 | dependencies: 7 | - _libgcc_mutex=0.1=main 8 | - _openmp_mutex=5.1=1_gnu 9 | - blas=1.0=mkl 10 | - ca-certificates=2023.08.22=h06a4308_0 11 | - certifi=2022.12.7=py37h06a4308_0 12 | - freetype=2.12.1=h4a9f257_0 13 | - giflib=5.2.1=h5eee18b_3 14 | - intel-openmp=2021.4.0=h06a4308_3561 15 | - jpeg=9e=h5eee18b_1 16 | - lcms2=2.12=h3be6417_0 17 | - ld_impl_linux-64=2.38=h1181459_1 18 | - lerc=3.0=h295c915_0 19 | - libdeflate=1.17=h5eee18b_1 20 | - libffi=3.4.4=h6a678d5_0 21 | - libgcc-ng=11.2.0=h1234567_1 22 | - libgfortran-ng=7.5.0=h14aa051_20 23 | - libgfortran4=7.5.0=h14aa051_20 24 | - libgomp=11.2.0=h1234567_1 25 | - libpng=1.6.39=h5eee18b_0 26 | - libstdcxx-ng=11.2.0=h1234567_1 27 | - libtiff=4.5.1=h6a678d5_0 28 | - libwebp=1.2.4=h11a3e52_1 29 | - libwebp-base=1.2.4=h5eee18b_1 30 | - lz4-c=1.9.4=h6a678d5_0 31 | - mkl=2021.4.0=h06a4308_640 32 | - mkl-service=2.4.0=py37h7f8727e_0 33 | - mkl_fft=1.3.1=py37hd3c417c_0 34 | - mkl_random=1.2.2=py37h51133e4_0 35 | - mpi=1.0=openmpi 36 | - mpi4py=3.1.4=py37h3e5f7c9_0 37 | - ncurses=6.4=h6a678d5_0 38 | - numpy=1.21.5=py37h6c91a56_3 39 | - numpy-base=1.21.5=py37ha15fc14_3 40 | - openmpi=4.0.4=hdf1f1ad_0 41 | - openssl=1.1.1w=h7f8727e_0 42 | - pillow=9.4.0=py37h6a678d5_0 43 | - python=3.7.16=h7a1cb2a_0 44 | - readline=8.2=h5eee18b_0 45 | - setuptools=65.6.3=py37h06a4308_0 46 | - six=1.16.0=pyhd3eb1b0_1 47 | - sqlite=3.41.2=h5eee18b_0 48 | - tk=8.6.12=h1ccaba5_0 49 | - wheel=0.38.4=py37h06a4308_0 50 | - xz=5.4.2=h5eee18b_0 51 | - zlib=1.2.13=h5eee18b_0 52 | - zstd=1.5.5=hc292b87_0 53 | - pip: 54 | - blobfile==2.0.2 55 | - filelock==3.12.2 56 | - lxml==4.9.3 57 | - nvidia-cublas-cu11==11.10.3.66 58 | - nvidia-cuda-nvrtc-cu11==11.7.99 59 | - nvidia-cuda-runtime-cu11==11.7.99 60 | - nvidia-cudnn-cu11==8.5.0.96 61 | - pip==23.3 62 | - pycryptodomex==3.19.0 63 | - torch==1.13.1 64 | - tqdm==4.66.1 65 | - typing-extensions==4.7.1 66 | - urllib3==2.0.7 67 | 68 | -------------------------------------------------------------------------------- /figs/model_framework.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Miche11eU/PET-Image-Denoising-Using-3D-Diffusion-Model/dd4bf6232e7adc9a998eafa769e1482c8075b1fb/figs/model_framework.png -------------------------------------------------------------------------------- /figs/results.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Miche11eU/PET-Image-Denoising-Using-3D-Diffusion-Model/dd4bf6232e7adc9a998eafa769e1482c8075b1fb/figs/results.png -------------------------------------------------------------------------------- /guided_diffusion/__init__.py: -------------------------------------------------------------------------------- 1 | """ 2 | Codebase for "Improved Denoising Diffusion Probabilistic Models". 3 | """ 4 | -------------------------------------------------------------------------------- /guided_diffusion/__pycache__/__init__.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Miche11eU/PET-Image-Denoising-Using-3D-Diffusion-Model/dd4bf6232e7adc9a998eafa769e1482c8075b1fb/guided_diffusion/__pycache__/__init__.cpython-37.pyc -------------------------------------------------------------------------------- /guided_diffusion/__pycache__/__init__.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Miche11eU/PET-Image-Denoising-Using-3D-Diffusion-Model/dd4bf6232e7adc9a998eafa769e1482c8075b1fb/guided_diffusion/__pycache__/__init__.cpython-38.pyc -------------------------------------------------------------------------------- /guided_diffusion/__pycache__/__init__.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Miche11eU/PET-Image-Denoising-Using-3D-Diffusion-Model/dd4bf6232e7adc9a998eafa769e1482c8075b1fb/guided_diffusion/__pycache__/__init__.cpython-39.pyc -------------------------------------------------------------------------------- /guided_diffusion/__pycache__/dist_util.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Miche11eU/PET-Image-Denoising-Using-3D-Diffusion-Model/dd4bf6232e7adc9a998eafa769e1482c8075b1fb/guided_diffusion/__pycache__/dist_util.cpython-37.pyc -------------------------------------------------------------------------------- /guided_diffusion/__pycache__/dist_util.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Miche11eU/PET-Image-Denoising-Using-3D-Diffusion-Model/dd4bf6232e7adc9a998eafa769e1482c8075b1fb/guided_diffusion/__pycache__/dist_util.cpython-38.pyc -------------------------------------------------------------------------------- /guided_diffusion/__pycache__/dist_util.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Miche11eU/PET-Image-Denoising-Using-3D-Diffusion-Model/dd4bf6232e7adc9a998eafa769e1482c8075b1fb/guided_diffusion/__pycache__/dist_util.cpython-39.pyc -------------------------------------------------------------------------------- /guided_diffusion/__pycache__/fp16_util.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Miche11eU/PET-Image-Denoising-Using-3D-Diffusion-Model/dd4bf6232e7adc9a998eafa769e1482c8075b1fb/guided_diffusion/__pycache__/fp16_util.cpython-37.pyc -------------------------------------------------------------------------------- /guided_diffusion/__pycache__/fp16_util.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Miche11eU/PET-Image-Denoising-Using-3D-Diffusion-Model/dd4bf6232e7adc9a998eafa769e1482c8075b1fb/guided_diffusion/__pycache__/fp16_util.cpython-39.pyc -------------------------------------------------------------------------------- /guided_diffusion/__pycache__/gaussian_diffusion.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Miche11eU/PET-Image-Denoising-Using-3D-Diffusion-Model/dd4bf6232e7adc9a998eafa769e1482c8075b1fb/guided_diffusion/__pycache__/gaussian_diffusion.cpython-37.pyc -------------------------------------------------------------------------------- /guided_diffusion/__pycache__/gaussian_diffusion.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Miche11eU/PET-Image-Denoising-Using-3D-Diffusion-Model/dd4bf6232e7adc9a998eafa769e1482c8075b1fb/guided_diffusion/__pycache__/gaussian_diffusion.cpython-39.pyc -------------------------------------------------------------------------------- /guided_diffusion/__pycache__/image_datasets.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Miche11eU/PET-Image-Denoising-Using-3D-Diffusion-Model/dd4bf6232e7adc9a998eafa769e1482c8075b1fb/guided_diffusion/__pycache__/image_datasets.cpython-37.pyc -------------------------------------------------------------------------------- /guided_diffusion/__pycache__/image_datasets.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Miche11eU/PET-Image-Denoising-Using-3D-Diffusion-Model/dd4bf6232e7adc9a998eafa769e1482c8075b1fb/guided_diffusion/__pycache__/image_datasets.cpython-39.pyc -------------------------------------------------------------------------------- /guided_diffusion/__pycache__/logger.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Miche11eU/PET-Image-Denoising-Using-3D-Diffusion-Model/dd4bf6232e7adc9a998eafa769e1482c8075b1fb/guided_diffusion/__pycache__/logger.cpython-37.pyc -------------------------------------------------------------------------------- /guided_diffusion/__pycache__/logger.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Miche11eU/PET-Image-Denoising-Using-3D-Diffusion-Model/dd4bf6232e7adc9a998eafa769e1482c8075b1fb/guided_diffusion/__pycache__/logger.cpython-39.pyc -------------------------------------------------------------------------------- /guided_diffusion/__pycache__/losses.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Miche11eU/PET-Image-Denoising-Using-3D-Diffusion-Model/dd4bf6232e7adc9a998eafa769e1482c8075b1fb/guided_diffusion/__pycache__/losses.cpython-37.pyc -------------------------------------------------------------------------------- /guided_diffusion/__pycache__/losses.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Miche11eU/PET-Image-Denoising-Using-3D-Diffusion-Model/dd4bf6232e7adc9a998eafa769e1482c8075b1fb/guided_diffusion/__pycache__/losses.cpython-39.pyc -------------------------------------------------------------------------------- /guided_diffusion/__pycache__/nn.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Miche11eU/PET-Image-Denoising-Using-3D-Diffusion-Model/dd4bf6232e7adc9a998eafa769e1482c8075b1fb/guided_diffusion/__pycache__/nn.cpython-37.pyc -------------------------------------------------------------------------------- /guided_diffusion/__pycache__/nn.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Miche11eU/PET-Image-Denoising-Using-3D-Diffusion-Model/dd4bf6232e7adc9a998eafa769e1482c8075b1fb/guided_diffusion/__pycache__/nn.cpython-39.pyc -------------------------------------------------------------------------------- /guided_diffusion/__pycache__/resample.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Miche11eU/PET-Image-Denoising-Using-3D-Diffusion-Model/dd4bf6232e7adc9a998eafa769e1482c8075b1fb/guided_diffusion/__pycache__/resample.cpython-37.pyc -------------------------------------------------------------------------------- /guided_diffusion/__pycache__/resample.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Miche11eU/PET-Image-Denoising-Using-3D-Diffusion-Model/dd4bf6232e7adc9a998eafa769e1482c8075b1fb/guided_diffusion/__pycache__/resample.cpython-39.pyc -------------------------------------------------------------------------------- /guided_diffusion/__pycache__/respace.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Miche11eU/PET-Image-Denoising-Using-3D-Diffusion-Model/dd4bf6232e7adc9a998eafa769e1482c8075b1fb/guided_diffusion/__pycache__/respace.cpython-37.pyc -------------------------------------------------------------------------------- /guided_diffusion/__pycache__/respace.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Miche11eU/PET-Image-Denoising-Using-3D-Diffusion-Model/dd4bf6232e7adc9a998eafa769e1482c8075b1fb/guided_diffusion/__pycache__/respace.cpython-39.pyc -------------------------------------------------------------------------------- /guided_diffusion/__pycache__/script_util.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Miche11eU/PET-Image-Denoising-Using-3D-Diffusion-Model/dd4bf6232e7adc9a998eafa769e1482c8075b1fb/guided_diffusion/__pycache__/script_util.cpython-37.pyc -------------------------------------------------------------------------------- /guided_diffusion/__pycache__/script_util.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Miche11eU/PET-Image-Denoising-Using-3D-Diffusion-Model/dd4bf6232e7adc9a998eafa769e1482c8075b1fb/guided_diffusion/__pycache__/script_util.cpython-39.pyc -------------------------------------------------------------------------------- /guided_diffusion/__pycache__/train_util.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Miche11eU/PET-Image-Denoising-Using-3D-Diffusion-Model/dd4bf6232e7adc9a998eafa769e1482c8075b1fb/guided_diffusion/__pycache__/train_util.cpython-37.pyc -------------------------------------------------------------------------------- /guided_diffusion/__pycache__/train_util.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Miche11eU/PET-Image-Denoising-Using-3D-Diffusion-Model/dd4bf6232e7adc9a998eafa769e1482c8075b1fb/guided_diffusion/__pycache__/train_util.cpython-39.pyc -------------------------------------------------------------------------------- /guided_diffusion/__pycache__/unet.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Miche11eU/PET-Image-Denoising-Using-3D-Diffusion-Model/dd4bf6232e7adc9a998eafa769e1482c8075b1fb/guided_diffusion/__pycache__/unet.cpython-37.pyc -------------------------------------------------------------------------------- /guided_diffusion/__pycache__/unet.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Miche11eU/PET-Image-Denoising-Using-3D-Diffusion-Model/dd4bf6232e7adc9a998eafa769e1482c8075b1fb/guided_diffusion/__pycache__/unet.cpython-39.pyc -------------------------------------------------------------------------------- /guided_diffusion/dist_util.py: -------------------------------------------------------------------------------- 1 | """ 2 | Helpers for distributed training. 3 | """ 4 | 5 | import io 6 | import os 7 | import socket 8 | 9 | import blobfile as bf 10 | from mpi4py import MPI 11 | import torch as th 12 | import torch.distributed as dist 13 | 14 | # Change this to reflect your cluster layout. 15 | # The GPU for a given rank is (rank % GPUS_PER_NODE). 16 | GPUS_PER_NODE = 8 17 | 18 | SETUP_RETRY_COUNT = 3 19 | 20 | 21 | def setup_dist(): 22 | """ 23 | Setup a distributed process group. 24 | """ 25 | if dist.is_initialized(): 26 | return 27 | os.environ["CUDA_VISIBLE_DEVICES"] = f"{MPI.COMM_WORLD.Get_rank() % GPUS_PER_NODE}" 28 | # os.environ['CUDA_VISIBLE_DEVICES'] = "3" 29 | 30 | comm = MPI.COMM_WORLD 31 | backend = "gloo" if not th.cuda.is_available() else "nccl" 32 | 33 | if backend == "gloo": 34 | hostname = "localhost" 35 | else: 36 | hostname = socket.gethostbyname(socket.getfqdn()) 37 | os.environ["MASTER_ADDR"] = comm.bcast(hostname, root=0) 38 | os.environ["RANK"] = str(comm.rank) 39 | os.environ["WORLD_SIZE"] = str(comm.size) 40 | 41 | port = comm.bcast(_find_free_port(), root=0) 42 | os.environ["MASTER_PORT"] = str(port) 43 | dist.init_process_group(backend=backend, init_method="env://") 44 | 45 | 46 | def dev(): 47 | """ 48 | Get the device to use for torch.distributed. 49 | """ 50 | if th.cuda.is_available(): 51 | return th.device(f"cuda") 52 | return th.device("cpu") 53 | 54 | 55 | def load_state_dict(path, **kwargs): 56 | """ 57 | Load a PyTorch file without redundant fetches across MPI ranks. 58 | """ 59 | chunk_size = 2 ** 30 # MPI has a relatively small size limit 60 | if MPI.COMM_WORLD.Get_rank() == 0: 61 | with bf.BlobFile(path, "rb") as f: 62 | data = f.read() 63 | num_chunks = len(data) // chunk_size 64 | if len(data) % chunk_size: 65 | num_chunks += 1 66 | MPI.COMM_WORLD.bcast(num_chunks) 67 | for i in range(0, len(data), chunk_size): 68 | MPI.COMM_WORLD.bcast(data[i : i + chunk_size]) 69 | else: 70 | num_chunks = MPI.COMM_WORLD.bcast(None) 71 | data = bytes() 72 | for _ in range(num_chunks): 73 | data += MPI.COMM_WORLD.bcast(None) 74 | 75 | return th.load(io.BytesIO(data), **kwargs) 76 | 77 | 78 | def sync_params(params): 79 | """ 80 | Synchronize a sequence of Tensors across ranks from rank 0. 81 | """ 82 | for p in params: 83 | with th.no_grad(): 84 | dist.broadcast(p, 0) 85 | 86 | 87 | def _find_free_port(): 88 | try: 89 | s = socket.socket(socket.AF_INET, socket.SOCK_STREAM) 90 | s.bind(("", 0)) 91 | s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) 92 | return s.getsockname()[1] 93 | finally: 94 | s.close() 95 | -------------------------------------------------------------------------------- /guided_diffusion/fp16_util.py: -------------------------------------------------------------------------------- 1 | """ 2 | Helpers to train with 16-bit precision. 3 | """ 4 | 5 | import numpy as np 6 | import torch as th 7 | import torch.nn as nn 8 | from torch._utils import _flatten_dense_tensors, _unflatten_dense_tensors 9 | 10 | from . import logger 11 | 12 | INITIAL_LOG_LOSS_SCALE = 20.0 13 | 14 | 15 | def convert_module_to_f16(l): 16 | """ 17 | Convert primitive modules to float16. 18 | """ 19 | if isinstance(l, (nn.Conv1d, nn.Conv2d, nn.Conv3d)): 20 | l.weight.data = l.weight.data.half() 21 | if l.bias is not None: 22 | l.bias.data = l.bias.data.half() 23 | 24 | 25 | def convert_module_to_f32(l): 26 | """ 27 | Convert primitive modules to float32, undoing convert_module_to_f16(). 28 | """ 29 | if isinstance(l, (nn.Conv1d, nn.Conv2d, nn.Conv3d)): 30 | l.weight.data = l.weight.data.float() 31 | if l.bias is not None: 32 | l.bias.data = l.bias.data.float() 33 | 34 | 35 | def make_master_params(param_groups_and_shapes): 36 | """ 37 | Copy model parameters into a (differently-shaped) list of full-precision 38 | parameters. 39 | """ 40 | master_params = [] 41 | for param_group, shape in param_groups_and_shapes: 42 | master_param = nn.Parameter( 43 | _flatten_dense_tensors( 44 | [param.detach().float() for (_, param) in param_group] 45 | ).view(shape) 46 | ) 47 | master_param.requires_grad = True 48 | master_params.append(master_param) 49 | return master_params 50 | 51 | 52 | def model_grads_to_master_grads(param_groups_and_shapes, master_params): 53 | """ 54 | Copy the gradients from the model parameters into the master parameters 55 | from make_master_params(). 56 | """ 57 | for master_param, (param_group, shape) in zip( 58 | master_params, param_groups_and_shapes 59 | ): 60 | master_param.grad = _flatten_dense_tensors( 61 | [param_grad_or_zeros(param) for (_, param) in param_group] 62 | ).view(shape) 63 | 64 | 65 | def master_params_to_model_params(param_groups_and_shapes, master_params): 66 | """ 67 | Copy the master parameter data back into the model parameters. 68 | """ 69 | # Without copying to a list, if a generator is passed, this will 70 | # silently not copy any parameters. 71 | for master_param, (param_group, _) in zip(master_params, param_groups_and_shapes): 72 | for (_, param), unflat_master_param in zip( 73 | param_group, unflatten_master_params(param_group, master_param.view(-1)) 74 | ): 75 | param.detach().copy_(unflat_master_param) 76 | 77 | 78 | def unflatten_master_params(param_group, master_param): 79 | return _unflatten_dense_tensors(master_param, [param for (_, param) in param_group]) 80 | 81 | 82 | def get_param_groups_and_shapes(named_model_params): 83 | named_model_params = list(named_model_params) 84 | scalar_vector_named_params = ( 85 | [(n, p) for (n, p) in named_model_params if p.ndim <= 1], 86 | (-1), 87 | ) 88 | matrix_named_params = ( 89 | [(n, p) for (n, p) in named_model_params if p.ndim > 1], 90 | (1, -1), 91 | ) 92 | return [scalar_vector_named_params, matrix_named_params] 93 | 94 | 95 | def master_params_to_state_dict( 96 | model, param_groups_and_shapes, master_params, use_fp16 97 | ): 98 | if use_fp16: 99 | state_dict = model.state_dict() 100 | for master_param, (param_group, _) in zip( 101 | master_params, param_groups_and_shapes 102 | ): 103 | for (name, _), unflat_master_param in zip( 104 | param_group, unflatten_master_params(param_group, master_param.view(-1)) 105 | ): 106 | assert name in state_dict 107 | state_dict[name] = unflat_master_param 108 | else: 109 | state_dict = model.state_dict() 110 | for i, (name, _value) in enumerate(model.named_parameters()): 111 | assert name in state_dict 112 | state_dict[name] = master_params[i] 113 | return state_dict 114 | 115 | 116 | def state_dict_to_master_params(model, state_dict, use_fp16): 117 | if use_fp16: 118 | named_model_params = [ 119 | (name, state_dict[name]) for name, _ in model.named_parameters() 120 | ] 121 | param_groups_and_shapes = get_param_groups_and_shapes(named_model_params) 122 | master_params = make_master_params(param_groups_and_shapes) 123 | else: 124 | master_params = [state_dict[name] for name, _ in model.named_parameters()] 125 | return master_params 126 | 127 | 128 | def zero_master_grads(master_params): 129 | for param in master_params: 130 | param.grad = None 131 | 132 | 133 | def zero_grad(model_params): 134 | for param in model_params: 135 | # Taken from https://pytorch.org/docs/stable/_modules/torch/optim/optimizer.html#Optimizer.add_param_group 136 | if param.grad is not None: 137 | param.grad.detach_() 138 | param.grad.zero_() 139 | 140 | 141 | def param_grad_or_zeros(param): 142 | if param.grad is not None: 143 | return param.grad.data.detach() 144 | else: 145 | return th.zeros_like(param) 146 | 147 | 148 | class MixedPrecisionTrainer: 149 | def __init__( 150 | self, 151 | *, 152 | model, 153 | use_fp16=False, 154 | fp16_scale_growth=1e-3, 155 | initial_lg_loss_scale=INITIAL_LOG_LOSS_SCALE, 156 | ): 157 | self.model = model 158 | self.use_fp16 = use_fp16 159 | self.fp16_scale_growth = fp16_scale_growth 160 | 161 | self.model_params = list(self.model.parameters()) 162 | self.master_params = self.model_params 163 | self.param_groups_and_shapes = None 164 | self.lg_loss_scale = initial_lg_loss_scale 165 | 166 | if self.use_fp16: 167 | self.param_groups_and_shapes = get_param_groups_and_shapes( 168 | self.model.named_parameters() 169 | ) 170 | self.master_params = make_master_params(self.param_groups_and_shapes) 171 | self.model.convert_to_fp16() 172 | 173 | def zero_grad(self): 174 | zero_grad(self.model_params) 175 | 176 | def backward(self, loss: th.Tensor): 177 | if self.use_fp16: 178 | loss_scale = 2 ** self.lg_loss_scale 179 | (loss * loss_scale).backward() 180 | else: 181 | loss.backward() 182 | 183 | def optimize(self, opt: th.optim.Optimizer): 184 | if self.use_fp16: 185 | return self._optimize_fp16(opt) 186 | else: 187 | return self._optimize_normal(opt) 188 | 189 | def _optimize_fp16(self, opt: th.optim.Optimizer): 190 | logger.logkv_mean("lg_loss_scale", self.lg_loss_scale) 191 | model_grads_to_master_grads(self.param_groups_and_shapes, self.master_params) 192 | grad_norm, param_norm = self._compute_norms(grad_scale=2 ** self.lg_loss_scale) 193 | if check_overflow(grad_norm): 194 | self.lg_loss_scale -= 1 195 | logger.log(f"Found NaN, decreased lg_loss_scale to {self.lg_loss_scale}") 196 | zero_master_grads(self.master_params) 197 | return False 198 | 199 | logger.logkv_mean("grad_norm", grad_norm) 200 | logger.logkv_mean("param_norm", param_norm) 201 | 202 | for p in self.master_params: 203 | p.grad.mul_(1.0 / (2 ** self.lg_loss_scale)) 204 | opt.step() 205 | zero_master_grads(self.master_params) 206 | master_params_to_model_params(self.param_groups_and_shapes, self.master_params) 207 | self.lg_loss_scale += self.fp16_scale_growth 208 | return True 209 | 210 | def _optimize_normal(self, opt: th.optim.Optimizer): 211 | grad_norm, param_norm = self._compute_norms() 212 | logger.logkv_mean("grad_norm", grad_norm) 213 | logger.logkv_mean("param_norm", param_norm) 214 | opt.step() 215 | return True 216 | 217 | def _compute_norms(self, grad_scale=1.0): 218 | grad_norm = 0.0 219 | param_norm = 0.0 220 | for p in self.master_params: 221 | with th.no_grad(): 222 | param_norm += th.norm(p, p=2, dtype=th.float32).item() ** 2 223 | if p.grad is not None: 224 | grad_norm += th.norm(p.grad, p=2, dtype=th.float32).item() ** 2 225 | return np.sqrt(grad_norm) / grad_scale, np.sqrt(param_norm) 226 | 227 | def master_params_to_state_dict(self, master_params): 228 | return master_params_to_state_dict( 229 | self.model, self.param_groups_and_shapes, master_params, self.use_fp16 230 | ) 231 | 232 | def state_dict_to_master_params(self, state_dict): 233 | return state_dict_to_master_params(self.model, state_dict, self.use_fp16) 234 | 235 | 236 | def check_overflow(value): 237 | return (value == float("inf")) or (value == -float("inf")) or (value != value) 238 | -------------------------------------------------------------------------------- /guided_diffusion/gaussian_diffusion.py: -------------------------------------------------------------------------------- 1 | """ 2 | This code started out as a PyTorch port of Ho et al's diffusion models: 3 | https://github.com/hojonathanho/diffusion/blob/1e0dceb3b3495bbe19116a5e1b3596cd0706c543/diffusion_tf/diffusion_utils_2.py 4 | 5 | Docstrings have been added, as well as DDIM sampling and a new collection of beta schedules. 6 | """ 7 | 8 | import enum 9 | import math 10 | 11 | import numpy as np 12 | import torch as th 13 | 14 | from .nn import mean_flat 15 | from .losses import normal_kl, discretized_gaussian_log_likelihood 16 | 17 | 18 | def get_named_beta_schedule(schedule_name, num_diffusion_timesteps): 19 | """ 20 | Get a pre-defined beta schedule for the given name. 21 | 22 | The beta schedule library consists of beta schedules which remain similar 23 | in the limit of num_diffusion_timesteps. 24 | Beta schedules may be added, but should not be removed or changed once 25 | they are committed to maintain backwards compatibility. 26 | """ 27 | if schedule_name == "linear": 28 | # Linear schedule from Ho et al, extended to work for any number of 29 | # diffusion steps. 30 | scale = 1000 / num_diffusion_timesteps 31 | beta_start = scale * 0.0001 32 | beta_end = scale * 0.02 33 | return np.linspace( 34 | beta_start, beta_end, num_diffusion_timesteps, dtype=np.float64 35 | ) 36 | elif schedule_name == "cosine": 37 | return betas_for_alpha_bar( 38 | num_diffusion_timesteps, 39 | lambda t: math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2, 40 | ) 41 | else: 42 | raise NotImplementedError(f"unknown beta schedule: {schedule_name}") 43 | 44 | 45 | def betas_for_alpha_bar(num_diffusion_timesteps, alpha_bar, max_beta=0.999): 46 | """ 47 | Create a beta schedule that discretizes the given alpha_t_bar function, 48 | which defines the cumulative product of (1-beta) over time from t = [0,1]. 49 | 50 | :param num_diffusion_timesteps: the number of betas to produce. 51 | :param alpha_bar: a lambda that takes an argument t from 0 to 1 and 52 | produces the cumulative product of (1-beta) up to that 53 | part of the diffusion process. 54 | :param max_beta: the maximum beta to use; use values lower than 1 to 55 | prevent singularities. 56 | """ 57 | betas = [] 58 | for i in range(num_diffusion_timesteps): 59 | t1 = i / num_diffusion_timesteps 60 | t2 = (i + 1) / num_diffusion_timesteps 61 | betas.append(min(1 - alpha_bar(t2) / alpha_bar(t1), max_beta)) 62 | return np.array(betas) 63 | 64 | 65 | class ModelMeanType(enum.Enum): 66 | """ 67 | Which type of output the model predicts. 68 | """ 69 | 70 | PREVIOUS_X = enum.auto() # the model predicts x_{t-1} 71 | START_X = enum.auto() # the model predicts x_0 72 | EPSILON = enum.auto() # the model predicts epsilon 73 | 74 | 75 | class ModelVarType(enum.Enum): 76 | """ 77 | What is used as the model's output variance. 78 | 79 | The LEARNED_RANGE option has been added to allow the model to predict 80 | values between FIXED_SMALL and FIXED_LARGE, making its job easier. 81 | """ 82 | 83 | LEARNED = enum.auto() 84 | FIXED_SMALL = enum.auto() 85 | FIXED_LARGE = enum.auto() 86 | LEARNED_RANGE = enum.auto() 87 | 88 | 89 | class LossType(enum.Enum): 90 | MSE = enum.auto() # use raw MSE loss (and KL when learning variances) 91 | RESCALED_MSE = ( 92 | enum.auto() 93 | ) # use raw MSE loss (with RESCALED_KL when learning variances) 94 | KL = enum.auto() # use the variational lower-bound 95 | RESCALED_KL = enum.auto() # like KL, but rescale to estimate the full VLB 96 | 97 | def is_vb(self): 98 | return self == LossType.KL or self == LossType.RESCALED_KL 99 | 100 | 101 | class GaussianDiffusion: 102 | """ 103 | Utilities for training and sampling diffusion models. 104 | 105 | Ported directly from here, and then adapted over time to further experimentation. 106 | https://github.com/hojonathanho/diffusion/blob/1e0dceb3b3495bbe19116a5e1b3596cd0706c543/diffusion_tf/diffusion_utils_2.py#L42 107 | 108 | :param betas: a 1-D numpy array of betas for each diffusion timestep, 109 | starting at T and going to 1. 110 | :param model_mean_type: a ModelMeanType determining what the model outputs. 111 | :param model_var_type: a ModelVarType determining how variance is output. 112 | :param loss_type: a LossType determining the loss function to use. 113 | :param rescale_timesteps: if True, pass floating point timesteps into the 114 | model so that they are always scaled like in the 115 | original paper (0 to 1000). 116 | """ 117 | 118 | def __init__( 119 | self, 120 | *, 121 | betas, 122 | model_mean_type, 123 | model_var_type, 124 | loss_type, 125 | rescale_timesteps=False, 126 | ): 127 | self.model_mean_type = model_mean_type 128 | self.model_var_type = model_var_type 129 | self.loss_type = loss_type 130 | self.rescale_timesteps = rescale_timesteps 131 | 132 | # Use float64 for accuracy. 133 | betas = np.array(betas, dtype=np.float64) 134 | self.betas = betas 135 | assert len(betas.shape) == 1, "betas must be 1-D" 136 | assert (betas > 0).all() and (betas <= 1).all() 137 | 138 | self.num_timesteps = int(betas.shape[0]) 139 | 140 | alphas = 1.0 - betas 141 | self.alphas_cumprod = np.cumprod(alphas, axis=0) 142 | self.alphas_cumprod_prev = np.append(1.0, self.alphas_cumprod[:-1]) 143 | self.alphas_cumprod_next = np.append(self.alphas_cumprod[1:], 0.0) 144 | assert self.alphas_cumprod_prev.shape == (self.num_timesteps,) 145 | 146 | # calculations for diffusion q(x_t | x_{t-1}) and others 147 | self.sqrt_alphas_cumprod = np.sqrt(self.alphas_cumprod) 148 | self.sqrt_one_minus_alphas_cumprod = np.sqrt(1.0 - self.alphas_cumprod) 149 | self.log_one_minus_alphas_cumprod = np.log(1.0 - self.alphas_cumprod) 150 | self.sqrt_recip_alphas_cumprod = np.sqrt(1.0 / self.alphas_cumprod) 151 | self.sqrt_recipm1_alphas_cumprod = np.sqrt(1.0 / self.alphas_cumprod - 1) 152 | 153 | # calculations for posterior q(x_{t-1} | x_t, x_0) 154 | self.posterior_variance = ( 155 | betas * (1.0 - self.alphas_cumprod_prev) / (1.0 - self.alphas_cumprod) 156 | ) 157 | # log calculation clipped because the posterior variance is 0 at the 158 | # beginning of the diffusion chain. 159 | self.posterior_log_variance_clipped = np.log( 160 | np.append(self.posterior_variance[1], self.posterior_variance[1:]) 161 | ) 162 | self.posterior_mean_coef1 = ( 163 | betas * np.sqrt(self.alphas_cumprod_prev) / (1.0 - self.alphas_cumprod) 164 | ) 165 | self.posterior_mean_coef2 = ( 166 | (1.0 - self.alphas_cumprod_prev) 167 | * np.sqrt(alphas) 168 | / (1.0 - self.alphas_cumprod) 169 | ) 170 | 171 | def q_mean_variance(self, x_start, t): 172 | """ 173 | Get the distribution q(x_t | x_0). 174 | 175 | :param x_start: the [N x C x ...] tensor of noiseless inputs. 176 | :param t: the number of diffusion steps (minus 1). Here, 0 means one step. 177 | :return: A tuple (mean, variance, log_variance), all of x_start's shape. 178 | """ 179 | mean = ( 180 | _extract_into_tensor(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start 181 | ) 182 | variance = _extract_into_tensor(1.0 - self.alphas_cumprod, t, x_start.shape) 183 | log_variance = _extract_into_tensor( 184 | self.log_one_minus_alphas_cumprod, t, x_start.shape 185 | ) 186 | return mean, variance, log_variance 187 | 188 | def q_sample(self, x_start, t, noise=None): 189 | """ 190 | Diffuse the data for a given number of diffusion steps. 191 | 192 | In other words, sample from q(x_t | x_0). 193 | 194 | :param x_start: the initial data batch. 195 | :param t: the number of diffusion steps (minus 1). Here, 0 means one step. 196 | :param noise: if specified, the split-out normal noise. 197 | :return: A noisy version of x_start. 198 | """ 199 | if noise is None: 200 | noise = th.randn_like(x_start) 201 | assert noise.shape == x_start.shape 202 | return ( 203 | _extract_into_tensor(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start 204 | + _extract_into_tensor(self.sqrt_one_minus_alphas_cumprod, t, x_start.shape) 205 | * noise 206 | ) 207 | 208 | def q_posterior_mean_variance(self, x_start, x_t, t): 209 | """ 210 | Compute the mean and variance of the diffusion posterior: 211 | 212 | q(x_{t-1} | x_t, x_0) 213 | 214 | """ 215 | assert x_start.shape == x_t.shape 216 | posterior_mean = ( 217 | _extract_into_tensor(self.posterior_mean_coef1, t, x_t.shape) * x_start 218 | + _extract_into_tensor(self.posterior_mean_coef2, t, x_t.shape) * x_t 219 | ) 220 | posterior_variance = _extract_into_tensor(self.posterior_variance, t, x_t.shape) 221 | posterior_log_variance_clipped = _extract_into_tensor( 222 | self.posterior_log_variance_clipped, t, x_t.shape 223 | ) 224 | assert ( 225 | posterior_mean.shape[0] 226 | == posterior_variance.shape[0] 227 | == posterior_log_variance_clipped.shape[0] 228 | == x_start.shape[0] 229 | ) 230 | return posterior_mean, posterior_variance, posterior_log_variance_clipped 231 | 232 | def p_mean_variance( 233 | self, model, x, t, clip_denoised=True, denoised_fn=None, model_kwargs=None 234 | ): 235 | """ 236 | Apply the model to get p(x_{t-1} | x_t), as well as a prediction of 237 | the initial x, x_0. 238 | 239 | :param model: the model, which takes a signal and a batch of timesteps 240 | as input. 241 | :param x: the [N x C x ...] tensor at time t. 242 | :param t: a 1-D Tensor of timesteps. 243 | :param clip_denoised: if True, clip the denoised signal into [-1, 1]. 244 | :param denoised_fn: if not None, a function which applies to the 245 | x_start prediction before it is used to sample. Applies before 246 | clip_denoised. 247 | :param model_kwargs: if not None, a dict of extra keyword arguments to 248 | pass to the model. This can be used for conditioning. 249 | :return: a dict with the following keys: 250 | - 'mean': the model mean output. 251 | - 'variance': the model variance output. 252 | - 'log_variance': the log of 'variance'. 253 | - 'pred_xstart': the prediction for x_0. 254 | """ 255 | if model_kwargs is None: 256 | model_kwargs = {} 257 | 258 | B, C = x.shape[:2] 259 | assert t.shape == (B,) 260 | model_output = model(x, self._scale_timesteps(t), **model_kwargs) 261 | 262 | if self.model_var_type in [ModelVarType.LEARNED, ModelVarType.LEARNED_RANGE]: 263 | assert model_output.shape == (B, C * 2, *x.shape[2:]) 264 | model_output, model_var_values = th.split(model_output, C, dim=1) 265 | if self.model_var_type == ModelVarType.LEARNED: 266 | model_log_variance = model_var_values 267 | model_variance = th.exp(model_log_variance) 268 | else: 269 | min_log = _extract_into_tensor( 270 | self.posterior_log_variance_clipped, t, x.shape 271 | ) 272 | max_log = _extract_into_tensor(np.log(self.betas), t, x.shape) 273 | # The model_var_values is [-1, 1] for [min_var, max_var]. 274 | frac = (model_var_values + 1) / 2 275 | model_log_variance = frac * max_log + (1 - frac) * min_log 276 | model_variance = th.exp(model_log_variance) 277 | else: 278 | model_variance, model_log_variance = { 279 | # for fixedlarge, we set the initial (log-)variance like so 280 | # to get a better decoder log likelihood. 281 | ModelVarType.FIXED_LARGE: ( 282 | np.append(self.posterior_variance[1], self.betas[1:]), 283 | np.log(np.append(self.posterior_variance[1], self.betas[1:])), 284 | ), 285 | ModelVarType.FIXED_SMALL: ( 286 | self.posterior_variance, 287 | self.posterior_log_variance_clipped, 288 | ), 289 | }[self.model_var_type] 290 | model_variance = _extract_into_tensor(model_variance, t, x.shape) 291 | model_log_variance = _extract_into_tensor(model_log_variance, t, x.shape) 292 | 293 | def process_xstart(x): 294 | if denoised_fn is not None: 295 | x = denoised_fn(x) 296 | if clip_denoised: 297 | return x.clamp(-1, 1) 298 | return x 299 | 300 | if self.model_mean_type == ModelMeanType.PREVIOUS_X: 301 | pred_xstart = process_xstart( 302 | self._predict_xstart_from_xprev(x_t=x, t=t, xprev=model_output) 303 | ) 304 | model_mean = model_output 305 | elif self.model_mean_type in [ModelMeanType.START_X, ModelMeanType.EPSILON]: 306 | if self.model_mean_type == ModelMeanType.START_X: 307 | pred_xstart = process_xstart(model_output) 308 | else: 309 | pred_xstart = process_xstart( 310 | self._predict_xstart_from_eps(x_t=x, t=t, eps=model_output) 311 | ) 312 | model_mean, _, _ = self.q_posterior_mean_variance( 313 | x_start=pred_xstart, x_t=x, t=t 314 | ) 315 | else: 316 | raise NotImplementedError(self.model_mean_type) 317 | 318 | assert ( 319 | model_mean.shape == model_log_variance.shape == pred_xstart.shape == x.shape 320 | ) 321 | return { 322 | "mean": model_mean, 323 | "variance": model_variance, 324 | "log_variance": model_log_variance, 325 | "pred_xstart": pred_xstart, 326 | } 327 | 328 | def _predict_xstart_from_eps(self, x_t, t, eps): 329 | assert x_t.shape == eps.shape 330 | return ( 331 | _extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t 332 | - _extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape) * eps 333 | ) 334 | 335 | def _predict_xstart_from_xprev(self, x_t, t, xprev): 336 | assert x_t.shape == xprev.shape 337 | return ( # (xprev - coef2*x_t) / coef1 338 | _extract_into_tensor(1.0 / self.posterior_mean_coef1, t, x_t.shape) * xprev 339 | - _extract_into_tensor( 340 | self.posterior_mean_coef2 / self.posterior_mean_coef1, t, x_t.shape 341 | ) 342 | * x_t 343 | ) 344 | 345 | def _predict_eps_from_xstart(self, x_t, t, pred_xstart): 346 | return ( 347 | _extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t 348 | - pred_xstart 349 | ) / _extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape) 350 | 351 | def _scale_timesteps(self, t): 352 | if self.rescale_timesteps: 353 | return t.float() * (1000.0 / self.num_timesteps) 354 | return t 355 | 356 | def condition_mean(self, cond_fn, p_mean_var, x, t, model_kwargs=None): 357 | """ 358 | Compute the mean for the previous step, given a function cond_fn that 359 | computes the gradient of a conditional log probability with respect to 360 | x. In particular, cond_fn computes grad(log(p(y|x))), and we want to 361 | condition on y. 362 | 363 | This uses the conditioning strategy from Sohl-Dickstein et al. (2015). 364 | """ 365 | gradient = cond_fn(x, self._scale_timesteps(t), **model_kwargs) 366 | new_mean = ( 367 | p_mean_var["mean"].float() + p_mean_var["variance"] * gradient.float() 368 | ) 369 | return new_mean 370 | 371 | def condition_score(self, cond_fn, p_mean_var, x, t, model_kwargs=None): 372 | """ 373 | Compute what the p_mean_variance output would have been, should the 374 | model's score function be conditioned by cond_fn. 375 | 376 | See condition_mean() for details on cond_fn. 377 | 378 | Unlike condition_mean(), this instead uses the conditioning strategy 379 | from Song et al (2020). 380 | """ 381 | alpha_bar = _extract_into_tensor(self.alphas_cumprod, t, x.shape) 382 | 383 | eps = self._predict_eps_from_xstart(x, t, p_mean_var["pred_xstart"]) 384 | eps = eps - (1 - alpha_bar).sqrt() * cond_fn( 385 | x, self._scale_timesteps(t), **model_kwargs 386 | ) 387 | 388 | out = p_mean_var.copy() 389 | out["pred_xstart"] = self._predict_xstart_from_eps(x, t, eps) 390 | out["mean"], _, _ = self.q_posterior_mean_variance( 391 | x_start=out["pred_xstart"], x_t=x, t=t 392 | ) 393 | return out 394 | 395 | def p_sample( 396 | self, 397 | model, 398 | x, 399 | t, 400 | clip_denoised=True, 401 | denoised_fn=None, 402 | cond_fn=None, 403 | model_kwargs=None, 404 | ): 405 | """ 406 | Sample x_{t-1} from the model at the given timestep. 407 | 408 | :param model: the model to sample from. 409 | :param x: the current tensor at x_{t-1}. 410 | :param t: the value of t, starting at 0 for the first diffusion step. 411 | :param clip_denoised: if True, clip the x_start prediction to [-1, 1]. 412 | :param denoised_fn: if not None, a function which applies to the 413 | x_start prediction before it is used to sample. 414 | :param cond_fn: if not None, this is a gradient function that acts 415 | similarly to the model. 416 | :param model_kwargs: if not None, a dict of extra keyword arguments to 417 | pass to the model. This can be used for conditioning. 418 | :return: a dict containing the following keys: 419 | - 'sample': a random sample from the model. 420 | - 'pred_xstart': a prediction of x_0. 421 | """ 422 | out = self.p_mean_variance( 423 | model, 424 | x, 425 | t, 426 | clip_denoised=clip_denoised, 427 | denoised_fn=denoised_fn, 428 | model_kwargs=model_kwargs, 429 | ) 430 | noise = th.randn_like(x) 431 | nonzero_mask = ( 432 | (t != 0).float().view(-1, *([1] * (len(x.shape) - 1))) 433 | ) # no noise when t == 0 434 | if cond_fn is not None: 435 | out["mean"] = self.condition_mean( 436 | cond_fn, out, x, t, model_kwargs=model_kwargs 437 | ) 438 | sample = out["mean"] + nonzero_mask * th.exp(0.5 * out["log_variance"]) * noise 439 | return {"sample": sample, "pred_xstart": out["pred_xstart"]} 440 | 441 | def p_sample_loop( 442 | self, 443 | model, 444 | shape, 445 | noise=None, 446 | clip_denoised=True, 447 | denoised_fn=None, 448 | cond_fn=None, 449 | model_kwargs=None, 450 | device=None, 451 | progress=False, 452 | ): 453 | """ 454 | Generate samples from the model. 455 | 456 | :param model: the model module. 457 | :param shape: the shape of the samples, (N, C, H, W). 458 | :param noise: if specified, the noise from the encoder to sample. 459 | Should be of the same shape as `shape`. 460 | :param clip_denoised: if True, clip x_start predictions to [-1, 1]. 461 | :param denoised_fn: if not None, a function which applies to the 462 | x_start prediction before it is used to sample. 463 | :param cond_fn: if not None, this is a gradient function that acts 464 | similarly to the model. 465 | :param model_kwargs: if not None, a dict of extra keyword arguments to 466 | pass to the model. This can be used for conditioning. 467 | :param device: if specified, the device to create the samples on. 468 | If not specified, use a model parameter's device. 469 | :param progress: if True, show a tqdm progress bar. 470 | :return: a non-differentiable batch of samples. 471 | """ 472 | final = None 473 | for sample in self.p_sample_loop_progressive( 474 | model, 475 | shape, 476 | noise=noise, 477 | clip_denoised=clip_denoised, 478 | denoised_fn=denoised_fn, 479 | cond_fn=cond_fn, 480 | model_kwargs=model_kwargs, 481 | device=device, 482 | progress=progress, 483 | ): 484 | final = sample 485 | return final["sample"] 486 | 487 | def p_sample_loop_progressive( 488 | self, 489 | model, 490 | shape, 491 | noise=None, 492 | clip_denoised=True, 493 | denoised_fn=None, 494 | cond_fn=None, 495 | model_kwargs=None, 496 | device=None, 497 | progress=False, 498 | ): 499 | """ 500 | Generate samples from the model and yield intermediate samples from 501 | each timestep of diffusion. 502 | 503 | Arguments are the same as p_sample_loop(). 504 | Returns a generator over dicts, where each dict is the return value of 505 | p_sample(). 506 | """ 507 | if device is None: 508 | device = next(model.parameters()).device 509 | assert isinstance(shape, (tuple, list)) 510 | if noise is not None: 511 | img = noise #here 512 | else: 513 | img = th.randn(*shape, device=device) 514 | indices = list(range(self.num_timesteps))[::-1] 515 | 516 | if progress: 517 | # Lazy import so that we don't depend on tqdm. 518 | from tqdm.auto import tqdm 519 | 520 | indices = tqdm(indices) 521 | 522 | for i in indices: 523 | t = th.tensor([i] * shape[0], device=device) 524 | with th.no_grad(): 525 | out = self.p_sample( 526 | model, 527 | img, 528 | t, 529 | clip_denoised=clip_denoised, 530 | denoised_fn=denoised_fn, 531 | cond_fn=cond_fn, 532 | model_kwargs=model_kwargs, 533 | ) 534 | yield out 535 | img = out["sample"] 536 | 537 | def ddim_sample( 538 | self, 539 | model, 540 | x, 541 | t, 542 | clip_denoised=True, 543 | denoised_fn=None, 544 | cond_fn=None, 545 | model_kwargs=None, 546 | eta=0.0, 547 | ): 548 | """ 549 | Sample x_{t-1} from the model using DDIM. 550 | 551 | Same usage as p_sample(). 552 | """ 553 | out = self.p_mean_variance( 554 | model, 555 | x, 556 | t, 557 | clip_denoised=clip_denoised, 558 | denoised_fn=denoised_fn, 559 | model_kwargs=model_kwargs, 560 | ) 561 | if cond_fn is not None: 562 | out = self.condition_score(cond_fn, out, x, t, model_kwargs=model_kwargs) 563 | 564 | # Usually our model outputs epsilon, but we re-derive it 565 | # in case we used x_start or x_prev prediction. 566 | eps = self._predict_eps_from_xstart(x, t, out["pred_xstart"]) 567 | 568 | alpha_bar = _extract_into_tensor(self.alphas_cumprod, t, x.shape) 569 | alpha_bar_prev = _extract_into_tensor(self.alphas_cumprod_prev, t, x.shape) 570 | sigma = ( 571 | eta 572 | * th.sqrt((1 - alpha_bar_prev) / (1 - alpha_bar)) 573 | * th.sqrt(1 - alpha_bar / alpha_bar_prev) 574 | ) 575 | # Equation 12. 576 | noise = th.randn_like(x) 577 | mean_pred = ( 578 | out["pred_xstart"] * th.sqrt(alpha_bar_prev) 579 | + th.sqrt(1 - alpha_bar_prev - sigma ** 2) * eps 580 | ) 581 | nonzero_mask = ( 582 | (t != 0).float().view(-1, *([1] * (len(x.shape) - 1))) 583 | ) # no noise when t == 0 584 | sample = mean_pred + nonzero_mask * sigma * noise 585 | return {"sample": sample, "pred_xstart": out["pred_xstart"]} 586 | 587 | def ddim_reverse_sample( 588 | self, 589 | model, 590 | x, 591 | t, 592 | clip_denoised=True, 593 | denoised_fn=None, 594 | model_kwargs=None, 595 | eta=0.0, 596 | ): 597 | """ 598 | Sample x_{t+1} from the model using DDIM reverse ODE. 599 | """ 600 | assert eta == 0.0, "Reverse ODE only for deterministic path" 601 | out = self.p_mean_variance( 602 | model, 603 | x, 604 | t, 605 | clip_denoised=clip_denoised, 606 | denoised_fn=denoised_fn, 607 | model_kwargs=model_kwargs, 608 | ) 609 | # Usually our model outputs epsilon, but we re-derive it 610 | # in case we used x_start or x_prev prediction. 611 | eps = ( 612 | _extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x.shape) * x 613 | - out["pred_xstart"] 614 | ) / _extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, x.shape) 615 | alpha_bar_next = _extract_into_tensor(self.alphas_cumprod_next, t, x.shape) 616 | 617 | # Equation 12. reversed 618 | mean_pred = ( 619 | out["pred_xstart"] * th.sqrt(alpha_bar_next) 620 | + th.sqrt(1 - alpha_bar_next) * eps 621 | ) 622 | 623 | return {"sample": mean_pred, "pred_xstart": out["pred_xstart"]} 624 | 625 | def ddim_sample_loop( 626 | self, 627 | model, 628 | shape, 629 | noise=None, 630 | clip_denoised=True, 631 | denoised_fn=None, 632 | cond_fn=None, 633 | model_kwargs=None, 634 | device=None, 635 | progress=False, 636 | eta=0.0, 637 | ): 638 | """ 639 | Generate samples from the model using DDIM. 640 | 641 | Same usage as p_sample_loop(). 642 | """ 643 | final = None 644 | for sample in self.ddim_sample_loop_progressive( 645 | model, 646 | shape, 647 | noise=noise, 648 | clip_denoised=clip_denoised, 649 | denoised_fn=denoised_fn, 650 | cond_fn=cond_fn, 651 | model_kwargs=model_kwargs, 652 | device=device, 653 | progress=progress, 654 | eta=eta, 655 | ): 656 | final = sample 657 | return final["sample"] 658 | 659 | def ddim_sample_loop_progressive( 660 | self, 661 | model, 662 | shape, 663 | noise=None, 664 | clip_denoised=True, 665 | denoised_fn=None, 666 | cond_fn=None, 667 | model_kwargs=None, 668 | device=None, 669 | progress=False, 670 | eta=0.0, 671 | ): 672 | """ 673 | Use DDIM to sample from the model and yield intermediate samples from 674 | each timestep of DDIM. 675 | 676 | Same usage as p_sample_loop_progressive(). 677 | """ 678 | if device is None: 679 | device = next(model.parameters()).device 680 | assert isinstance(shape, (tuple, list)) 681 | if noise is not None: 682 | img = noise 683 | else: 684 | img = th.randn(*shape, device=device) 685 | indices = list(range(self.num_timesteps))[::-1] 686 | 687 | if progress: 688 | # Lazy import so that we don't depend on tqdm. 689 | from tqdm.auto import tqdm 690 | 691 | indices = tqdm(indices) 692 | 693 | for i in indices: 694 | t = th.tensor([i] * shape[0], device=device) 695 | with th.no_grad(): 696 | out = self.ddim_sample( 697 | model, 698 | img, 699 | t, 700 | clip_denoised=clip_denoised, 701 | denoised_fn=denoised_fn, 702 | cond_fn=cond_fn, 703 | model_kwargs=model_kwargs, 704 | eta=eta, 705 | ) 706 | yield out 707 | img = out["sample"] 708 | 709 | def _vb_terms_bpd( 710 | self, model, x_start, x_t, t, clip_denoised=True, model_kwargs=None 711 | ): 712 | """ 713 | Get a term for the variational lower-bound. 714 | 715 | The resulting units are bits (rather than nats, as one might expect). 716 | This allows for comparison to other papers. 717 | 718 | :return: a dict with the following keys: 719 | - 'output': a shape [N] tensor of NLLs or KLs. 720 | - 'pred_xstart': the x_0 predictions. 721 | """ 722 | true_mean, _, true_log_variance_clipped = self.q_posterior_mean_variance( 723 | x_start=x_start, x_t=x_t, t=t 724 | ) 725 | out = self.p_mean_variance( 726 | model, x_t, t, clip_denoised=clip_denoised, model_kwargs=model_kwargs 727 | ) 728 | kl = normal_kl( 729 | true_mean, true_log_variance_clipped, out["mean"], out["log_variance"] 730 | ) 731 | kl = mean_flat(kl) / np.log(2.0) 732 | 733 | decoder_nll = -discretized_gaussian_log_likelihood( 734 | x_start, means=out["mean"], log_scales=0.5 * out["log_variance"] 735 | ) 736 | assert decoder_nll.shape == x_start.shape 737 | decoder_nll = mean_flat(decoder_nll) / np.log(2.0) 738 | 739 | # At the first timestep return the decoder NLL, 740 | # otherwise return KL(q(x_{t-1}|x_t,x_0) || p(x_{t-1}|x_t)) 741 | output = th.where((t == 0), decoder_nll, kl) 742 | return {"output": output, "pred_xstart": out["pred_xstart"]} 743 | 744 | def training_losses(self, model, x_start, t, model_kwargs=None, noise=None): 745 | """ 746 | Compute training losses for a single timestep. 747 | 748 | :param model: the model to evaluate loss on. 749 | :param x_start: the [N x C x ...] tensor of inputs. 750 | :param t: a batch of timestep indices. 751 | :param model_kwargs: if not None, a dict of extra keyword arguments to 752 | pass to the model. This can be used for conditioning. 753 | :param noise: if specified, the specific Gaussian noise to try to remove. 754 | :return: a dict with the key "loss" containing a tensor of shape [N]. 755 | Some mean or variance settings may also have other keys. 756 | """ 757 | if model_kwargs is None: 758 | model_kwargs = {} 759 | if noise is None: 760 | noise = th.randn_like(x_start) 761 | x_t = self.q_sample(x_start, t, noise=noise) 762 | 763 | terms = {} 764 | 765 | if self.loss_type == LossType.KL or self.loss_type == LossType.RESCALED_KL: 766 | terms["loss"] = self._vb_terms_bpd( 767 | model=model, 768 | x_start=x_start, 769 | x_t=x_t, 770 | t=t, 771 | clip_denoised=False, 772 | model_kwargs=model_kwargs, 773 | )["output"] 774 | if self.loss_type == LossType.RESCALED_KL: 775 | terms["loss"] *= self.num_timesteps 776 | elif self.loss_type == LossType.MSE or self.loss_type == LossType.RESCALED_MSE: 777 | model_output = model(x_t, self._scale_timesteps(t), **model_kwargs) 778 | # print(model_output.shape) 779 | # print(x_t.shape) 780 | 781 | if self.model_var_type in [ 782 | ModelVarType.LEARNED, 783 | ModelVarType.LEARNED_RANGE, 784 | ]: 785 | B, C = x_t.shape[:2] 786 | assert model_output.shape == (B, C * 2, *x_t.shape[2:]) 787 | model_output, model_var_values = th.split(model_output, C, dim=1) 788 | # Learn the variance using the variational bound, but don't let 789 | # it affect our mean prediction. 790 | frozen_out = th.cat([model_output.detach(), model_var_values], dim=1) 791 | terms["vb"] = self._vb_terms_bpd( 792 | model=lambda *args, r=frozen_out: r, 793 | x_start=x_start, 794 | x_t=x_t, 795 | t=t, 796 | clip_denoised=False, 797 | )["output"] 798 | if self.loss_type == LossType.RESCALED_MSE: 799 | # Divide by 1000 for equivalence with initial implementation. 800 | # Without a factor of 1/1000, the VB term hurts the MSE term. 801 | terms["vb"] *= self.num_timesteps / 1000.0 802 | 803 | target = { 804 | ModelMeanType.PREVIOUS_X: self.q_posterior_mean_variance( 805 | x_start=x_start, x_t=x_t, t=t 806 | )[0], 807 | ModelMeanType.START_X: x_start, 808 | ModelMeanType.EPSILON: noise, 809 | }[self.model_mean_type] 810 | assert model_output.shape == target.shape == x_start.shape 811 | terms["mse"] = mean_flat((target - model_output) ** 2) 812 | if "vb" in terms: 813 | terms["loss"] = terms["mse"] + terms["vb"] 814 | else: 815 | terms["loss"] = terms["mse"] 816 | else: 817 | raise NotImplementedError(self.loss_type) 818 | 819 | return terms 820 | 821 | def _prior_bpd(self, x_start): 822 | """ 823 | Get the prior KL term for the variational lower-bound, measured in 824 | bits-per-dim. 825 | 826 | This term can't be optimized, as it only depends on the encoder. 827 | 828 | :param x_start: the [N x C x ...] tensor of inputs. 829 | :return: a batch of [N] KL values (in bits), one per batch element. 830 | """ 831 | batch_size = x_start.shape[0] 832 | t = th.tensor([self.num_timesteps - 1] * batch_size, device=x_start.device) 833 | qt_mean, _, qt_log_variance = self.q_mean_variance(x_start, t) 834 | kl_prior = normal_kl( 835 | mean1=qt_mean, logvar1=qt_log_variance, mean2=0.0, logvar2=0.0 836 | ) 837 | return mean_flat(kl_prior) / np.log(2.0) 838 | 839 | def calc_bpd_loop(self, model, x_start, clip_denoised=True, model_kwargs=None): 840 | """ 841 | Compute the entire variational lower-bound, measured in bits-per-dim, 842 | as well as other related quantities. 843 | 844 | :param model: the model to evaluate loss on. 845 | :param x_start: the [N x C x ...] tensor of inputs. 846 | :param clip_denoised: if True, clip denoised samples. 847 | :param model_kwargs: if not None, a dict of extra keyword arguments to 848 | pass to the model. This can be used for conditioning. 849 | 850 | :return: a dict containing the following keys: 851 | - total_bpd: the total variational lower-bound, per batch element. 852 | - prior_bpd: the prior term in the lower-bound. 853 | - vb: an [N x T] tensor of terms in the lower-bound. 854 | - xstart_mse: an [N x T] tensor of x_0 MSEs for each timestep. 855 | - mse: an [N x T] tensor of epsilon MSEs for each timestep. 856 | """ 857 | device = x_start.device 858 | batch_size = x_start.shape[0] 859 | 860 | vb = [] 861 | xstart_mse = [] 862 | mse = [] 863 | for t in list(range(self.num_timesteps))[::-1]: 864 | t_batch = th.tensor([t] * batch_size, device=device) 865 | noise = th.randn_like(x_start) 866 | x_t = self.q_sample(x_start=x_start, t=t_batch, noise=noise) 867 | # Calculate VLB term at the current timestep 868 | with th.no_grad(): 869 | out = self._vb_terms_bpd( 870 | model, 871 | x_start=x_start, 872 | x_t=x_t, 873 | t=t_batch, 874 | clip_denoised=clip_denoised, 875 | model_kwargs=model_kwargs, 876 | ) 877 | vb.append(out["output"]) 878 | xstart_mse.append(mean_flat((out["pred_xstart"] - x_start) ** 2)) 879 | eps = self._predict_eps_from_xstart(x_t, t_batch, out["pred_xstart"]) 880 | mse.append(mean_flat((eps - noise) ** 2)) 881 | 882 | vb = th.stack(vb, dim=1) 883 | xstart_mse = th.stack(xstart_mse, dim=1) 884 | mse = th.stack(mse, dim=1) 885 | 886 | prior_bpd = self._prior_bpd(x_start) 887 | total_bpd = vb.sum(dim=1) + prior_bpd 888 | return { 889 | "total_bpd": total_bpd, 890 | "prior_bpd": prior_bpd, 891 | "vb": vb, 892 | "xstart_mse": xstart_mse, 893 | "mse": mse, 894 | } 895 | 896 | 897 | def _extract_into_tensor(arr, timesteps, broadcast_shape): 898 | """ 899 | Extract values from a 1-D numpy array for a batch of indices. 900 | 901 | :param arr: the 1-D numpy array. 902 | :param timesteps: a tensor of indices into the array to extract. 903 | :param broadcast_shape: a larger shape of K dimensions with the batch 904 | dimension equal to the length of timesteps. 905 | :return: a tensor of shape [batch_size, 1, ...] where the shape has K dims. 906 | """ 907 | res = th.from_numpy(arr).to(device=timesteps.device)[timesteps].float() 908 | while len(res.shape) < len(broadcast_shape): 909 | res = res[..., None] 910 | return res.expand(broadcast_shape) 911 | -------------------------------------------------------------------------------- /guided_diffusion/image_datasets.py: -------------------------------------------------------------------------------- 1 | import math 2 | import random 3 | import os 4 | 5 | from PIL import Image 6 | import blobfile as bf 7 | from mpi4py import MPI 8 | import numpy as np 9 | from torch.utils.data import DataLoader, Dataset 10 | 11 | def load_data( 12 | *, 13 | data_dir, 14 | batch_size, 15 | image_size, 16 | class_cond=False, 17 | deterministic=False, 18 | random_crop=False, 19 | random_flip=False, 20 | ): 21 | """ 22 | For a dataset, create a generator over (images, kwargs) pairs. 23 | 24 | Each images is an NCHW float tensor, and the kwargs dict contains zero or 25 | more keys, each of which map to a batched Tensor of their own. 26 | The kwargs dict can be used for class labels, in which case the key is "y" 27 | and the values are integer tensors of class labels. 28 | 29 | :param data_dir: a dataset directory. 30 | :param batch_size: the batch size of each returned pair. 31 | :param image_size: the size to which images are resized. 32 | :param class_cond: if True, include a "y" key in returned dicts for class 33 | label. If classes are not available and this is true, an 34 | exception will be raised. 35 | :param deterministic: if True, yield results in a deterministic order. 36 | :param random_crop: if True, randomly crop the images for augmentation. 37 | :param random_flip: if True, randomly flip the images for augmentation. 38 | """ 39 | if not data_dir: 40 | raise ValueError("unspecified data directory") 41 | all_files = _list_image_files_recursively(data_dir) 42 | classes = None 43 | if class_cond: 44 | # Assume classes are the first part of the filename, 45 | # before an underscore. 46 | class_names = [bf.basename(path).split("_")[0] for path in all_files] 47 | sorted_classes = {x: i for i, x in enumerate(sorted(set(class_names)))} 48 | classes = [sorted_classes[x] for x in class_names] 49 | dataset = ImageDataset( 50 | image_size, 51 | all_files, 52 | classes=classes, 53 | shard=MPI.COMM_WORLD.Get_rank(), 54 | num_shards=MPI.COMM_WORLD.Get_size(), 55 | random_crop=random_crop, 56 | random_flip=random_flip, 57 | ) 58 | if deterministic: 59 | loader = DataLoader( 60 | dataset, batch_size=batch_size, shuffle=False, num_workers=1, drop_last=True 61 | ) 62 | else: 63 | loader = DataLoader( 64 | dataset, batch_size=batch_size, shuffle=True, num_workers=1, drop_last=True 65 | ) 66 | while True: 67 | yield from loader 68 | 69 | 70 | def _list_image_files_recursively(data_dir): 71 | results = [] 72 | for entry in sorted(bf.listdir(data_dir)): 73 | full_path = bf.join(data_dir, entry) 74 | ext = entry.split(".")[-1] 75 | if "." in entry and ext.lower() in ["npz"]: 76 | results.append(full_path) 77 | elif bf.isdir(full_path): 78 | results.extend(_list_image_files_recursively(full_path)) 79 | return results 80 | 81 | 82 | 83 | # pet and ct data 3D 84 | # normalize to (-1,1) 85 | class ImageDataset(Dataset): 86 | def __init__( 87 | self, 88 | resolution, 89 | image_paths, 90 | classes=None, 91 | shard=0, 92 | num_shards=1, 93 | random_crop=False, 94 | random_flip=False, 95 | ): 96 | super().__init__() 97 | self.resolution = resolution 98 | self.local_images = image_paths[shard:][::num_shards] 99 | self.local_classes = None if classes is None else classes[shard:][::num_shards] 100 | self.random_crop = random_crop 101 | self.random_flip = random_flip 102 | 103 | def __len__(self): 104 | return len(self.local_images) 105 | 106 | def __getitem__(self, idx): 107 | path = self.local_images[idx] 108 | 109 | pet_data = np.load(path)['arr_0'] 110 | # pet_data = pet_data.astype(np.float32) * 2.0 - 1.0 111 | pet_data = pet_data.astype(np.float32) 112 | 113 | pet_data = pet_data/4 # add nmlz 114 | # pet_data = pet_data * 2.0 - 1.0 115 | 116 | # input 96*96*96 117 | # size = 96 # patch size 118 | # rand_xyz = np.random.randint(0, 144-size+1, 3) 119 | # pet_low = pet_data[0:2,rand_xyz[0]:rand_xyz[0]+size,rand_xyz[1]:rand_xyz[1]+size,rand_xyz[2]:rand_xyz[2]+size].copy() 120 | # label = pet_data[2,rand_xyz[0]:rand_xyz[0]+size,rand_xyz[1]:rand_xyz[1]+size,rand_xyz[2]:rand_xyz[2]+size].copy() 121 | # while label.max() == 0: 122 | # rand_xyz = np.random.randint(0, 144-size+1, 3) 123 | # pet_low = pet_data[0:2,rand_xyz[0]:rand_xyz[0]+size,rand_xyz[1]:rand_xyz[1]+size,rand_xyz[2]:rand_xyz[2]+size].copy() 124 | # label = pet_data[2,rand_xyz[0]:rand_xyz[0]+size,rand_xyz[1]:rand_xyz[1]+size,rand_xyz[2]:rand_xyz[2]+size].copy() 125 | # C, H, W, T = pet_low.shape 126 | 127 | # input 96*96*32 128 | size_xy = 96 # patch size 129 | size_z = 96 #32 130 | rand_x = np.random.randint(0, 180-size_xy+1) 131 | rand_y = np.random.randint(0, 280-size_xy+1) 132 | rand_z = np.random.randint(0, 520-size_z+1) 133 | pet_low = pet_data[0, rand_x:rand_x+size_xy, rand_y:rand_y+size_xy, rand_z:rand_z+size_z].copy() 134 | label = pet_data[1, rand_x:rand_x+size_xy, rand_y:rand_y+size_xy, rand_z:rand_z+size_z].copy() 135 | # while label.max() == label.min(): 136 | # rand_xy = np.random.randint(0, 144-size_xy+1, 2) 137 | # rand_z = np.random.randint(0, 144-size_z+1) 138 | # pet_low = pet_data[0:2, rand_xy[0]:rand_xy[0]+size_xy, rand_xy[1]:rand_xy[1]+size_xy, rand_z:rand_z+size_z].copy() 139 | # label = pet_data[2, rand_xy[0]:rand_xy[0]+size_xy, rand_xy[1]:rand_xy[1]+size_xy, rand_z:rand_z+size_z].copy() 140 | H, W, T = pet_low.shape 141 | 142 | out_dict = {} 143 | if self.local_classes is not None: 144 | out_dict["y"] = np.array(self.local_classes[idx], dtype=np.int64) 145 | return np.transpose(pet_low.reshape((1, H, W, T)), [0, 3, 1, 2]), np.transpose(label.reshape((1, H, W, T)), [0, 3, 1, 2]), out_dict 146 | 147 | 148 | 149 | def center_crop_arr(pil_image, image_size): 150 | # We are not on a new enough PIL to support the `reducing_gap` 151 | # argument, which uses BOX downsampling at powers of two first. 152 | # Thus, we do it by hand to improve downsample quality. 153 | while min(*pil_image.size) >= 2 * image_size: 154 | pil_image = pil_image.resize( 155 | tuple(x // 2 for x in pil_image.size), resample=Image.BOX 156 | ) 157 | 158 | scale = image_size / min(*pil_image.size) 159 | pil_image = pil_image.resize( 160 | tuple(round(x * scale) for x in pil_image.size), resample=Image.BICUBIC 161 | ) 162 | 163 | arr = np.array(pil_image) 164 | crop_y = (arr.shape[0] - image_size) // 2 165 | crop_x = (arr.shape[1] - image_size) // 2 166 | return arr[crop_y : crop_y + image_size, crop_x : crop_x + image_size] 167 | 168 | 169 | def random_crop_arr(pil_image, image_size, min_crop_frac=0.8, max_crop_frac=1.0): 170 | min_smaller_dim_size = math.ceil(image_size / max_crop_frac) 171 | max_smaller_dim_size = math.ceil(image_size / min_crop_frac) 172 | smaller_dim_size = random.randrange(min_smaller_dim_size, max_smaller_dim_size + 1) 173 | 174 | # We are not on a new enough PIL to support the `reducing_gap` 175 | # argument, which uses BOX downsampling at powers of two first. 176 | # Thus, we do it by hand to improve downsample quality. 177 | while min(*pil_image.size) >= 2 * smaller_dim_size: 178 | pil_image = pil_image.resize( 179 | tuple(x // 2 for x in pil_image.size), resample=Image.BOX 180 | ) 181 | 182 | scale = smaller_dim_size / min(*pil_image.size) 183 | pil_image = pil_image.resize( 184 | tuple(round(x * scale) for x in pil_image.size), resample=Image.BICUBIC 185 | ) 186 | 187 | arr = np.array(pil_image) 188 | crop_y = random.randrange(arr.shape[0] - image_size + 1) 189 | crop_x = random.randrange(arr.shape[1] - image_size + 1) 190 | return arr[crop_y : crop_y + image_size, crop_x : crop_x + image_size] 191 | -------------------------------------------------------------------------------- /guided_diffusion/logger.py: -------------------------------------------------------------------------------- 1 | """ 2 | Logger copied from OpenAI baselines to avoid extra RL-based dependencies: 3 | https://github.com/openai/baselines/blob/ea25b9e8b234e6ee1bca43083f8f3cf974143998/baselines/logger.py 4 | """ 5 | 6 | import os 7 | import sys 8 | import shutil 9 | import os.path as osp 10 | import json 11 | import time 12 | import datetime 13 | import tempfile 14 | import warnings 15 | from collections import defaultdict 16 | from contextlib import contextmanager 17 | 18 | DEBUG = 10 19 | INFO = 20 20 | WARN = 30 21 | ERROR = 40 22 | 23 | DISABLED = 50 24 | 25 | 26 | class KVWriter(object): 27 | def writekvs(self, kvs): 28 | raise NotImplementedError 29 | 30 | 31 | class SeqWriter(object): 32 | def writeseq(self, seq): 33 | raise NotImplementedError 34 | 35 | 36 | class HumanOutputFormat(KVWriter, SeqWriter): 37 | def __init__(self, filename_or_file): 38 | if isinstance(filename_or_file, str): 39 | self.file = open(filename_or_file, "wt") 40 | self.own_file = True 41 | else: 42 | assert hasattr(filename_or_file, "read"), ( 43 | "expected file or str, got %s" % filename_or_file 44 | ) 45 | self.file = filename_or_file 46 | self.own_file = False 47 | 48 | def writekvs(self, kvs): 49 | # Create strings for printing 50 | key2str = {} 51 | for (key, val) in sorted(kvs.items()): 52 | if hasattr(val, "__float__"): 53 | valstr = "%-8.3g" % val 54 | else: 55 | valstr = str(val) 56 | key2str[self._truncate(key)] = self._truncate(valstr) 57 | 58 | # Find max widths 59 | if len(key2str) == 0: 60 | print("WARNING: tried to write empty key-value dict") 61 | return 62 | else: 63 | keywidth = max(map(len, key2str.keys())) 64 | valwidth = max(map(len, key2str.values())) 65 | 66 | # Write out the data 67 | dashes = "-" * (keywidth + valwidth + 7) 68 | lines = [dashes] 69 | for (key, val) in sorted(key2str.items(), key=lambda kv: kv[0].lower()): 70 | lines.append( 71 | "| %s%s | %s%s |" 72 | % (key, " " * (keywidth - len(key)), val, " " * (valwidth - len(val))) 73 | ) 74 | lines.append(dashes) 75 | self.file.write("\n".join(lines) + "\n") 76 | 77 | # Flush the output to the file 78 | self.file.flush() 79 | 80 | def _truncate(self, s): 81 | maxlen = 30 82 | return s[: maxlen - 3] + "..." if len(s) > maxlen else s 83 | 84 | def writeseq(self, seq): 85 | seq = list(seq) 86 | for (i, elem) in enumerate(seq): 87 | self.file.write(elem) 88 | if i < len(seq) - 1: # add space unless this is the last one 89 | self.file.write(" ") 90 | self.file.write("\n") 91 | self.file.flush() 92 | 93 | def close(self): 94 | if self.own_file: 95 | self.file.close() 96 | 97 | 98 | class JSONOutputFormat(KVWriter): 99 | def __init__(self, filename): 100 | self.file = open(filename, "wt") 101 | 102 | def writekvs(self, kvs): 103 | for k, v in sorted(kvs.items()): 104 | if hasattr(v, "dtype"): 105 | kvs[k] = float(v) 106 | self.file.write(json.dumps(kvs) + "\n") 107 | self.file.flush() 108 | 109 | def close(self): 110 | self.file.close() 111 | 112 | 113 | class CSVOutputFormat(KVWriter): 114 | def __init__(self, filename): 115 | self.file = open(filename, "w+t") 116 | self.keys = [] 117 | self.sep = "," 118 | 119 | def writekvs(self, kvs): 120 | # Add our current row to the history 121 | extra_keys = list(kvs.keys() - self.keys) 122 | extra_keys.sort() 123 | if extra_keys: 124 | self.keys.extend(extra_keys) 125 | self.file.seek(0) 126 | lines = self.file.readlines() 127 | self.file.seek(0) 128 | for (i, k) in enumerate(self.keys): 129 | if i > 0: 130 | self.file.write(",") 131 | self.file.write(k) 132 | self.file.write("\n") 133 | for line in lines[1:]: 134 | self.file.write(line[:-1]) 135 | self.file.write(self.sep * len(extra_keys)) 136 | self.file.write("\n") 137 | for (i, k) in enumerate(self.keys): 138 | if i > 0: 139 | self.file.write(",") 140 | v = kvs.get(k) 141 | if v is not None: 142 | self.file.write(str(v)) 143 | self.file.write("\n") 144 | self.file.flush() 145 | 146 | def close(self): 147 | self.file.close() 148 | 149 | 150 | class TensorBoardOutputFormat(KVWriter): 151 | """ 152 | Dumps key/value pairs into TensorBoard's numeric format. 153 | """ 154 | 155 | def __init__(self, dir): 156 | os.makedirs(dir, exist_ok=True) 157 | self.dir = dir 158 | self.step = 1 159 | prefix = "events" 160 | path = osp.join(osp.abspath(dir), prefix) 161 | import tensorflow as tf 162 | from tensorflow.python import pywrap_tensorflow 163 | from tensorflow.core.util import event_pb2 164 | from tensorflow.python.util import compat 165 | 166 | self.tf = tf 167 | self.event_pb2 = event_pb2 168 | self.pywrap_tensorflow = pywrap_tensorflow 169 | self.writer = pywrap_tensorflow.EventsWriter(compat.as_bytes(path)) 170 | 171 | def writekvs(self, kvs): 172 | def summary_val(k, v): 173 | kwargs = {"tag": k, "simple_value": float(v)} 174 | return self.tf.Summary.Value(**kwargs) 175 | 176 | summary = self.tf.Summary(value=[summary_val(k, v) for k, v in kvs.items()]) 177 | event = self.event_pb2.Event(wall_time=time.time(), summary=summary) 178 | event.step = ( 179 | self.step 180 | ) # is there any reason why you'd want to specify the step? 181 | self.writer.WriteEvent(event) 182 | self.writer.Flush() 183 | self.step += 1 184 | 185 | def close(self): 186 | if self.writer: 187 | self.writer.Close() 188 | self.writer = None 189 | 190 | 191 | def make_output_format(format, ev_dir, log_suffix=""): 192 | os.makedirs(ev_dir, exist_ok=True) 193 | if format == "stdout": 194 | return HumanOutputFormat(sys.stdout) 195 | elif format == "log": 196 | return HumanOutputFormat(osp.join(ev_dir, "log%s.txt" % log_suffix)) 197 | elif format == "json": 198 | return JSONOutputFormat(osp.join(ev_dir, "progress%s.json" % log_suffix)) 199 | elif format == "csv": 200 | return CSVOutputFormat(osp.join(ev_dir, "progress%s.csv" % log_suffix)) 201 | elif format == "tensorboard": 202 | return TensorBoardOutputFormat(osp.join(ev_dir, "tb%s" % log_suffix)) 203 | else: 204 | raise ValueError("Unknown format specified: %s" % (format,)) 205 | 206 | 207 | # ================================================================ 208 | # API 209 | # ================================================================ 210 | 211 | 212 | def logkv(key, val): 213 | """ 214 | Log a value of some diagnostic 215 | Call this once for each diagnostic quantity, each iteration 216 | If called many times, last value will be used. 217 | """ 218 | get_current().logkv(key, val) 219 | 220 | 221 | def logkv_mean(key, val): 222 | """ 223 | The same as logkv(), but if called many times, values averaged. 224 | """ 225 | get_current().logkv_mean(key, val) 226 | 227 | 228 | def logkvs(d): 229 | """ 230 | Log a dictionary of key-value pairs 231 | """ 232 | for (k, v) in d.items(): 233 | logkv(k, v) 234 | 235 | 236 | def dumpkvs(): 237 | """ 238 | Write all of the diagnostics from the current iteration 239 | """ 240 | return get_current().dumpkvs() 241 | 242 | 243 | def getkvs(): 244 | return get_current().name2val 245 | 246 | 247 | def log(*args, level=INFO): 248 | """ 249 | Write the sequence of args, with no separators, to the console and output files (if you've configured an output file). 250 | """ 251 | get_current().log(*args, level=level) 252 | 253 | 254 | def debug(*args): 255 | log(*args, level=DEBUG) 256 | 257 | 258 | def info(*args): 259 | log(*args, level=INFO) 260 | 261 | 262 | def warn(*args): 263 | log(*args, level=WARN) 264 | 265 | 266 | def error(*args): 267 | log(*args, level=ERROR) 268 | 269 | 270 | def set_level(level): 271 | """ 272 | Set logging threshold on current logger. 273 | """ 274 | get_current().set_level(level) 275 | 276 | 277 | def set_comm(comm): 278 | get_current().set_comm(comm) 279 | 280 | 281 | def get_dir(): 282 | """ 283 | Get directory that log files are being written to. 284 | will be None if there is no output directory (i.e., if you didn't call start) 285 | """ 286 | return get_current().get_dir() 287 | 288 | 289 | record_tabular = logkv 290 | dump_tabular = dumpkvs 291 | 292 | 293 | @contextmanager 294 | def profile_kv(scopename): 295 | logkey = "wait_" + scopename 296 | tstart = time.time() 297 | try: 298 | yield 299 | finally: 300 | get_current().name2val[logkey] += time.time() - tstart 301 | 302 | 303 | def profile(n): 304 | """ 305 | Usage: 306 | @profile("my_func") 307 | def my_func(): code 308 | """ 309 | 310 | def decorator_with_name(func): 311 | def func_wrapper(*args, **kwargs): 312 | with profile_kv(n): 313 | return func(*args, **kwargs) 314 | 315 | return func_wrapper 316 | 317 | return decorator_with_name 318 | 319 | 320 | # ================================================================ 321 | # Backend 322 | # ================================================================ 323 | 324 | 325 | def get_current(): 326 | if Logger.CURRENT is None: 327 | _configure_default_logger() 328 | 329 | return Logger.CURRENT 330 | 331 | 332 | class Logger(object): 333 | DEFAULT = None # A logger with no output files. (See right below class definition) 334 | # So that you can still log to the terminal without setting up any output files 335 | CURRENT = None # Current logger being used by the free functions above 336 | 337 | def __init__(self, dir, output_formats, comm=None): 338 | self.name2val = defaultdict(float) # values this iteration 339 | self.name2cnt = defaultdict(int) 340 | self.level = INFO 341 | self.dir = dir 342 | self.output_formats = output_formats 343 | self.comm = comm 344 | 345 | # Logging API, forwarded 346 | # ---------------------------------------- 347 | def logkv(self, key, val): 348 | self.name2val[key] = val 349 | 350 | def logkv_mean(self, key, val): 351 | oldval, cnt = self.name2val[key], self.name2cnt[key] 352 | self.name2val[key] = oldval * cnt / (cnt + 1) + val / (cnt + 1) 353 | self.name2cnt[key] = cnt + 1 354 | 355 | def dumpkvs(self): 356 | if self.comm is None: 357 | d = self.name2val 358 | else: 359 | d = mpi_weighted_mean( 360 | self.comm, 361 | { 362 | name: (val, self.name2cnt.get(name, 1)) 363 | for (name, val) in self.name2val.items() 364 | }, 365 | ) 366 | if self.comm.rank != 0: 367 | d["dummy"] = 1 # so we don't get a warning about empty dict 368 | out = d.copy() # Return the dict for unit testing purposes 369 | for fmt in self.output_formats: 370 | if isinstance(fmt, KVWriter): 371 | fmt.writekvs(d) 372 | self.name2val.clear() 373 | self.name2cnt.clear() 374 | return out 375 | 376 | def log(self, *args, level=INFO): 377 | if self.level <= level: 378 | self._do_log(args) 379 | 380 | # Configuration 381 | # ---------------------------------------- 382 | def set_level(self, level): 383 | self.level = level 384 | 385 | def set_comm(self, comm): 386 | self.comm = comm 387 | 388 | def get_dir(self): 389 | return self.dir 390 | 391 | def close(self): 392 | for fmt in self.output_formats: 393 | fmt.close() 394 | 395 | # Misc 396 | # ---------------------------------------- 397 | def _do_log(self, args): 398 | for fmt in self.output_formats: 399 | if isinstance(fmt, SeqWriter): 400 | fmt.writeseq(map(str, args)) 401 | 402 | 403 | def get_rank_without_mpi_import(): 404 | # check environment variables here instead of importing mpi4py 405 | # to avoid calling MPI_Init() when this module is imported 406 | for varname in ["PMI_RANK", "OMPI_COMM_WORLD_RANK"]: 407 | if varname in os.environ: 408 | return int(os.environ[varname]) 409 | return 0 410 | 411 | 412 | def mpi_weighted_mean(comm, local_name2valcount): 413 | """ 414 | Copied from: https://github.com/openai/baselines/blob/ea25b9e8b234e6ee1bca43083f8f3cf974143998/baselines/common/mpi_util.py#L110 415 | Perform a weighted average over dicts that are each on a different node 416 | Input: local_name2valcount: dict mapping key -> (value, count) 417 | Returns: key -> mean 418 | """ 419 | all_name2valcount = comm.gather(local_name2valcount) 420 | if comm.rank == 0: 421 | name2sum = defaultdict(float) 422 | name2count = defaultdict(float) 423 | for n2vc in all_name2valcount: 424 | for (name, (val, count)) in n2vc.items(): 425 | try: 426 | val = float(val) 427 | except ValueError: 428 | if comm.rank == 0: 429 | warnings.warn( 430 | "WARNING: tried to compute mean on non-float {}={}".format( 431 | name, val 432 | ) 433 | ) 434 | else: 435 | name2sum[name] += val * count 436 | name2count[name] += count 437 | return {name: name2sum[name] / name2count[name] for name in name2sum} 438 | else: 439 | return {} 440 | 441 | 442 | def configure(dir=None, format_strs=None, comm=None, log_suffix=""): 443 | """ 444 | If comm is provided, average all numerical stats across that comm 445 | """ 446 | if dir is None: 447 | dir = os.getenv("OPENAI_LOGDIR") 448 | if dir is None: 449 | dir = osp.join( 450 | tempfile.gettempdir(), 451 | datetime.datetime.now().strftime("openai-%Y-%m-%d-%H-%M-%S-%f"), 452 | ) 453 | assert isinstance(dir, str) 454 | dir = os.path.expanduser(dir) 455 | os.makedirs(os.path.expanduser(dir), exist_ok=True) 456 | 457 | rank = get_rank_without_mpi_import() 458 | if rank > 0: 459 | log_suffix = log_suffix + "-rank%03i" % rank 460 | 461 | if format_strs is None: 462 | if rank == 0: 463 | format_strs = os.getenv("OPENAI_LOG_FORMAT", "stdout,log,csv").split(",") 464 | else: 465 | format_strs = os.getenv("OPENAI_LOG_FORMAT_MPI", "log").split(",") 466 | format_strs = filter(None, format_strs) 467 | output_formats = [make_output_format(f, dir, log_suffix) for f in format_strs] 468 | 469 | Logger.CURRENT = Logger(dir=dir, output_formats=output_formats, comm=comm) 470 | if output_formats: 471 | log("Logging to %s" % dir) 472 | 473 | 474 | def _configure_default_logger(): 475 | configure() 476 | Logger.DEFAULT = Logger.CURRENT 477 | 478 | 479 | def reset(): 480 | if Logger.CURRENT is not Logger.DEFAULT: 481 | Logger.CURRENT.close() 482 | Logger.CURRENT = Logger.DEFAULT 483 | log("Reset logger") 484 | 485 | 486 | @contextmanager 487 | def scoped_configure(dir=None, format_strs=None, comm=None): 488 | prevlogger = Logger.CURRENT 489 | configure(dir=dir, format_strs=format_strs, comm=comm) 490 | try: 491 | yield 492 | finally: 493 | Logger.CURRENT.close() 494 | Logger.CURRENT = prevlogger 495 | 496 | -------------------------------------------------------------------------------- /guided_diffusion/losses.py: -------------------------------------------------------------------------------- 1 | """ 2 | Helpers for various likelihood-based losses. These are ported from the original 3 | Ho et al. diffusion models codebase: 4 | https://github.com/hojonathanho/diffusion/blob/1e0dceb3b3495bbe19116a5e1b3596cd0706c543/diffusion_tf/utils.py 5 | """ 6 | 7 | import numpy as np 8 | 9 | import torch as th 10 | 11 | 12 | def normal_kl(mean1, logvar1, mean2, logvar2): 13 | """ 14 | Compute the KL divergence between two gaussians. 15 | 16 | Shapes are automatically broadcasted, so batches can be compared to 17 | scalars, among other use cases. 18 | """ 19 | tensor = None 20 | for obj in (mean1, logvar1, mean2, logvar2): 21 | if isinstance(obj, th.Tensor): 22 | tensor = obj 23 | break 24 | assert tensor is not None, "at least one argument must be a Tensor" 25 | 26 | # Force variances to be Tensors. Broadcasting helps convert scalars to 27 | # Tensors, but it does not work for th.exp(). 28 | logvar1, logvar2 = [ 29 | x if isinstance(x, th.Tensor) else th.tensor(x).to(tensor) 30 | for x in (logvar1, logvar2) 31 | ] 32 | 33 | return 0.5 * ( 34 | -1.0 35 | + logvar2 36 | - logvar1 37 | + th.exp(logvar1 - logvar2) 38 | + ((mean1 - mean2) ** 2) * th.exp(-logvar2) 39 | ) 40 | 41 | 42 | def approx_standard_normal_cdf(x): 43 | """ 44 | A fast approximation of the cumulative distribution function of the 45 | standard normal. 46 | """ 47 | return 0.5 * (1.0 + th.tanh(np.sqrt(2.0 / np.pi) * (x + 0.044715 * th.pow(x, 3)))) 48 | 49 | 50 | def discretized_gaussian_log_likelihood(x, *, means, log_scales): 51 | """ 52 | Compute the log-likelihood of a Gaussian distribution discretizing to a 53 | given image. 54 | 55 | :param x: the target images. It is assumed that this was uint8 values, 56 | rescaled to the range [-1, 1]. 57 | :param means: the Gaussian mean Tensor. 58 | :param log_scales: the Gaussian log stddev Tensor. 59 | :return: a tensor like x of log probabilities (in nats). 60 | """ 61 | assert x.shape == means.shape == log_scales.shape 62 | centered_x = x - means 63 | inv_stdv = th.exp(-log_scales) 64 | plus_in = inv_stdv * (centered_x + 1.0 / 255.0) 65 | cdf_plus = approx_standard_normal_cdf(plus_in) 66 | min_in = inv_stdv * (centered_x - 1.0 / 255.0) 67 | cdf_min = approx_standard_normal_cdf(min_in) 68 | log_cdf_plus = th.log(cdf_plus.clamp(min=1e-12)) 69 | log_one_minus_cdf_min = th.log((1.0 - cdf_min).clamp(min=1e-12)) 70 | cdf_delta = cdf_plus - cdf_min 71 | log_probs = th.where( 72 | x < -0.999, 73 | log_cdf_plus, 74 | th.where(x > 0.999, log_one_minus_cdf_min, th.log(cdf_delta.clamp(min=1e-12))), 75 | ) 76 | assert log_probs.shape == x.shape 77 | return log_probs 78 | -------------------------------------------------------------------------------- /guided_diffusion/nn.py: -------------------------------------------------------------------------------- 1 | """ 2 | Various utilities for neural networks. 3 | """ 4 | 5 | import math 6 | 7 | import torch as th 8 | import torch.nn as nn 9 | 10 | 11 | # PyTorch 1.7 has SiLU, but we support PyTorch 1.5. 12 | class SiLU(nn.Module): 13 | def forward(self, x): 14 | return x * th.sigmoid(x) 15 | 16 | 17 | class GroupNorm32(nn.GroupNorm): 18 | def forward(self, x): 19 | return super().forward(x.float()).type(x.dtype) 20 | 21 | 22 | def conv_nd(dims, *args, **kwargs): 23 | """ 24 | Create a 1D, 2D, or 3D convolution module. 25 | """ 26 | if dims == 1: 27 | return nn.Conv1d(*args, **kwargs) 28 | elif dims == 2: 29 | return nn.Conv2d(*args, **kwargs) 30 | elif dims == 3: 31 | return nn.Conv3d(*args, **kwargs) 32 | raise ValueError(f"unsupported dimensions: {dims}") 33 | 34 | 35 | def linear(*args, **kwargs): 36 | """ 37 | Create a linear module. 38 | """ 39 | return nn.Linear(*args, **kwargs) 40 | 41 | 42 | def avg_pool_nd(dims, *args, **kwargs): 43 | """ 44 | Create a 1D, 2D, or 3D average pooling module. 45 | """ 46 | if dims == 1: 47 | return nn.AvgPool1d(*args, **kwargs) 48 | elif dims == 2: 49 | return nn.AvgPool2d(*args, **kwargs) 50 | elif dims == 3: 51 | return nn.AvgPool3d(*args, **kwargs) 52 | raise ValueError(f"unsupported dimensions: {dims}") 53 | 54 | 55 | def update_ema(target_params, source_params, rate=0.99): 56 | """ 57 | Update target parameters to be closer to those of source parameters using 58 | an exponential moving average. 59 | 60 | :param target_params: the target parameter sequence. 61 | :param source_params: the source parameter sequence. 62 | :param rate: the EMA rate (closer to 1 means slower). 63 | """ 64 | for targ, src in zip(target_params, source_params): 65 | targ.detach().mul_(rate).add_(src, alpha=1 - rate) 66 | 67 | 68 | def zero_module(module): 69 | """ 70 | Zero out the parameters of a module and return it. 71 | """ 72 | for p in module.parameters(): 73 | p.detach().zero_() 74 | return module 75 | 76 | 77 | def scale_module(module, scale): 78 | """ 79 | Scale the parameters of a module and return it. 80 | """ 81 | for p in module.parameters(): 82 | p.detach().mul_(scale) 83 | return module 84 | 85 | 86 | def mean_flat(tensor): 87 | """ 88 | Take the mean over all non-batch dimensions. 89 | """ 90 | return tensor.mean(dim=list(range(1, len(tensor.shape)))) 91 | 92 | 93 | def normalization(channels): 94 | """ 95 | Make a standard normalization layer. 96 | 97 | :param channels: number of input channels. 98 | :return: an nn.Module for normalization. 99 | """ 100 | return GroupNorm32(32, channels) 101 | 102 | 103 | def timestep_embedding(timesteps, dim, max_period=10000): 104 | """ 105 | Create sinusoidal timestep embeddings. 106 | 107 | :param timesteps: a 1-D Tensor of N indices, one per batch element. 108 | These may be fractional. 109 | :param dim: the dimension of the output. 110 | :param max_period: controls the minimum frequency of the embeddings. 111 | :return: an [N x dim] Tensor of positional embeddings. 112 | """ 113 | half = dim // 2 114 | freqs = th.exp( 115 | -math.log(max_period) * th.arange(start=0, end=half, dtype=th.float32) / half 116 | ).to(device=timesteps.device) 117 | args = timesteps[:, None].float() * freqs[None] 118 | embedding = th.cat([th.cos(args), th.sin(args)], dim=-1) 119 | if dim % 2: 120 | embedding = th.cat([embedding, th.zeros_like(embedding[:, :1])], dim=-1) 121 | return embedding 122 | 123 | 124 | def checkpoint(func, inputs, params, flag): 125 | """ 126 | Evaluate a function without caching intermediate activations, allowing for 127 | reduced memory at the expense of extra compute in the backward pass. 128 | 129 | :param func: the function to evaluate. 130 | :param inputs: the argument sequence to pass to `func`. 131 | :param params: a sequence of parameters `func` depends on but does not 132 | explicitly take as arguments. 133 | :param flag: if False, disable gradient checkpointing. 134 | """ 135 | if flag: 136 | args = tuple(inputs) + tuple(params) 137 | return CheckpointFunction.apply(func, len(inputs), *args) 138 | else: 139 | return func(*inputs) 140 | 141 | 142 | class CheckpointFunction(th.autograd.Function): 143 | @staticmethod 144 | def forward(ctx, run_function, length, *args): 145 | ctx.run_function = run_function 146 | ctx.input_tensors = list(args[:length]) 147 | ctx.input_params = list(args[length:]) 148 | with th.no_grad(): 149 | output_tensors = ctx.run_function(*ctx.input_tensors) 150 | return output_tensors 151 | 152 | @staticmethod 153 | def backward(ctx, *output_grads): 154 | ctx.input_tensors = [x.detach().requires_grad_(True) for x in ctx.input_tensors] 155 | with th.enable_grad(): 156 | # Fixes a bug where the first op in run_function modifies the 157 | # Tensor storage in place, which is not allowed for detach()'d 158 | # Tensors. 159 | shallow_copies = [x.view_as(x) for x in ctx.input_tensors] 160 | output_tensors = ctx.run_function(*shallow_copies) 161 | input_grads = th.autograd.grad( 162 | output_tensors, 163 | ctx.input_tensors + ctx.input_params, 164 | output_grads, 165 | allow_unused=True, 166 | ) 167 | del ctx.input_tensors 168 | del ctx.input_params 169 | del output_tensors 170 | return (None, None) + input_grads 171 | -------------------------------------------------------------------------------- /guided_diffusion/resample.py: -------------------------------------------------------------------------------- 1 | from abc import ABC, abstractmethod 2 | 3 | import numpy as np 4 | import torch as th 5 | import torch.distributed as dist 6 | 7 | 8 | def create_named_schedule_sampler(name, diffusion): 9 | """ 10 | Create a ScheduleSampler from a library of pre-defined samplers. 11 | 12 | :param name: the name of the sampler. 13 | :param diffusion: the diffusion object to sample for. 14 | """ 15 | if name == "uniform": 16 | return UniformSampler(diffusion) 17 | elif name == "loss-second-moment": 18 | return LossSecondMomentResampler(diffusion) 19 | else: 20 | raise NotImplementedError(f"unknown schedule sampler: {name}") 21 | 22 | 23 | class ScheduleSampler(ABC): 24 | """ 25 | A distribution over timesteps in the diffusion process, intended to reduce 26 | variance of the objective. 27 | 28 | By default, samplers perform unbiased importance sampling, in which the 29 | objective's mean is unchanged. 30 | However, subclasses may override sample() to change how the resampled 31 | terms are reweighted, allowing for actual changes in the objective. 32 | """ 33 | 34 | @abstractmethod 35 | def weights(self): 36 | """ 37 | Get a numpy array of weights, one per diffusion step. 38 | 39 | The weights needn't be normalized, but must be positive. 40 | """ 41 | 42 | def sample(self, batch_size, device): 43 | """ 44 | Importance-sample timesteps for a batch. 45 | 46 | :param batch_size: the number of timesteps. 47 | :param device: the torch device to save to. 48 | :return: a tuple (timesteps, weights): 49 | - timesteps: a tensor of timestep indices. 50 | - weights: a tensor of weights to scale the resulting losses. 51 | """ 52 | w = self.weights() 53 | p = w / np.sum(w) 54 | indices_np = np.random.choice(len(p), size=(batch_size,), p=p) 55 | indices = th.from_numpy(indices_np).long().to(device) 56 | weights_np = 1 / (len(p) * p[indices_np]) 57 | weights = th.from_numpy(weights_np).float().to(device) 58 | return indices, weights 59 | 60 | 61 | class UniformSampler(ScheduleSampler): 62 | def __init__(self, diffusion): 63 | self.diffusion = diffusion 64 | self._weights = np.ones([diffusion.num_timesteps]) 65 | 66 | def weights(self): 67 | return self._weights 68 | 69 | 70 | class LossAwareSampler(ScheduleSampler): 71 | def update_with_local_losses(self, local_ts, local_losses): 72 | """ 73 | Update the reweighting using losses from a model. 74 | 75 | Call this method from each rank with a batch of timesteps and the 76 | corresponding losses for each of those timesteps. 77 | This method will perform synchronization to make sure all of the ranks 78 | maintain the exact same reweighting. 79 | 80 | :param local_ts: an integer Tensor of timesteps. 81 | :param local_losses: a 1D Tensor of losses. 82 | """ 83 | batch_sizes = [ 84 | th.tensor([0], dtype=th.int32, device=local_ts.device) 85 | for _ in range(dist.get_world_size()) 86 | ] 87 | dist.all_gather( 88 | batch_sizes, 89 | th.tensor([len(local_ts)], dtype=th.int32, device=local_ts.device), 90 | ) 91 | 92 | # Pad all_gather batches to be the maximum batch size. 93 | batch_sizes = [x.item() for x in batch_sizes] 94 | max_bs = max(batch_sizes) 95 | 96 | timestep_batches = [th.zeros(max_bs).to(local_ts) for bs in batch_sizes] 97 | loss_batches = [th.zeros(max_bs).to(local_losses) for bs in batch_sizes] 98 | dist.all_gather(timestep_batches, local_ts) 99 | dist.all_gather(loss_batches, local_losses) 100 | timesteps = [ 101 | x.item() for y, bs in zip(timestep_batches, batch_sizes) for x in y[:bs] 102 | ] 103 | losses = [x.item() for y, bs in zip(loss_batches, batch_sizes) for x in y[:bs]] 104 | self.update_with_all_losses(timesteps, losses) 105 | 106 | @abstractmethod 107 | def update_with_all_losses(self, ts, losses): 108 | """ 109 | Update the reweighting using losses from a model. 110 | 111 | Sub-classes should override this method to update the reweighting 112 | using losses from the model. 113 | 114 | This method directly updates the reweighting without synchronizing 115 | between workers. It is called by update_with_local_losses from all 116 | ranks with identical arguments. Thus, it should have deterministic 117 | behavior to maintain state across workers. 118 | 119 | :param ts: a list of int timesteps. 120 | :param losses: a list of float losses, one per timestep. 121 | """ 122 | 123 | 124 | class LossSecondMomentResampler(LossAwareSampler): 125 | def __init__(self, diffusion, history_per_term=10, uniform_prob=0.001): 126 | self.diffusion = diffusion 127 | self.history_per_term = history_per_term 128 | self.uniform_prob = uniform_prob 129 | self._loss_history = np.zeros( 130 | [diffusion.num_timesteps, history_per_term], dtype=np.float64 131 | ) 132 | self._loss_counts = np.zeros([diffusion.num_timesteps], dtype=np.int) 133 | 134 | def weights(self): 135 | if not self._warmed_up(): 136 | return np.ones([self.diffusion.num_timesteps], dtype=np.float64) 137 | weights = np.sqrt(np.mean(self._loss_history ** 2, axis=-1)) 138 | weights /= np.sum(weights) 139 | weights *= 1 - self.uniform_prob 140 | weights += self.uniform_prob / len(weights) 141 | return weights 142 | 143 | def update_with_all_losses(self, ts, losses): 144 | for t, loss in zip(ts, losses): 145 | if self._loss_counts[t] == self.history_per_term: 146 | # Shift out the oldest loss term. 147 | self._loss_history[t, :-1] = self._loss_history[t, 1:] 148 | self._loss_history[t, -1] = loss 149 | else: 150 | self._loss_history[t, self._loss_counts[t]] = loss 151 | self._loss_counts[t] += 1 152 | 153 | def _warmed_up(self): 154 | return (self._loss_counts == self.history_per_term).all() 155 | -------------------------------------------------------------------------------- /guided_diffusion/respace.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import torch as th 3 | 4 | from .gaussian_diffusion import GaussianDiffusion 5 | 6 | 7 | def space_timesteps(num_timesteps, section_counts): 8 | """ 9 | Create a list of timesteps to use from an original diffusion process, 10 | given the number of timesteps we want to take from equally-sized portions 11 | of the original process. 12 | 13 | For example, if there's 300 timesteps and the section counts are [10,15,20] 14 | then the first 100 timesteps are strided to be 10 timesteps, the second 100 15 | are strided to be 15 timesteps, and the final 100 are strided to be 20. 16 | 17 | If the stride is a string starting with "ddim", then the fixed striding 18 | from the DDIM paper is used, and only one section is allowed. 19 | 20 | :param num_timesteps: the number of diffusion steps in the original 21 | process to divide up. 22 | :param section_counts: either a list of numbers, or a string containing 23 | comma-separated numbers, indicating the step count 24 | per section. As a special case, use "ddimN" where N 25 | is a number of steps to use the striding from the 26 | DDIM paper. 27 | :return: a set of diffusion steps from the original process to use. 28 | """ 29 | if isinstance(section_counts, str): 30 | if section_counts.startswith("ddim"): 31 | desired_count = int(section_counts[len("ddim") :]) 32 | for i in range(1, num_timesteps): 33 | if len(range(0, num_timesteps, i)) == desired_count: 34 | return set(range(0, num_timesteps, i)) 35 | raise ValueError( 36 | f"cannot create exactly {num_timesteps} steps with an integer stride" 37 | ) 38 | section_counts = [int(x) for x in section_counts.split(",")] 39 | size_per = num_timesteps // len(section_counts) 40 | extra = num_timesteps % len(section_counts) 41 | start_idx = 0 42 | all_steps = [] 43 | for i, section_count in enumerate(section_counts): 44 | size = size_per + (1 if i < extra else 0) 45 | if size < section_count: 46 | raise ValueError( 47 | f"cannot divide section of {size} steps into {section_count}" 48 | ) 49 | if section_count <= 1: 50 | frac_stride = 1 51 | else: 52 | frac_stride = (size - 1) / (section_count - 1) 53 | cur_idx = 0.0 54 | taken_steps = [] 55 | for _ in range(section_count): 56 | taken_steps.append(start_idx + round(cur_idx)) 57 | cur_idx += frac_stride 58 | all_steps += taken_steps 59 | start_idx += size 60 | return set(all_steps) 61 | 62 | 63 | class SpacedDiffusion(GaussianDiffusion): 64 | """ 65 | A diffusion process which can skip steps in a base diffusion process. 66 | 67 | :param use_timesteps: a collection (sequence or set) of timesteps from the 68 | original diffusion process to retain. 69 | :param kwargs: the kwargs to create the base diffusion process. 70 | """ 71 | 72 | def __init__(self, use_timesteps, **kwargs): 73 | self.use_timesteps = set(use_timesteps) 74 | self.timestep_map = [] 75 | self.original_num_steps = len(kwargs["betas"]) 76 | 77 | base_diffusion = GaussianDiffusion(**kwargs) # pylint: disable=missing-kwoa 78 | last_alpha_cumprod = 1.0 79 | new_betas = [] 80 | for i, alpha_cumprod in enumerate(base_diffusion.alphas_cumprod): 81 | if i in self.use_timesteps: 82 | new_betas.append(1 - alpha_cumprod / last_alpha_cumprod) 83 | last_alpha_cumprod = alpha_cumprod 84 | self.timestep_map.append(i) 85 | kwargs["betas"] = np.array(new_betas) 86 | super().__init__(**kwargs) 87 | 88 | def p_mean_variance( 89 | self, model, *args, **kwargs 90 | ): # pylint: disable=signature-differs 91 | return super().p_mean_variance(self._wrap_model(model), *args, **kwargs) 92 | 93 | def training_losses( 94 | self, model, *args, **kwargs 95 | ): # pylint: disable=signature-differs 96 | return super().training_losses(self._wrap_model(model), *args, **kwargs) 97 | 98 | def condition_mean(self, cond_fn, *args, **kwargs): 99 | return super().condition_mean(self._wrap_model(cond_fn), *args, **kwargs) 100 | 101 | def condition_score(self, cond_fn, *args, **kwargs): 102 | return super().condition_score(self._wrap_model(cond_fn), *args, **kwargs) 103 | 104 | def _wrap_model(self, model): 105 | if isinstance(model, _WrappedModel): 106 | return model 107 | return _WrappedModel( 108 | model, self.timestep_map, self.rescale_timesteps, self.original_num_steps 109 | ) 110 | 111 | def _scale_timesteps(self, t): 112 | # Scaling is done by the wrapped model. 113 | return t 114 | 115 | 116 | class _WrappedModel: 117 | def __init__(self, model, timestep_map, rescale_timesteps, original_num_steps): 118 | self.model = model 119 | self.timestep_map = timestep_map 120 | self.rescale_timesteps = rescale_timesteps 121 | self.original_num_steps = original_num_steps 122 | 123 | def __call__(self, x, ts, **kwargs): 124 | map_tensor = th.tensor(self.timestep_map, device=ts.device, dtype=ts.dtype) 125 | new_ts = map_tensor[ts] 126 | if self.rescale_timesteps: 127 | new_ts = new_ts.float() * (1000.0 / self.original_num_steps) 128 | return self.model(x, new_ts, **kwargs) 129 | -------------------------------------------------------------------------------- /guided_diffusion/script_util.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import inspect 3 | 4 | from . import gaussian_diffusion as gd 5 | from .respace import SpacedDiffusion, space_timesteps 6 | from .unet import SuperResModel, UNetModel, EncoderUNetModel, SegModelv2, SegModelv2_6c, SegModelv3_6c, SuperResModel_noatt, SegModelv2_3d_noatt, SegModel_3d_noatt_midcat 7 | 8 | NUM_CLASSES = 1000 9 | 10 | 11 | def diffusion_defaults(): 12 | """ 13 | Defaults for image and classifier training. 14 | """ 15 | return dict( 16 | learn_sigma=False, 17 | diffusion_steps=1000, 18 | noise_schedule="linear", 19 | timestep_respacing="", 20 | use_kl=False, 21 | predict_xstart=False, 22 | rescale_timesteps=False, 23 | rescale_learned_sigmas=False, 24 | ) 25 | 26 | 27 | def classifier_defaults(): 28 | """ 29 | Defaults for classifier models. 30 | """ 31 | return dict( 32 | image_size=64, 33 | classifier_use_fp16=False, 34 | classifier_width=128, 35 | classifier_depth=2, 36 | classifier_attention_resolutions="32,16,8", # 16 37 | classifier_use_scale_shift_norm=True, # False 38 | classifier_resblock_updown=True, # False 39 | classifier_pool="attention", 40 | ) 41 | 42 | 43 | def model_and_diffusion_defaults(): 44 | """ 45 | Defaults for image training. 46 | """ 47 | res = dict( 48 | image_size=64, 49 | num_channels=128, 50 | num_res_blocks=2, 51 | num_heads=4, 52 | num_heads_upsample=-1, 53 | num_head_channels=-1, 54 | attention_resolutions="16,8", 55 | channel_mult="", 56 | dropout=0.0, 57 | class_cond=False, 58 | use_checkpoint=False, 59 | use_scale_shift_norm=True, 60 | resblock_updown=False, 61 | use_fp16=False, 62 | use_new_attention_order=False, 63 | ) 64 | res.update(diffusion_defaults()) 65 | return res 66 | 67 | 68 | def classifier_and_diffusion_defaults(): 69 | res = classifier_defaults() 70 | res.update(diffusion_defaults()) 71 | return res 72 | 73 | 74 | def create_model_and_diffusion( 75 | image_size, 76 | class_cond, 77 | learn_sigma, 78 | num_channels, 79 | num_res_blocks, 80 | channel_mult, 81 | num_heads, 82 | num_head_channels, 83 | num_heads_upsample, 84 | attention_resolutions, 85 | dropout, 86 | diffusion_steps, 87 | noise_schedule, 88 | timestep_respacing, 89 | use_kl, 90 | predict_xstart, 91 | rescale_timesteps, 92 | rescale_learned_sigmas, 93 | use_checkpoint, 94 | use_scale_shift_norm, 95 | resblock_updown, 96 | use_fp16, 97 | use_new_attention_order, 98 | ): 99 | model = create_model( 100 | image_size, 101 | num_channels, 102 | num_res_blocks, 103 | channel_mult=channel_mult, 104 | learn_sigma=learn_sigma, 105 | class_cond=class_cond, 106 | use_checkpoint=use_checkpoint, 107 | attention_resolutions=attention_resolutions, 108 | num_heads=num_heads, 109 | num_head_channels=num_head_channels, 110 | num_heads_upsample=num_heads_upsample, 111 | use_scale_shift_norm=use_scale_shift_norm, 112 | dropout=dropout, 113 | resblock_updown=resblock_updown, 114 | use_fp16=use_fp16, 115 | use_new_attention_order=use_new_attention_order, 116 | ) 117 | diffusion = create_gaussian_diffusion( 118 | steps=diffusion_steps, 119 | learn_sigma=learn_sigma, 120 | noise_schedule=noise_schedule, 121 | use_kl=use_kl, 122 | predict_xstart=predict_xstart, 123 | rescale_timesteps=rescale_timesteps, 124 | rescale_learned_sigmas=rescale_learned_sigmas, 125 | timestep_respacing=timestep_respacing, 126 | ) 127 | return model, diffusion 128 | 129 | 130 | def create_model( 131 | image_size, 132 | num_channels, 133 | num_res_blocks, 134 | channel_mult="", 135 | learn_sigma=False, 136 | class_cond=False, 137 | use_checkpoint=False, 138 | attention_resolutions="16", 139 | num_heads=1, 140 | num_head_channels=-1, 141 | num_heads_upsample=-1, 142 | use_scale_shift_norm=False, 143 | dropout=0, 144 | resblock_updown=False, 145 | use_fp16=False, 146 | use_new_attention_order=False, 147 | ): 148 | if channel_mult == "": 149 | if image_size == 512: 150 | channel_mult = (0.5, 1, 1, 2, 2, 4, 4) 151 | elif image_size == 256: 152 | channel_mult = (1, 1, 2, 2, 4, 4) 153 | elif image_size == 128: 154 | channel_mult = (1, 1, 2, 3, 4) 155 | elif image_size == 64: 156 | channel_mult = (1, 2, 3, 4) 157 | else: 158 | raise ValueError(f"unsupported image size: {image_size}") 159 | else: 160 | channel_mult = tuple(int(ch_mult) for ch_mult in channel_mult.split(",")) 161 | 162 | attention_ds = [] 163 | for res in attention_resolutions.split(","): 164 | attention_ds.append(image_size // int(res)) 165 | 166 | return UNetModel( 167 | image_size=image_size, 168 | in_channels=3, 169 | model_channels=num_channels, 170 | out_channels=(3 if not learn_sigma else 6), 171 | num_res_blocks=num_res_blocks, 172 | attention_resolutions=tuple(attention_ds), 173 | dropout=dropout, 174 | channel_mult=channel_mult, 175 | num_classes=(NUM_CLASSES if class_cond else None), 176 | use_checkpoint=use_checkpoint, 177 | use_fp16=use_fp16, 178 | num_heads=num_heads, 179 | num_head_channels=num_head_channels, 180 | num_heads_upsample=num_heads_upsample, 181 | use_scale_shift_norm=use_scale_shift_norm, 182 | resblock_updown=resblock_updown, 183 | use_new_attention_order=use_new_attention_order, 184 | ) 185 | 186 | 187 | def create_classifier_and_diffusion( 188 | image_size, 189 | classifier_use_fp16, 190 | classifier_width, 191 | classifier_depth, 192 | classifier_attention_resolutions, 193 | classifier_use_scale_shift_norm, 194 | classifier_resblock_updown, 195 | classifier_pool, 196 | learn_sigma, 197 | diffusion_steps, 198 | noise_schedule, 199 | timestep_respacing, 200 | use_kl, 201 | predict_xstart, 202 | rescale_timesteps, 203 | rescale_learned_sigmas, 204 | ): 205 | classifier = create_classifier( 206 | image_size, 207 | classifier_use_fp16, 208 | classifier_width, 209 | classifier_depth, 210 | classifier_attention_resolutions, 211 | classifier_use_scale_shift_norm, 212 | classifier_resblock_updown, 213 | classifier_pool, 214 | ) 215 | diffusion = create_gaussian_diffusion( 216 | steps=diffusion_steps, 217 | learn_sigma=learn_sigma, 218 | noise_schedule=noise_schedule, 219 | use_kl=use_kl, 220 | predict_xstart=predict_xstart, 221 | rescale_timesteps=rescale_timesteps, 222 | rescale_learned_sigmas=rescale_learned_sigmas, 223 | timestep_respacing=timestep_respacing, 224 | ) 225 | return classifier, diffusion 226 | 227 | 228 | def create_classifier( 229 | image_size, 230 | classifier_use_fp16, 231 | classifier_width, 232 | classifier_depth, 233 | classifier_attention_resolutions, 234 | classifier_use_scale_shift_norm, 235 | classifier_resblock_updown, 236 | classifier_pool, 237 | ): 238 | if image_size == 512: 239 | channel_mult = (0.5, 1, 1, 2, 2, 4, 4) 240 | elif image_size == 256: 241 | channel_mult = (1, 1, 2, 2, 4, 4) 242 | elif image_size == 128: 243 | channel_mult = (1, 1, 2, 3, 4) 244 | elif image_size == 64: 245 | channel_mult = (1, 2, 3, 4) 246 | else: 247 | raise ValueError(f"unsupported image size: {image_size}") 248 | 249 | attention_ds = [] 250 | for res in classifier_attention_resolutions.split(","): 251 | attention_ds.append(image_size // int(res)) 252 | 253 | return EncoderUNetModel( 254 | image_size=image_size, 255 | in_channels=3, 256 | model_channels=classifier_width, 257 | out_channels=1000, 258 | num_res_blocks=classifier_depth, 259 | attention_resolutions=tuple(attention_ds), 260 | channel_mult=channel_mult, 261 | use_fp16=classifier_use_fp16, 262 | num_head_channels=64, 263 | use_scale_shift_norm=classifier_use_scale_shift_norm, 264 | resblock_updown=classifier_resblock_updown, 265 | pool=classifier_pool, 266 | ) 267 | 268 | 269 | def sr_model_and_diffusion_defaults(): 270 | res = model_and_diffusion_defaults() 271 | res["large_size"] = 256 272 | res["small_size"] = 64 273 | arg_names = inspect.getfullargspec(sr_create_model_and_diffusion)[0] 274 | for k in res.copy().keys(): 275 | if k not in arg_names: 276 | del res[k] 277 | return res 278 | 279 | 280 | def sr_create_model_and_diffusion( 281 | large_size, 282 | small_size, 283 | class_cond, 284 | learn_sigma, 285 | num_channels, 286 | num_res_blocks, 287 | num_heads, 288 | num_head_channels, 289 | num_heads_upsample, 290 | attention_resolutions, 291 | dropout, 292 | diffusion_steps, 293 | noise_schedule, 294 | timestep_respacing, 295 | use_kl, 296 | predict_xstart, 297 | rescale_timesteps, 298 | rescale_learned_sigmas, 299 | use_checkpoint, 300 | use_scale_shift_norm, 301 | resblock_updown, 302 | use_fp16, 303 | ): 304 | model = sr_create_model( 305 | large_size, 306 | small_size, 307 | num_channels, 308 | num_res_blocks, 309 | learn_sigma=learn_sigma, 310 | class_cond=class_cond, 311 | use_checkpoint=use_checkpoint, 312 | attention_resolutions=attention_resolutions, 313 | num_heads=num_heads, 314 | num_head_channels=num_head_channels, 315 | num_heads_upsample=num_heads_upsample, 316 | use_scale_shift_norm=use_scale_shift_norm, 317 | dropout=dropout, 318 | resblock_updown=resblock_updown, 319 | use_fp16=use_fp16, 320 | ) 321 | diffusion = create_gaussian_diffusion( 322 | steps=diffusion_steps, 323 | learn_sigma=learn_sigma, 324 | noise_schedule=noise_schedule, 325 | use_kl=use_kl, 326 | predict_xstart=predict_xstart, 327 | rescale_timesteps=rescale_timesteps, 328 | rescale_learned_sigmas=rescale_learned_sigmas, 329 | timestep_respacing=timestep_respacing, 330 | ) 331 | return model, diffusion 332 | 333 | 334 | def sr_create_model( 335 | large_size, 336 | small_size, 337 | num_channels, 338 | num_res_blocks, 339 | learn_sigma, 340 | class_cond, 341 | use_checkpoint, 342 | attention_resolutions, 343 | num_heads, 344 | num_head_channels, 345 | num_heads_upsample, 346 | use_scale_shift_norm, 347 | dropout, 348 | resblock_updown, 349 | use_fp16, 350 | ): 351 | _ = small_size # hack to prevent unused variable 352 | 353 | if large_size == 512: 354 | channel_mult = (1, 1, 2, 2, 4, 4) 355 | elif large_size == 256: 356 | channel_mult = (1, 1, 2, 2, 4, 4) 357 | elif large_size == 64: 358 | channel_mult = (1, 2, 3, 4) 359 | else: 360 | channel_mult = (1, 1, 2, 3, 4) 361 | # raise ValueError(f"unsupported large size: {large_size}") 362 | 363 | attention_ds = [] 364 | for res in attention_resolutions.split(","): 365 | attention_ds.append(large_size // int(res)) 366 | 367 | # return SuperResModel( 368 | # image_size=large_size, 369 | # in_channels=3, 370 | # model_channels=num_channels, 371 | # out_channels=(3 if not learn_sigma else 6), 372 | # num_res_blocks=num_res_blocks, 373 | # attention_resolutions=tuple(attention_ds), 374 | # dropout=dropout, 375 | # channel_mult=channel_mult, 376 | # num_classes=(NUM_CLASSES if class_cond else None), 377 | # use_checkpoint=use_checkpoint, 378 | # num_heads=num_heads, 379 | # num_head_channels=num_head_channels, 380 | # num_heads_upsample=num_heads_upsample, 381 | # use_scale_shift_norm=use_scale_shift_norm, 382 | # resblock_updown=resblock_updown, 383 | # use_fp16=use_fp16, 384 | # ) 385 | 386 | 387 | # 3D input PET CT, input channel is 1.5*2, output channel is 1 or 2 388 | # return SuperResModel( 389 | # image_size=large_size, 390 | # in_channels=1.5, 391 | # model_channels=num_channels, 392 | # out_channels=(1 if not learn_sigma else 2), 393 | # num_res_blocks=num_res_blocks, 394 | # attention_resolutions=tuple(attention_ds), 395 | # dropout=dropout, 396 | # channel_mult=channel_mult, 397 | # dims=3, # 3D conv 398 | # num_classes=(NUM_CLASSES if class_cond else None), 399 | # use_checkpoint=use_checkpoint, 400 | # num_heads=num_heads, 401 | # num_head_channels=num_head_channels, 402 | # num_heads_upsample=num_heads_upsample, 403 | # use_scale_shift_norm=use_scale_shift_norm, 404 | # resblock_updown=resblock_updown, 405 | # use_fp16=use_fp16, 406 | # ) 407 | 408 | 409 | # # 3D input PET CT, input channel is 1.5*2, output channel is 1 or 2 410 | # return SuperResModel_noatt( 411 | # image_size=large_size, 412 | # in_channels=1.5, 413 | # model_channels=num_channels, 414 | # out_channels=(1 if not learn_sigma else 2), 415 | # num_res_blocks=num_res_blocks, 416 | # attention_resolutions=tuple(attention_ds), 417 | # dropout=dropout, 418 | # channel_mult=channel_mult, 419 | # dims=3, # 3D conv 420 | # num_classes=(NUM_CLASSES if class_cond else None), 421 | # use_checkpoint=use_checkpoint, 422 | # num_heads=num_heads, 423 | # num_head_channels=num_head_channels, 424 | # num_heads_upsample=num_heads_upsample, 425 | # use_scale_shift_norm=use_scale_shift_norm, 426 | # resblock_updown=resblock_updown, 427 | # use_fp16=use_fp16, 428 | # ) 429 | 430 | 431 | # 3D input PET or CT only, input channel is 1*2, output channel is 1 or 2 432 | return SuperResModel_noatt( 433 | image_size=large_size, 434 | in_channels=1, 435 | model_channels=num_channels, 436 | out_channels=(1 if not learn_sigma else 2), 437 | num_res_blocks=num_res_blocks, 438 | attention_resolutions=tuple(attention_ds), 439 | dropout=dropout, 440 | channel_mult=channel_mult, 441 | dims=3, # 3D conv 442 | num_classes=(NUM_CLASSES if class_cond else None), 443 | use_checkpoint=use_checkpoint, 444 | num_heads=num_heads, 445 | num_head_channels=num_head_channels, 446 | num_heads_upsample=num_heads_upsample, 447 | use_scale_shift_norm=use_scale_shift_norm, 448 | resblock_updown=resblock_updown, 449 | use_fp16=use_fp16, 450 | ) 451 | 452 | 453 | # 3D input PET CT, input channel is 1.5*2, output channel is 1 or 2 454 | # return SegModelv2_3d_noatt( 455 | # image_size=large_size, 456 | # in_channels=1.5, 457 | # model_channels=num_channels, 458 | # out_channels=(1 if not learn_sigma else 2), 459 | # num_res_blocks=num_res_blocks, 460 | # attention_resolutions=tuple(attention_ds), 461 | # dropout=dropout, 462 | # channel_mult=channel_mult, 463 | # dims=3, # 3D conv 464 | # num_classes=(NUM_CLASSES if class_cond else None), 465 | # use_checkpoint=use_checkpoint, 466 | # num_heads=num_heads, 467 | # num_head_channels=num_head_channels, 468 | # num_heads_upsample=num_heads_upsample, 469 | # use_scale_shift_norm=use_scale_shift_norm, 470 | # resblock_updown=resblock_updown, 471 | # use_fp16=use_fp16, 472 | # ) 473 | 474 | 475 | # 3D input PET CT, input channel is 1.5*2, output channel is 1 or 2 476 | # return SegModel_3d_noatt_midcat( 477 | # image_size=large_size, 478 | # in_channels=1.5, 479 | # model_channels=num_channels, 480 | # out_channels=(1 if not learn_sigma else 2), 481 | # num_res_blocks=num_res_blocks, 482 | # attention_resolutions=tuple(attention_ds), 483 | # dropout=dropout, 484 | # channel_mult=channel_mult, 485 | # dims=3, # 3D conv 486 | # num_classes=(NUM_CLASSES if class_cond else None), 487 | # use_checkpoint=use_checkpoint, 488 | # num_heads=num_heads, 489 | # num_head_channels=num_head_channels, 490 | # num_heads_upsample=num_heads_upsample, 491 | # use_scale_shift_norm=use_scale_shift_norm, 492 | # resblock_updown=resblock_updown, 493 | # use_fp16=use_fp16, 494 | # ) 495 | 496 | 497 | # input channel 4.5*2 (pet and ct and noise) 498 | # return SuperResModel( 499 | # image_size=large_size, 500 | # in_channels=4.5, 501 | # model_channels=num_channels, 502 | # out_channels=(3 if not learn_sigma else 6), 503 | # num_res_blocks=num_res_blocks, 504 | # attention_resolutions=tuple(attention_ds), 505 | # dropout=dropout, 506 | # channel_mult=channel_mult, 507 | # num_classes=(NUM_CLASSES if class_cond else None), 508 | # use_checkpoint=use_checkpoint, 509 | # num_heads=num_heads, 510 | # num_head_channels=num_head_channels, 511 | # num_heads_upsample=num_heads_upsample, 512 | # use_scale_shift_norm=use_scale_shift_norm, 513 | # resblock_updown=resblock_updown, 514 | # use_fp16=use_fp16, 515 | # ) 516 | 517 | # return SegModelv2( 518 | # image_size=large_size, 519 | # in_channels=3, 520 | # model_channels=num_channels, 521 | # out_channels=(3 if not learn_sigma else 6), 522 | # num_res_blocks=num_res_blocks, 523 | # attention_resolutions=tuple(attention_ds), 524 | # dropout=dropout, 525 | # channel_mult=channel_mult, 526 | # num_classes=(NUM_CLASSES if class_cond else None), 527 | # use_checkpoint=use_checkpoint, 528 | # num_heads=num_heads, 529 | # num_head_channels=num_head_channels, 530 | # num_heads_upsample=num_heads_upsample, 531 | # use_scale_shift_norm=use_scale_shift_norm, 532 | # resblock_updown=resblock_updown, 533 | # use_fp16=use_fp16, 534 | # ) 535 | 536 | # input channel 6 (pet and ct) 537 | # return SegModelv2_6c( 538 | # image_size=large_size, 539 | # in_channels=6, 540 | # model_channels=num_channels, 541 | # out_channels=(3 if not learn_sigma else 6), 542 | # num_res_blocks=num_res_blocks, 543 | # attention_resolutions=tuple(attention_ds), 544 | # dropout=dropout, 545 | # channel_mult=channel_mult, 546 | # num_classes=(NUM_CLASSES if class_cond else None), 547 | # use_checkpoint=use_checkpoint, 548 | # num_heads=num_heads, 549 | # num_head_channels=num_head_channels, 550 | # num_heads_upsample=num_heads_upsample, 551 | # use_scale_shift_norm=use_scale_shift_norm, 552 | # resblock_updown=resblock_updown, 553 | # use_fp16=use_fp16, 554 | # ) 555 | 556 | # input channel 6 (pet and ct) 557 | # use concat to fuse and a 1*1 conv to reduce channel 558 | # return SegModelv3_6c( 559 | # image_size=large_size, 560 | # in_channels=6, 561 | # model_channels=num_channels, 562 | # out_channels=(3 if not learn_sigma else 6), 563 | # num_res_blocks=num_res_blocks, 564 | # attention_resolutions=tuple(attention_ds), 565 | # dropout=dropout, 566 | # channel_mult=channel_mult, 567 | # num_classes=(NUM_CLASSES if class_cond else None), 568 | # use_checkpoint=use_checkpoint, 569 | # num_heads=num_heads, 570 | # num_head_channels=num_head_channels, 571 | # num_heads_upsample=num_heads_upsample, 572 | # use_scale_shift_norm=use_scale_shift_norm, 573 | # resblock_updown=resblock_updown, 574 | # use_fp16=use_fp16, 575 | # ) 576 | 577 | 578 | def create_gaussian_diffusion( 579 | *, 580 | steps=1000, 581 | learn_sigma=False, 582 | sigma_small=False, 583 | noise_schedule="linear", 584 | use_kl=False, 585 | predict_xstart=False, 586 | rescale_timesteps=False, 587 | rescale_learned_sigmas=False, 588 | timestep_respacing="", 589 | ): 590 | betas = gd.get_named_beta_schedule(noise_schedule, steps) 591 | if use_kl: 592 | loss_type = gd.LossType.RESCALED_KL 593 | elif rescale_learned_sigmas: 594 | loss_type = gd.LossType.RESCALED_MSE 595 | else: 596 | loss_type = gd.LossType.MSE 597 | if not timestep_respacing: 598 | timestep_respacing = [steps] 599 | return SpacedDiffusion( 600 | use_timesteps=space_timesteps(steps, timestep_respacing), 601 | betas=betas, 602 | model_mean_type=( 603 | gd.ModelMeanType.EPSILON if not predict_xstart else gd.ModelMeanType.START_X 604 | ), 605 | model_var_type=( 606 | ( 607 | gd.ModelVarType.FIXED_LARGE 608 | if not sigma_small 609 | else gd.ModelVarType.FIXED_SMALL 610 | ) 611 | if not learn_sigma 612 | else gd.ModelVarType.LEARNED_RANGE 613 | ), 614 | loss_type=loss_type, 615 | rescale_timesteps=rescale_timesteps, 616 | ) 617 | 618 | 619 | def add_dict_to_argparser(parser, default_dict): 620 | for k, v in default_dict.items(): 621 | v_type = type(v) 622 | if v is None: 623 | v_type = str 624 | elif isinstance(v, bool): 625 | v_type = str2bool 626 | parser.add_argument(f"--{k}", default=v, type=v_type) 627 | 628 | 629 | def args_to_dict(args, keys): 630 | return {k: getattr(args, k) for k in keys} 631 | 632 | 633 | def str2bool(v): 634 | """ 635 | https://stackoverflow.com/questions/15008758/parsing-boolean-values-with-argparse 636 | """ 637 | if isinstance(v, bool): 638 | return v 639 | if v.lower() in ("yes", "true", "t", "y", "1"): 640 | return True 641 | elif v.lower() in ("no", "false", "f", "n", "0"): 642 | return False 643 | else: 644 | raise argparse.ArgumentTypeError("boolean value expected") 645 | -------------------------------------------------------------------------------- /guided_diffusion/train_util.py: -------------------------------------------------------------------------------- 1 | import copy 2 | import functools 3 | import os 4 | 5 | import blobfile as bf 6 | import torch as th 7 | import torch.distributed as dist 8 | from torch.nn.parallel.distributed import DistributedDataParallel as DDP 9 | from torch.optim import AdamW 10 | 11 | from . import dist_util, logger 12 | from .fp16_util import MixedPrecisionTrainer 13 | from .nn import update_ema 14 | from .resample import LossAwareSampler, UniformSampler 15 | 16 | # For ImageNet experiments, this was a good default value. 17 | # We found that the lg_loss_scale quickly climbed to 18 | # 20-21 within the first ~1K steps of training. 19 | INITIAL_LOG_LOSS_SCALE = 20.0 20 | 21 | 22 | class TrainLoop: 23 | def __init__( 24 | self, 25 | *, 26 | model, 27 | diffusion, 28 | data, 29 | batch_size, 30 | microbatch, 31 | lr, 32 | ema_rate, 33 | log_interval, 34 | save_interval, 35 | resume_checkpoint, 36 | use_fp16=False, 37 | fp16_scale_growth=1e-3, 38 | schedule_sampler=None, 39 | weight_decay=0.0, 40 | lr_anneal_steps=0, 41 | ): 42 | self.model = model 43 | self.diffusion = diffusion 44 | self.data = data 45 | self.batch_size = batch_size 46 | self.microbatch = microbatch if microbatch > 0 else batch_size 47 | self.lr = lr 48 | self.ema_rate = ( 49 | [ema_rate] 50 | if isinstance(ema_rate, float) 51 | else [float(x) for x in ema_rate.split(",")] 52 | ) 53 | self.log_interval = log_interval 54 | self.save_interval = save_interval 55 | self.resume_checkpoint = resume_checkpoint 56 | self.use_fp16 = use_fp16 57 | self.fp16_scale_growth = fp16_scale_growth 58 | self.schedule_sampler = schedule_sampler or UniformSampler(diffusion) 59 | self.weight_decay = weight_decay 60 | self.lr_anneal_steps = lr_anneal_steps 61 | 62 | self.step = 0 63 | self.resume_step = 0 64 | self.global_batch = self.batch_size * dist.get_world_size() 65 | 66 | self.sync_cuda = th.cuda.is_available() 67 | 68 | self._load_and_sync_parameters() 69 | self.mp_trainer = MixedPrecisionTrainer( 70 | model=self.model, 71 | use_fp16=self.use_fp16, 72 | fp16_scale_growth=fp16_scale_growth, 73 | ) 74 | 75 | self.opt = AdamW( 76 | self.mp_trainer.master_params, lr=self.lr, weight_decay=self.weight_decay 77 | ) 78 | if self.resume_step: 79 | self._load_optimizer_state() 80 | # Model was resumed, either due to a restart or a checkpoint 81 | # being specified at the command line. 82 | self.ema_params = [ 83 | self._load_ema_parameters(rate) for rate in self.ema_rate 84 | ] 85 | else: 86 | self.ema_params = [ 87 | copy.deepcopy(self.mp_trainer.master_params) 88 | for _ in range(len(self.ema_rate)) 89 | ] 90 | 91 | if th.cuda.is_available(): 92 | self.use_ddp = True 93 | self.ddp_model = DDP( 94 | self.model, 95 | device_ids=[dist_util.dev()], 96 | output_device=dist_util.dev(), 97 | broadcast_buffers=False, 98 | bucket_cap_mb=128, 99 | find_unused_parameters=False, 100 | ) 101 | else: 102 | if dist.get_world_size() > 1: 103 | logger.warn( 104 | "Distributed training requires CUDA. " 105 | "Gradients will not be synchronized properly!" 106 | ) 107 | self.use_ddp = False 108 | self.ddp_model = self.model 109 | 110 | def _load_and_sync_parameters(self): 111 | resume_checkpoint = find_resume_checkpoint() or self.resume_checkpoint 112 | 113 | if resume_checkpoint: 114 | self.resume_step = parse_resume_step_from_filename(resume_checkpoint) 115 | #if dist.get_rank() == 0: 116 | logger.log(f"loading model from checkpoint: {resume_checkpoint}...") 117 | self.model.load_state_dict( 118 | dist_util.load_state_dict( 119 | resume_checkpoint, map_location=dist_util.dev() 120 | ) 121 | ) 122 | 123 | dist_util.sync_params(self.model.parameters()) 124 | 125 | def _load_ema_parameters(self, rate): 126 | ema_params = copy.deepcopy(self.mp_trainer.master_params) 127 | 128 | main_checkpoint = find_resume_checkpoint() or self.resume_checkpoint 129 | ema_checkpoint = find_ema_checkpoint(main_checkpoint, self.resume_step, rate) 130 | if ema_checkpoint: 131 | #if dist.get_rank() == 0: 132 | logger.log(f"loading EMA from checkpoint: {ema_checkpoint}...") 133 | state_dict = dist_util.load_state_dict( 134 | ema_checkpoint, map_location=dist_util.dev() 135 | ) 136 | ema_params = self.mp_trainer.state_dict_to_master_params(state_dict) 137 | 138 | dist_util.sync_params(ema_params) 139 | return ema_params 140 | 141 | def _load_optimizer_state(self): 142 | main_checkpoint = find_resume_checkpoint() or self.resume_checkpoint 143 | opt_checkpoint = bf.join( 144 | bf.dirname(main_checkpoint), f"opt{self.resume_step:06}.pt" 145 | ) 146 | if bf.exists(opt_checkpoint): 147 | logger.log(f"loading optimizer state from checkpoint: {opt_checkpoint}") 148 | state_dict = dist_util.load_state_dict( 149 | opt_checkpoint, map_location=dist_util.dev() 150 | ) 151 | self.opt.load_state_dict(state_dict) 152 | 153 | def run_loop(self): 154 | while ( 155 | not self.lr_anneal_steps 156 | or self.step + self.resume_step < self.lr_anneal_steps 157 | ): 158 | batch, cond = next(self.data) 159 | # print(batch.shape) 160 | # print(cond["low_res"].shape) 161 | self.run_step(batch, cond) 162 | if self.step % self.log_interval == 0: 163 | logger.dumpkvs() 164 | if self.step % self.save_interval == 0: 165 | self.save() 166 | # Run for a finite amount of time in integration tests. 167 | if os.environ.get("DIFFUSION_TRAINING_TEST", "") and self.step > 0: 168 | return 169 | self.step += 1 170 | # Save the last checkpoint if it wasn't already saved. 171 | if (self.step - 1) % self.save_interval != 0: 172 | self.save() 173 | 174 | def run_step(self, batch, cond): 175 | self.forward_backward(batch, cond) 176 | took_step = self.mp_trainer.optimize(self.opt) 177 | if took_step: 178 | self._update_ema() 179 | self._anneal_lr() 180 | self.log_step() 181 | 182 | def forward_backward(self, batch, cond): 183 | self.mp_trainer.zero_grad() 184 | for i in range(0, batch.shape[0], self.microbatch): 185 | micro = batch[i : i + self.microbatch].to(dist_util.dev()) 186 | micro_cond = { 187 | k: v[i : i + self.microbatch].to(dist_util.dev()) 188 | for k, v in cond.items() 189 | } 190 | last_batch = (i + self.microbatch) >= batch.shape[0] 191 | t, weights = self.schedule_sampler.sample(micro.shape[0], dist_util.dev()) 192 | 193 | compute_losses = functools.partial( 194 | self.diffusion.training_losses, 195 | self.ddp_model, 196 | micro, 197 | t, 198 | model_kwargs=micro_cond, 199 | ) 200 | 201 | if last_batch or not self.use_ddp: 202 | losses = compute_losses() 203 | else: 204 | with self.ddp_model.no_sync(): 205 | losses = compute_losses() 206 | 207 | if isinstance(self.schedule_sampler, LossAwareSampler): 208 | self.schedule_sampler.update_with_local_losses( 209 | t, losses["loss"].detach() 210 | ) 211 | 212 | loss = (losses["loss"] * weights).mean() 213 | log_loss_dict( 214 | self.diffusion, t, {k: v * weights for k, v in losses.items()} 215 | ) 216 | self.mp_trainer.backward(loss) 217 | 218 | def _update_ema(self): 219 | for rate, params in zip(self.ema_rate, self.ema_params): 220 | update_ema(params, self.mp_trainer.master_params, rate=rate) 221 | 222 | def _anneal_lr(self): 223 | if not self.lr_anneal_steps: 224 | return 225 | frac_done = (self.step + self.resume_step) / self.lr_anneal_steps 226 | lr = self.lr * (1 - frac_done) 227 | for param_group in self.opt.param_groups: 228 | param_group["lr"] = lr 229 | 230 | def log_step(self): 231 | logger.logkv("step", self.step + self.resume_step) 232 | logger.logkv("samples", (self.step + self.resume_step + 1) * self.global_batch) 233 | 234 | def save(self): 235 | def save_checkpoint(rate, params): 236 | state_dict = self.mp_trainer.master_params_to_state_dict(params) 237 | if dist.get_rank() == 0: 238 | logger.log(f"saving model {rate}...") 239 | if not rate: 240 | filename = f"model{(self.step+self.resume_step):06d}.pt" 241 | else: 242 | filename = f"ema_{rate}_{(self.step+self.resume_step):06d}.pt" 243 | with bf.BlobFile(bf.join(get_blob_logdir(), filename), "wb") as f: 244 | th.save(state_dict, f) 245 | 246 | save_checkpoint(0, self.mp_trainer.master_params) 247 | for rate, params in zip(self.ema_rate, self.ema_params): 248 | save_checkpoint(rate, params) 249 | 250 | if dist.get_rank() == 0: 251 | with bf.BlobFile( 252 | bf.join(get_blob_logdir(), f"opt{(self.step+self.resume_step):06d}.pt"), 253 | "wb", 254 | ) as f: 255 | th.save(self.opt.state_dict(), f) 256 | 257 | dist.barrier() 258 | 259 | 260 | def parse_resume_step_from_filename(filename): 261 | """ 262 | Parse filenames of the form path/to/modelNNNNNN.pt, where NNNNNN is the 263 | checkpoint's number of steps. 264 | """ 265 | split = filename.split("model") 266 | if len(split) < 2: 267 | return 0 268 | split1 = split[-1].split(".")[0] 269 | try: 270 | return int(split1) 271 | except ValueError: 272 | return 0 273 | 274 | 275 | def get_blob_logdir(): 276 | # You can change this to be a separate path to save checkpoints to 277 | # a blobstore or some external drive. 278 | return logger.get_dir() 279 | # return '/autofs/vast/kggp/ydong/seg_ddpm/save_models/3d/pet_ct/ddpm96_32_noatt2_-1to1_2' 280 | 281 | 282 | def find_resume_checkpoint(): 283 | # On your infrastructure, you may want to override this to automatically 284 | # discover the latest checkpoint on your blob storage, etc. 285 | return None 286 | 287 | 288 | def find_ema_checkpoint(main_checkpoint, step, rate): 289 | if main_checkpoint is None: 290 | return None 291 | filename = f"ema_{rate}_{(step):06d}.pt" 292 | path = bf.join(bf.dirname(main_checkpoint), filename) 293 | if bf.exists(path): 294 | return path 295 | return None 296 | 297 | 298 | def log_loss_dict(diffusion, ts, losses): 299 | for key, values in losses.items(): 300 | logger.logkv_mean(key, values.mean().item()) 301 | # Log the quantiles (four quartiles, in particular). 302 | for sub_t, sub_loss in zip(ts.cpu().numpy(), values.detach().cpu().numpy()): 303 | quartile = int(4 * sub_t / diffusion.num_timesteps) 304 | logger.logkv_mean(f"{key}_q{quartile}", sub_loss) 305 | -------------------------------------------------------------------------------- /scripts/test.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import os 3 | import sys 4 | sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))) 5 | 6 | import blobfile as bf 7 | import numpy as np 8 | import torch as th 9 | import torch.distributed as dist 10 | import datetime 11 | 12 | from guided_diffusion import dist_util, logger 13 | from guided_diffusion.script_util import ( 14 | sr_model_and_diffusion_defaults, 15 | sr_create_model_and_diffusion, 16 | args_to_dict, 17 | add_dict_to_argparser, 18 | ) 19 | 20 | 21 | def main(): 22 | args = create_argparser().parse_args() 23 | 24 | dist_util.setup_dist() 25 | logger.configure(dir=args.save_dir) 26 | 27 | logger.log("creating model...") 28 | model, diffusion = sr_create_model_and_diffusion( 29 | **args_to_dict(args, sr_model_and_diffusion_defaults().keys()) 30 | ) 31 | model.load_state_dict( 32 | dist_util.load_state_dict(args.model_path, map_location="cpu") 33 | ) 34 | model.to(dist_util.dev()) 35 | if args.use_fp16: 36 | model.convert_to_fp16() 37 | model.eval() 38 | 39 | logger.log("loading data...") 40 | data = load_data_for_worker(args.base_samples, args.batch_size, args.class_cond) 41 | logger.log("creating samples...") 42 | all_images = [] 43 | 44 | device = th.device("cuda" if th.cuda.is_available() else "cpu") 45 | 46 | 47 | while len(all_images) * args.batch_size < args.num_samples: 48 | model_kwargs = next(data) 49 | model_kwargs = {k: v.to(dist_util.dev()) for k, v in model_kwargs.items()} 50 | 51 | shape = (args.batch_size, 1, model_kwargs['low_res'].shape[2], model_kwargs['low_res'].shape[3], model_kwargs['low_res'].shape[4]) 52 | 53 | if device == "cuda": 54 | th.cuda.manual_seed_all(10) 55 | else: 56 | th.manual_seed(10) 57 | noise = th.randn(*shape, device=device) 58 | 59 | sample = diffusion.p_sample_loop( 60 | model, 61 | shape, 62 | noise, 63 | clip_denoised=args.clip_denoised, 64 | model_kwargs=model_kwargs, 65 | ) 66 | 67 | sample = sample.permute(0, 1, 3, 4, 2) 68 | sample = sample.contiguous() 69 | 70 | all_samples = [th.zeros_like(sample) for _ in range(dist.get_world_size())] 71 | dist.all_gather(all_samples, sample) 72 | for sample in all_samples: 73 | all_images.append(sample.cpu().numpy()) 74 | logger.log(f"created {len(all_images) * args.batch_size} samples\n") 75 | 76 | arr = np.concatenate(all_images, axis=0) 77 | arr = arr[: args.num_samples] 78 | arr = arr.reshape((arr.shape[0],arr.shape[2],arr.shape[3],arr.shape[4])) 79 | 80 | arr_result = np.zeros((192,288,576)) 81 | index = 0 82 | for i in range(6): 83 | arr_result[:, :, i*96:(i+1)*96] = arr[index,:,:,:] 84 | index += 1 85 | 86 | if dist.get_rank() == 0: 87 | shape_str = "x".join([str(x) for x in arr_result.shape]) 88 | out_path = os.path.join(logger.get_dir(), f"samples_{shape_str}_{datetime.datetime.now().strftime('%H%M%S%f')}.npz") 89 | logger.log(f"saving to {out_path}") 90 | np.savez(out_path, arr_result) 91 | 92 | dist.barrier() 93 | logger.log("sampling complete") 94 | 95 | 96 | def load_data_for_worker(base_samples, batch_size, class_cond): 97 | with bf.BlobFile(base_samples, "rb") as f: 98 | obj = np.load(f) 99 | low_pet = obj["arr_0"][0] 100 | image_arr = np.zeros((6,192,288,96)) 101 | index = 0 102 | for i in (0, 86, 172, 258, 344, 424): 103 | image_arr[index,:,:,:] = low_pet[:, :, i:i+96] 104 | index += 1 105 | 106 | image_arr[image_arr>4] = 4 107 | image_arr = image_arr/4 108 | 109 | rank = dist.get_rank() 110 | logger.log('rank:{%d}' % (rank)) 111 | 112 | num_ranks = dist.get_world_size() 113 | logger.log('num_ranks:{%d}' % (num_ranks)) 114 | 115 | buffer = [] 116 | label_buffer = [] 117 | while True: 118 | for i in range(rank, len(image_arr), num_ranks): 119 | buffer.append(image_arr[i]) 120 | logger.log('rank:{%d}, i:{%d}, buffer_len:{%d}' % (rank, i, len(buffer))) 121 | 122 | if len(buffer) == batch_size: 123 | batch = th.from_numpy(np.stack(buffer)).float() 124 | batch = batch.unsqueeze(0) 125 | batch = batch.permute(0, 1, 4, 2, 3) 126 | logger.log('batch_shape:{%d,%d,%d,%d,%d}' % (batch.shape[0], batch.shape[1], batch.shape[2], batch.shape[3], batch.shape[4])) 127 | res = dict(low_res=batch) 128 | yield res 129 | buffer, label_buffer = [], [] 130 | 131 | 132 | def create_argparser(): 133 | defaults = dict( 134 | save_dir='', 135 | clip_denoised=True, 136 | num_samples=10000, 137 | batch_size=16, 138 | use_ddim=False, 139 | base_samples="", 140 | model_path="", 141 | ) 142 | defaults.update(sr_model_and_diffusion_defaults()) 143 | parser = argparse.ArgumentParser() 144 | add_dict_to_argparser(parser, defaults) 145 | return parser 146 | 147 | 148 | if __name__ == "__main__": 149 | main() 150 | -------------------------------------------------------------------------------- /scripts/train.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import os 3 | import torch.nn.functional as F 4 | import sys 5 | sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))) 6 | 7 | from guided_diffusion import dist_util, logger 8 | from guided_diffusion.image_datasets import load_data 9 | from guided_diffusion.resample import create_named_schedule_sampler 10 | from guided_diffusion.script_util import ( 11 | sr_model_and_diffusion_defaults, 12 | sr_create_model_and_diffusion, 13 | args_to_dict, 14 | add_dict_to_argparser, 15 | ) 16 | from guided_diffusion.train_util import TrainLoop 17 | 18 | 19 | def main(): 20 | args = create_argparser().parse_args() 21 | 22 | dist_util.setup_dist() 23 | # logger.configure(dir='/autofs/vast/kggp/ydong/seg_ddpm/logs/3d/pet_ct/ddpm96_32_noatt2_-1to1_2') 24 | logger.configure(args.result_folder) 25 | 26 | logger.log("creating model...") 27 | model, diffusion = sr_create_model_and_diffusion( 28 | **args_to_dict(args, sr_model_and_diffusion_defaults().keys()) 29 | ) 30 | model.to(dist_util.dev()) 31 | schedule_sampler = create_named_schedule_sampler(args.schedule_sampler, diffusion) 32 | for name, param in model.named_parameters(): 33 | logger.log('%s:%s' % (name, str(param.shape))) 34 | logger.log('parameters:{%d}' % (sum(param.numel() for param in model.parameters()))) 35 | logger.log('attention_resolutions:{%s}' % args.attention_resolutions) 36 | logger.log('num_channels:{%s}' % str(args.num_channels)) 37 | logger.log('num_res_blocks:{%s}' % str(args.num_res_blocks)) 38 | logger.log('num_head_channels:{%s}' % str(args.num_head_channels)) 39 | logger.log('channel_mult:{%s}' % str(model.channel_mult)) 40 | 41 | logger.log("creating data loader...") 42 | data = load_superres_data( 43 | args.data_dir, 44 | args.batch_size, 45 | large_size=args.large_size, 46 | small_size=args.small_size, 47 | class_cond=args.class_cond, 48 | ) 49 | 50 | logger.log("training...") 51 | 52 | TrainLoop( 53 | model=model, 54 | diffusion=diffusion, 55 | data=data, 56 | batch_size=args.batch_size, 57 | microbatch=args.microbatch, 58 | lr=args.lr, 59 | ema_rate=args.ema_rate, 60 | log_interval=args.log_interval, 61 | save_interval=args.save_interval, 62 | resume_checkpoint=args.resume_checkpoint, 63 | use_fp16=args.use_fp16, 64 | fp16_scale_growth=args.fp16_scale_growth, 65 | schedule_sampler=schedule_sampler, 66 | weight_decay=args.weight_decay, 67 | lr_anneal_steps=args.lr_anneal_steps, 68 | ).run_loop() 69 | 70 | 71 | def load_superres_data(data_dir, batch_size, large_size, small_size, class_cond=False): 72 | data = load_data( 73 | data_dir=data_dir, 74 | batch_size=batch_size, 75 | image_size=large_size, 76 | class_cond=class_cond, 77 | ) 78 | for pet_batch, label_batch, model_kwargs in data: 79 | model_kwargs["low_res"] = pet_batch.clone() 80 | yield label_batch, model_kwargs 81 | 82 | ''' 83 | def load_superres_data(data_dir, batch_size, large_size, small_size, class_cond=False): 84 | data = load_data( 85 | data_dir=data_dir, 86 | batch_size=batch_size, 87 | image_size=large_size, 88 | class_cond=class_cond, 89 | ) 90 | for large_batch, model_kwargs in data: 91 | model_kwargs["low_res"] = F.interpolate(large_batch, small_size, mode="area") 92 | yield large_batch, model_kwargs 93 | ''' 94 | 95 | def create_argparser(): 96 | defaults = dict( 97 | data_dir="", 98 | schedule_sampler="uniform", 99 | lr=1e-4, 100 | weight_decay=0.0, 101 | lr_anneal_steps=0, 102 | batch_size=1, 103 | microbatch=-1, 104 | ema_rate="0.9999", 105 | log_interval=10, 106 | save_interval=10000, 107 | resume_checkpoint="", 108 | use_fp16=False, 109 | fp16_scale_growth=1e-3, 110 | result_folder = None, 111 | ) 112 | defaults.update(sr_model_and_diffusion_defaults()) 113 | parser = argparse.ArgumentParser() 114 | add_dict_to_argparser(parser, defaults) 115 | return parser 116 | 117 | 118 | if __name__ == "__main__": 119 | main() 120 | -------------------------------------------------------------------------------- /test_DDPM_3d_mpi.sh: -------------------------------------------------------------------------------- 1 | SAMPLE_FLAGS="--batch_size 1 --num_samples 6" 2 | MODEL_FLAGS="--attention_resolutions 1000 --large_size 96 --small_size 96 --num_channels 128 --use_fp16 True --num_head_channels 64 --learn_sigma True --resblock_updown True --use_scale_shift_norm True" 3 | DIFFUSION_FLAGS="--diffusion_steps 1000 --noise_schedule linear --rescale_learned_sigmas False --rescale_timesteps False" 4 | 5 | mpiexec -n 6 python ./scripts/test.py $MODEL_FLAGS --model_path ./checkpoints/model.pt --base_samples sample_PET.npz --save_dir ./results/ $SAMPLE_FLAGS 6 | 7 | --------------------------------------------------------------------------------