The Nuts and Bolts of Parallel-UNet: Implementation Details

The implementation of TryOnDiffusion with Parallel-UNet includes a 256×256 architecture with key changes for improved performance. Trained using JAX on TPU-v4 for 500K iterations, the inference process is efficient, taking around 18 seconds for a batch of four.


This content originally appeared on HackerNoon and was authored by Backpropagation

:::info Authors:

(1) Luyang Zhu, University of Washington and Google Research, and work done while the author was an intern at Google;

(2) Dawei Yang, Google Research;

(3) Tyler Zhu, Google Research;

(4) Fitsum Reda, Google Research;

(5) William Chan, Google Research;

(6) Chitwan Saharia, Google Research;

(7) Mohammad Norouzi, Google Research;

(8) Ira Kemelmacher-Shlizerman, University of Washington and Google Research.

:::

Abstract and 1. Introduction

2. Related Work

3. Method

3.1. Cascaded Diffusion Models for Try-On

3.2. Parallel-UNet

4. Experiments

5. Summary and Future Work and References

\ Appendix

A. Implementation Details

B. Additional Results

A. Implementation Details

A.1. Parallel-UNet

A.2. Training and Inference

TryOnDiffusion was implemented in JAX [4]. All three diffusion models are trained on 32 TPU-v4 chips for 500K iterations (around 3 days for each diffusion model). After trained, we run the inference of the whole pipeline on 4 TPU-v4 chips with batch size 4, which takes around 18 seconds for one batch.

\

:::info This paper is available on arxiv under CC BY-NC-ND 4.0 DEED license.

:::

\


This content originally appeared on HackerNoon and was authored by Backpropagation


Print Share Comment Cite Upload Translate Updates
APA

Backpropagation | Sciencx (2024-10-06T20:04:33+00:00) The Nuts and Bolts of Parallel-UNet: Implementation Details. Retrieved from https://www.scien.cx/2024/10/06/the-nuts-and-bolts-of-parallel-unet-implementation-details/

MLA
" » The Nuts and Bolts of Parallel-UNet: Implementation Details." Backpropagation | Sciencx - Sunday October 6, 2024, https://www.scien.cx/2024/10/06/the-nuts-and-bolts-of-parallel-unet-implementation-details/
HARVARD
Backpropagation | Sciencx Sunday October 6, 2024 » The Nuts and Bolts of Parallel-UNet: Implementation Details., viewed ,<https://www.scien.cx/2024/10/06/the-nuts-and-bolts-of-parallel-unet-implementation-details/>
VANCOUVER
Backpropagation | Sciencx - » The Nuts and Bolts of Parallel-UNet: Implementation Details. [Internet]. [Accessed ]. Available from: https://www.scien.cx/2024/10/06/the-nuts-and-bolts-of-parallel-unet-implementation-details/
CHICAGO
" » The Nuts and Bolts of Parallel-UNet: Implementation Details." Backpropagation | Sciencx - Accessed . https://www.scien.cx/2024/10/06/the-nuts-and-bolts-of-parallel-unet-implementation-details/
IEEE
" » The Nuts and Bolts of Parallel-UNet: Implementation Details." Backpropagation | Sciencx [Online]. Available: https://www.scien.cx/2024/10/06/the-nuts-and-bolts-of-parallel-unet-implementation-details/. [Accessed: ]
rf:citation
» The Nuts and Bolts of Parallel-UNet: Implementation Details | Backpropagation | Sciencx | https://www.scien.cx/2024/10/06/the-nuts-and-bolts-of-parallel-unet-implementation-details/ |

Please log in to upload a file.




There are no updates yet.
Click the Upload button above to add an update.

You must be logged in to translate posts. Please log in or register.