Update unet3d.py

Author: anna-grim

src/aind_exaspim_image_compression/machine_learning/unet3d.py

@@ -115,7 +115,7 @@ class DoubleConv(nn.Module):
         activations.
     """
 
-    def __init__(self, in_channels, out_channels, mid_channels=None):
+    def __init__(self, in_channels, out_channels, kernel_size=3, mid_channels=None):
         """
         Instantiates a DoubleConv object.
 
@@ -125,6 +125,8 @@ def __init__(self, in_channels, out_channels, mid_channels=None):
             Number of input channels to this module.
         out_channels : int
             Number of output channels produced by this module.
+        kernel_size : int, optional
+            Size of kernel used in convolutional layers. Default is 3.
         mid_channels : int, optional
             Number of channels in the intermediate convolution. Default is
             None.
@@ -138,10 +140,10 @@ def __init__(self, in_channels, out_channels, mid_channels=None):
 
         # Instance attributes
         self.double_conv = nn.Sequential(
-            nn.Conv3d(in_channels, mid_channels, kernel_size=4, padding=1),
+            nn.Conv3d(in_channels, mid_channels, kernel_size=kernel_size, padding=1),
             nn.BatchNorm3d(mid_channels),
             nn.LeakyReLU(negative_slope=0.01, inplace=True),
-            nn.Conv3d(mid_channels, out_channels, kernel_size=4, padding=1),
+            nn.Conv3d(mid_channels, out_channels, kernel_size=kernel_size, padding=1),
             nn.BatchNorm3d(out_channels),
             nn.LeakyReLU(negative_slope=0.01, inplace=True)
         )