Update body model

body/Body.pt

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:0a4652d0c8139e2d3ce86247bf3e050671de369fe982a701bf6bf76304131c7a
-size 423981609
+oid sha256:be2218b50786e326e38c65c5cf2e2745f354911acbb4ae86bd1ca702db316d1c
+size 47484709

body/Model.py

@@ -1,10 +1,6 @@
-import torch
-from torch import nn
-
 from dynamic_network_architectures.architectures.unet import ResidualEncoderUNet
 from konfai.network import network
 
-
 def _replace_unpicklable_identities(module: nn.Module) -> None:
     """Replace library lambdas such as ``lambda x: x`` with ``nn.Identity``."""
     for child in module.modules():
@@ -13,7 +9,7 @@ def _replace_unpicklable_identities(module: nn.Module) -> None:
         if hasattr(child, "nonlin2") and callable(child.nonlin2) and not isinstance(child.nonlin2, nn.Module):
             child.nonlin2 = nn.Identity()
     
-class ResEncMCT(network.Network):
+class ResEnc(network.Network):
     def __init__(
         self,
         optimizer: network.OptimizerLoader = network.OptimizerLoader(),
@@ -32,22 +28,22 @@ class ResEncMCT(network.Network):
             dim=2,
         )
         self.add_module("DecoderOutputs", ResidualEncoderUNet(
-            input_channels=in_channels,
-            n_stages=7,
-            features_per_stage=(32, 64, 128, 256, 512, 512, 512),
-            conv_op=nn.Conv2d,
-            kernel_sizes=3,
-            strides=(1, 2, 2, 2, 2, 2, 2),
-            n_blocks_per_stage=(1, 3, 4, 6, 6, 6, 6),
-            num_classes=nb_class,
-            n_conv_per_stage_decoder=(1, 1, 1, 1, 1, 1),
-            conv_bias=True,
-            norm_op=nn.InstanceNorm2d,
-            norm_op_kwargs={"eps": 1e-5, "affine": True},
-            dropout_op=None,
-            dropout_op_kwargs=None,
-            nonlin=nn.LeakyReLU,
-            nonlin_kwargs={"inplace": True},
-            deep_supervision=False,
-        ))
+                input_channels=in_channels,
+                n_stages=6,
+                features_per_stage=(24, 48, 96, 192, 256, 256),
+                conv_op=nn.Conv2d,
+                kernel_sizes=3,
+                strides=(1, 2, 2, 2, 2, 2),
+                n_blocks_per_stage=(1, 2, 2, 3, 3, 3),
+                num_classes=nb_class,
+                n_conv_per_stage_decoder=(1, 1, 1, 1, 1),
+                conv_bias=True,
+                norm_op=nn.InstanceNorm2d,
+                norm_op_kwargs={"eps": 1e-5, "affine": True},
+                dropout_op=None,
+                dropout_op_kwargs=None,
+                nonlin=nn.LeakyReLU,
+                nonlin_kwargs={"inplace": True},
+                deep_supervision=False,
+            ))
         _replace_unpicklable_identities(self.DecoderOutputs)

body/Prediction.yml

@@ -1,7 +1,7 @@
 Predictor:
   Model:
-    classpath: Model:ResEncMCT
-    ResEncMCT:
+    classpath: Model:ResEnc
+    ResEnc:
       in_channels: 5
       nb_class: 12
       outputs_criterions: None
@@ -18,7 +18,7 @@ Predictor:
                 spacing:
                 - 3
                 - 3
-                - 0
+                - 3
                 inverse: true
               Standardize:
                 lazy: false