| """Encoder construction matching nnssl's ResEncL preset, not a Transformers AutoModel. |
| |
| Architecture reference: MIC-DKFZ/nnssl, architectures/architecture_registry.py. |
| Its local source carries CC-BY-SA-4.0; this does not assign a license to weights. |
| """ |
| from pathlib import Path |
| import torch |
| from torch import nn |
| from dynamic_network_architectures.architectures.unet import ResidualEncoderUNet |
| from safetensors.torch import load_model |
|
|
| def build_encoder(): |
| network = ResidualEncoderUNet( |
| input_channels=1, n_stages=6, |
| features_per_stage=[32,64,128,256,320,320],conv_op=nn.Conv3d, |
| kernel_sizes=[[3,3,3]]*6, |
| strides=[[1,1,1]]+[[2,2,2]]*5, |
| n_blocks_per_stage=[1,3,4,6,6,6],num_classes=1, |
| n_conv_per_stage_decoder=[1]*5,conv_bias=True, |
| norm_op=nn.InstanceNorm3d,norm_op_kwargs={'eps':1e-5,'affine':True}, |
| nonlin=nn.LeakyReLU,nonlin_kwargs={'inplace':True},deep_supervision=False, |
| ) |
| return network.encoder.eval() |
|
|
| def load_encoder(directory): |
| encoder=build_encoder() |
| load_model(encoder,str(Path(directory)/'encoder.safetensors'),strict=True) |
| return encoder.eval() |
|
|
| if __name__=='__main__': |
| import sys |
| torch.set_num_threads(2) |
| model=load_encoder(sys.argv[1]) |
| with torch.inference_mode(): |
| out=model(torch.randn(1,1,64,64,64)) |
| tensors=out if isinstance(out,(list,tuple)) else [out] |
| assert all(torch.isfinite(x).all() for x in tensors) |
| print([list(x.shape) for x in tensors]) |
|
|
