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

upgrade to best downsample

Browse Source
This commit is contained in:
Phil Wang
2022-08-25 10:37:02 -07:00
parent 59fa101c4d
commit 1cc5d0afa7
4 changed files with 19 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
dalle2_pytorch/dalle2_pytorch.py
View File

@@ -1479,9 +1479,14 @@ class PixelShuffleUpsample(nn.Module):
def forward(self, x):
return self.net(x)
def Downsample(dim, *, dim_out = None):
def Downsample(dim, dim_out = None):
# https://arxiv.org/abs/2208.03641 shows this is the most optimal way to downsample
# named SP-conv in the paper, but basically a pixel unshuffle
dim_out = default(dim_out, dim)
return nn.Conv2d(dim, dim_out, 4, 2, 1)
return nn.Sequential(
Rearrange('b c (h s1) (w s2) -> b (c s1 s2) h w', s1 = 2, s2 = 2),
nn.Conv2d(dim * 4, dim_out, 1)
)
class WeightStandardizedConv2d(nn.Conv2d):
"""