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

add some training fns

Browse Source
This commit is contained in:
Phil Wang
2022-04-29 09:44:55 -07:00
parent fb662a62f3
commit f4a54e475e
1 changed files with 9 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
dalle2_pytorch/train.py
View File

@@ -2,6 +2,15 @@ import copy
import torch import torch
from torch import nn from torch import nn
# image related normalizations
# ddpms expect images to be in the range of -1 to 1
def normalize_img(img):
return img * 2 - 1
def unnormalize_img(normed_img):
return (normed_img + 1) * 0.5
# exponential moving average wrapper # exponential moving average wrapper
class EMA(nn.Module): class EMA(nn.Module):