Fix for Issue deepinv#647 – UNet uses zero-padding even when it shouldn't (deepinv#653)

vsechaud · web-flow · commit 99526f099236 · 2025-07-31T17:42:00.000+02:00
* fix bug unet padding mode

* black unet file

* change log
diff --git a/CHANGELOG.rst b/CHANGELOG.rst
@@ -16,7 +16,7 @@ Changed
 Fixed
 ^^^^^
 - Fix memory leak in `deepinv.physics.tomography` when using autograd (:gh:`651` by `Minh Hai Nguyen`_)
-
+- Fix the circular padded UNet (:gh:`653` by `Victor Sechaud`_)
 
 v0.3.2
 ------
diff --git a/deepinv/models/unet.py b/deepinv/models/unet.py
@@ -116,7 +116,13 @@ def conv_block(ch_in, ch_out):
                     ),
                     nn.ReLU(inplace=True),
                     nn.Conv2d(
-                        ch_out, ch_out, kernel_size=3, stride=1, padding=1, bias=bias
+                        ch_out,
+                        ch_out,
+                        kernel_size=3,
+                        stride=1,
+                        padding=1,
+                        bias=bias,
+                        padding_mode="circular" if circular_padding else "zeros",
                     ),
                     (
                         BFBatchNorm2d(ch_out, use_bias=bias)
@@ -138,7 +144,13 @@ def conv_block(ch_in, ch_out):
                     ),
                     nn.ReLU(inplace=True),
                     nn.Conv2d(
-                        ch_out, ch_out, kernel_size=3, stride=1, padding=1, bias=bias
+                        ch_out,
+                        ch_out,
+                        kernel_size=3,
+                        stride=1,
+                        padding=1,
+                        bias=bias,
+                        padding_mode="circular" if circular_padding else "zeros",
                     ),
                     nn.ReLU(inplace=True),
                 )
@@ -148,7 +160,13 @@ def up_conv(ch_in, ch_out):
                 return nn.Sequential(
                     nn.Upsample(scale_factor=2),
                     nn.Conv2d(
-                        ch_in, ch_out, kernel_size=3, stride=1, padding=1, bias=bias
+                        ch_in,
+                        ch_out,
+                        kernel_size=3,
+                        stride=1,
+                        padding=1,
+                        bias=bias,
+                        padding_mode="circular" if circular_padding else "zeros",
                     ),
                     (
                         BFBatchNorm2d(ch_out, use_bias=bias)
@@ -161,7 +179,13 @@ def up_conv(ch_in, ch_out):
                 return nn.Sequential(
                     nn.Upsample(scale_factor=2),
                     nn.Conv2d(
-                        ch_in, ch_out, kernel_size=3, stride=1, padding=1, bias=bias
+                        ch_in,
+                        ch_out,
+                        kernel_size=3,
+                        stride=1,
+                        padding=1,
+                        bias=bias,
+                        padding_mode="circular" if circular_padding else "zeros",
                     ),
                     nn.ReLU(inplace=True),
                 )
