use sigmoid postprocessing in unet2d example

example_descriptions/models/unet2d_nuclei_broad/bioimageio.yaml

@@ -93,6 +93,11 @@ outputs:
     sample_tensor:
       source: test_output.png
       sha256: 7bce8b53bcd0a12487a61f953aafe0f3700652848980d1083964c5bcb9555eec
+    postprocessing:
+      - id: sigmoid
+      - id: ensure_dtype
+        kwargs:
+          dtype: float32
 
 weights:
   pytorch_state_dict:
@@ -105,7 +110,7 @@ weights:
     architecture:
       callable: UNet2d
       source: unet2d.py
-      sha256: 7cdd8332dc3e3735e71c328f81b63a9ac86c028f80522312484ca9a4027d4ce1
+      sha256: 589f0c9e60fa00f015213cd251541bcbf9582644f3ecffb2e6f3a30d2af1931a
       kwargs: { input_channels: 1, output_channels: 1 }
     dependencies:
       source: environment.yaml

example_descriptions/models/unet2d_nuclei_broad/expand_output_shape_v0_4.bioimageio.yaml

@@ -54,6 +54,8 @@ outputs:
       reference_tensor: raw
       scale: [1.0, 0.0, null, 1.0, 1.0]
       offset: [0.0, 0.5, 0.5, 0.0, 0.0]
+    postprocessing:
+      - name: sigmoid
 
 test_inputs: [test_input.npy]
 test_outputs: [test_output_expanded.npy]
@@ -67,7 +69,7 @@ weights:
     sha256: e4d3885bccbe41cbf6c1d825f3cd2b707c7021ead5593156007e407a16b27cf2
     source: https://zenodo.org/records/3446812/files/unet2d_weights.torch
     architecture: unet2d_expand_output_shape.py:UNet2d
-    architecture_sha256: 80a886acc734f848a8e018d8063cfd7e003d7e20076583b28326bfdd6136be32
+    architecture_sha256: 1441e8cfaf387a98a1c0bb937d59a2e9d6c311a8912cd88b39c11ecff503ccfe
     kwargs: { input_channels: 1, output_channels: 1 }
     dependencies: conda:environment.yaml
 

example_descriptions/models/unet2d_nuclei_broad/invalid_v0_4_0_duplicate_tensor_names.bioimageio.yaml

@@ -48,6 +48,8 @@ outputs:
     data_range: [-.inf, .inf]
     halo: [0, 0, 32, 32]
     shape: [1, 1, 512, 512]
+    postprocessing:
+      - name: sigmoid
 
 dependencies: conda:environment.yaml
 
@@ -63,7 +65,7 @@ weights:
     sha256: e4d3885bccbe41cbf6c1d825f3cd2b707c7021ead5593156007e407a16b27cf2
     source: https://zenodo.org/records/3446812/files/unet2d_weights.torch
     architecture: unet2d.py:UNet2d
-    architecture_sha256: 7cdd8332dc3e3735e71c328f81b63a9ac86c028f80522312484ca9a4027d4ce1
+    architecture_sha256: 589f0c9e60fa00f015213cd251541bcbf9582644f3ecffb2e6f3a30d2af1931a
     kwargs: { input_channels: 1, output_channels: 1 }
   onnx:
     sha256: f1f086d5e340f9d4d7001a1b62a2b835f9b87a2fb5452c4fe7d8cc821bdf539c

example_descriptions/models/unet2d_nuclei_broad/unet2d.py

@@ -16,7 +16,7 @@ def forward(self, input):
 
 
 class UNet2d(nn.Module):
-    def __init__(self, input_channels, output_channels, training=False):
+    def __init__(self, input_channels, output_channels):
         super().__init__()
         self.input_channels = input_channels
         self.output_channels = output_channels
@@ -41,7 +41,6 @@ def __init__(self, input_channels, output_channels, training=False):
         )
 
         self.output = nn.Conv2d(16, self.output_channels, 1)
-        self.training = training
 
     def conv_layer(self, in_channels, out_channels):
         kernel_size = 3
@@ -78,9 +77,4 @@ def forward(self, input):
 
         x = self.output(x)
 
-        # apply a sigmoid directly if we are in inference mode
-        if not self.training:
-            # postprocessing
-            x = torch.sigmoid(x)
-
         return x

example_descriptions/models/unet2d_nuclei_broad/unet2d_expand_output_shape.py

@@ -16,7 +16,7 @@ def forward(self, input):
 
 
 class UNet2d(nn.Module):
-    def __init__(self, input_channels, output_channels, training=False):
+    def __init__(self, input_channels, output_channels):
         super().__init__()
         self.input_channels = input_channels
         self.output_channels = output_channels
@@ -41,7 +41,6 @@ def __init__(self, input_channels, output_channels, training=False):
         )
 
         self.output = nn.Conv2d(16, self.output_channels, 1)
-        self.training = training
 
     def conv_layer(self, in_channels, out_channels):
         kernel_size = 3
@@ -78,11 +77,6 @@ def forward(self, input):
 
         x = self.output(x)
 
-        # apply a sigmoid directly if we are in inference mode
-        if not self.training:
-            # postprocessing
-            x = torch.sigmoid(x)
-
         # expand the shape across z
         out_shape = tuple(x.shape)
         expanded_shape = out_shape[:2] + (1,) + out_shape[2:]

example_descriptions/models/unet2d_nuclei_broad/v0_4_0.bioimageio.yaml

@@ -54,6 +54,8 @@ outputs:
       reference_tensor: raw
       scale: [1.0, 0.0, 1.0, 1.0]
       offset: [0.0, 0.5, 0.0, 0.0]
+    postprocessing:
+      - name: sigmoid
 
 dependencies: conda:environment.yaml
 
@@ -72,7 +74,7 @@ weights:
     source: https://zenodo.org/records/3446812/files/unet2d_weights.torch
     sha256: e4d3885bccbe41cbf6c1d825f3cd2b707c7021ead5593156007e407a16b27cf2
     architecture: unet2d.py:UNet2d
-    architecture_sha256: 7cdd8332dc3e3735e71c328f81b63a9ac86c028f80522312484ca9a4027d4ce1
+    architecture_sha256: 589f0c9e60fa00f015213cd251541bcbf9582644f3ecffb2e6f3a30d2af1931a
     kwargs: { input_channels: 1, output_channels: 1 }
   onnx:
     source: weights.onnx

example_descriptions/models/unet2d_nuclei_broad/v0_4_9.bioimageio.yaml

@@ -53,6 +53,8 @@ outputs:
       reference_tensor: raw
       scale: [1.0, 0.0, 1.0, 1.0]
       offset: [0.0, 0.5, 0.0, 0.0]
+    postprocessing:
+      - name: sigmoid
 
 test_inputs: [test_input.npy]
 test_outputs: [test_output.npy]
@@ -69,7 +71,7 @@ weights:
     source: https://zenodo.org/records/3446812/files/unet2d_weights.torch
     sha256: e4d3885bccbe41cbf6c1d825f3cd2b707c7021ead5593156007e407a16b27cf2
     architecture: unet2d.py:UNet2d
-    architecture_sha256: 7cdd8332dc3e3735e71c328f81b63a9ac86c028f80522312484ca9a4027d4ce1
+    architecture_sha256: 589f0c9e60fa00f015213cd251541bcbf9582644f3ecffb2e6f3a30d2af1931a
     kwargs: { input_channels: 1, output_channels: 1 }
     dependencies: conda:environment.yaml
   onnx: