upgrade to best type of downsample

lucidrains · Aug 25, 2022 · 392a0b1 · 392a0b1
1 parent 65c8bb9
commit 392a0b1
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 2 deletions.
diff --git a/README.md b/README.md
@@ -101,3 +101,13 @@ out = unet(video) # (1, 3, 10, 128, 128)
     volume  = {abs/2110.09456}
 }
 ```
+
+```bibtex
+@article{Sunkara2022NoMS,
+    title   = {No More Strided Convolutions or Pooling: A New CNN Building Block for Low-Resolution Images and Small Objects},
+    author  = {Raja Sunkara and Tie Luo},
+    journal = {ArXiv},
+    year    = {2022},
+    volume  = {abs/2208.03641}
+}
+```
diff --git a/setup.py b/setup.py
@@ -3,7 +3,7 @@
 setup(
   name = 'x-unet',
   packages = find_packages(exclude=[]),
-  version = '0.0.20',
+  version = '0.0.21',
   license='MIT',
   description = 'X-Unet',
   long_description_content_type = 'text/markdown',

diff --git a/x_unet/x_unet.py b/x_unet/x_unet.py
@@ -39,7 +39,10 @@ def Upsample(dim, dim_out):
     return nn.ConvTranspose3d(dim, dim_out, (1, 4, 4), (1, 2, 2), (0, 1, 1))
 
 def Downsample(dim, dim_out):
-    return nn.Conv3d(dim, dim_out, (1, 4, 4), (1, 2, 2), (0, 1, 1))
+    return nn.Sequential(
+        Rearrange('b c f (h s1) (w s2) -> b (c s1 s2) f h w', s1 = 2, s2 = 2),
+        nn.Conv3d(dim * 4, dim_out, 1)
+    )
 
 # normalization