aljaz/DALLE2-pytorch
1
1
Fork 0
mirror of https://github.com/lucidrains/DALLE2-pytorch.git synced 2025-12-20 02:04:19 +01:00
Code Issues Packages Projects Releases Wiki Activity

update working unconditional decoder example

Browse Source
This commit is contained in:
Phil Wang
2022-05-16 12:50:07 -07:00
parent 164d9be444
commit c3d4a7ffe4
1 changed files with 10 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
README.md
View File

@@ -867,7 +867,7 @@ ex.
```python ```python
import torch import torch
from dalle2_pytorch import Unet, Decoder from dalle2_pytorch import Unet, Decoder, DecoderTrainer
# unet for the cascading ddpm # unet for the cascading ddpm
@@ -890,20 +890,24 @@ decoder = Decoder(
unconditional = True unconditional = True
).cuda() ).cuda()
# mock images (get a lot of this) # decoder trainer
decoder_trainer = DecoderTrainer(decoder)
# images (get a lot of this)
images = torch.randn(1, 3, 512, 512).cuda() images = torch.randn(1, 3, 512, 512).cuda()
# feed images into decoder # feed images into decoder
for i in (1, 2): for i in (1, 2):
loss = decoder(images, unet_number = i) loss = decoder_trainer(images, unet_number = i)
loss.backward() decoder_trainer.update(unet_number = i)
# do the above for many many many many steps # do the above for many many many many images
# then it will learn to generate images # then it will learn to generate images
images = decoder.sample(batch_size = 2) # (2, 3, 512, 512) images = decoder_trainer.sample(batch_size = 36, max_batch_size = 4) # (36, 3, 512, 512)
``` ```
## Dataloaders ## Dataloaders