| 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(): |
| if hasattr(child, "skip") and callable(child.skip) and not isinstance(child.skip, nn.Module): |
| child.skip = nn.Identity() |
| if hasattr(child, "nonlin2") and callable(child.nonlin2) and not isinstance(child.nonlin2, nn.Module): |
| child.nonlin2 = nn.Identity() |
| |
| class ResEnc(network.Network): |
| def __init__( |
| self, |
| optimizer: network.OptimizerLoader = network.OptimizerLoader(), |
| schedulers: dict[str, network.LRSchedulersLoader] = { |
| "default:ReduceLROnPlateau": network.LRSchedulersLoader(0) |
| }, |
| outputs_criterions: dict[str, network.TargetCriterionsLoader] = {"default": network.TargetCriterionsLoader()}, |
| in_channels: int = 5, |
| nb_class: int = 132, |
| ) -> None: |
| super().__init__( |
| in_channels=in_channels, |
| optimizer=optimizer, |
| schedulers=schedulers, |
| outputs_criterions=outputs_criterions, |
| dim=2, |
| ) |
| self.add_module("DecoderOutputs", ResidualEncoderUNet( |
| 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) |
