File size: 1,491 Bytes
6ef5ffa | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 | """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])
|