Change output channels of transfer learning mmar models

[frmrz]

Hi everyone!
I'd like to test some of the pretrained models available through the MMAR module to perform 3D segmentation of MRI volumes. I have both a binary and a multi-label task.
Is it possible to change the last Sequential layer in the pretrained Unet models to allow for a multi-label task? How? I'll leave the print of the pretrained model below.

Thank you all!
Francesco

mmar = {
RemoteMMARKeys.ID: "clara_pt_spleen_ct_segmentation_1",
RemoteMMARKeys.NAME: "clara_pt_spleen_ct_segmentation",
RemoteMMARKeys.FILE_TYPE: "zip",
RemoteMMARKeys.HASH_TYPE: "md5",
RemoteMMARKeys.HASH_VAL: None,
RemoteMMARKeys.MODEL_FILE: os.path.join("models", "model.pt"),
RemoteMMARKeys.CONFIG_FILE: os.path.join("config", "config_train.json"),
RemoteMMARKeys.VERSION: 1,
}

unet_model = load_from_mmar(
item=mmar[RemoteMMARKeys.NAME],
mmar_dir=root_dir,
map_location=device,
pretrained=True)

UNet(
(model): Sequential(
(0): ResidualUnit(
(conv): Sequential(
(unit0): Convolution(
(conv): Conv3d(1, 16, kernel_size=(3, 3, 3), stride=(2, 2, 2), padding=(1, 1, 1))
(adn): ADN(
(N): BatchNorm3d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(D): Dropout(p=0.0, inplace=False)
(A): PReLU(num_parameters=1)
)
)
(unit1): Convolution(
(conv): Conv3d(16, 16, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))
(adn): ADN(
(N): BatchNorm3d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(D): Dropout(p=0.0, inplace=False)
(A): PReLU(num_parameters=1)
)
)
)
(residual): Conv3d(1, 16, kernel_size=(3, 3, 3), stride=(2, 2, 2), padding=(1, 1, 1))
)
(1): SkipConnection(
(submodule): Sequential(
(0): ResidualUnit(
(conv): Sequential(
(unit0): Convolution(
(conv): Conv3d(16, 32, kernel_size=(3, 3, 3), stride=(2, 2, 2), padding=(1, 1, 1))
(adn): ADN(
(N): BatchNorm3d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(D): Dropout(p=0.0, inplace=False)
(A): PReLU(num_parameters=1)
)
)
(unit1): Convolution(
(conv): Conv3d(32, 32, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))
(adn): ADN(
(N): BatchNorm3d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(D): Dropout(p=0.0, inplace=False)
(A): PReLU(num_parameters=1)
)
)
)
(residual): Conv3d(16, 32, kernel_size=(3, 3, 3), stride=(2, 2, 2), padding=(1, 1, 1))
)
(1): SkipConnection(
(submodule): Sequential(
(0): ResidualUnit(
(conv): Sequential(
(unit0): Convolution(
(conv): Conv3d(32, 64, kernel_size=(3, 3, 3), stride=(2, 2, 2), padding=(1, 1, 1))
(adn): ADN(
(N): BatchNorm3d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(D): Dropout(p=0.0, inplace=False)
(A): PReLU(num_parameters=1)
)
)
(unit1): Convolution(
(conv): Conv3d(64, 64, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))
(adn): ADN(
(N): BatchNorm3d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(D): Dropout(p=0.0, inplace=False)
(A): PReLU(num_parameters=1)
)
)
)
(residual): Conv3d(32, 64, kernel_size=(3, 3, 3), stride=(2, 2, 2), padding=(1, 1, 1))
)
(1): SkipConnection(
(submodule): Sequential(
(0): ResidualUnit(
(conv): Sequential(
(unit0): Convolution(
(conv): Conv3d(64, 128, kernel_size=(3, 3, 3), stride=(2, 2, 2), padding=(1, 1, 1))
(adn): ADN(
(N): BatchNorm3d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(D): Dropout(p=0.0, inplace=False)
(A): PReLU(num_parameters=1)
)
)
(unit1): Convolution(
(conv): Conv3d(128, 128, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))
(adn): ADN(
(N): BatchNorm3d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(D): Dropout(p=0.0, inplace=False)
(A): PReLU(num_parameters=1)
)
)
)
(residual): Conv3d(64, 128, kernel_size=(3, 3, 3), stride=(2, 2, 2), padding=(1, 1, 1))
)
(1): SkipConnection(
(submodule): ResidualUnit(
(conv): Sequential(
(unit0): Convolution(
(conv): Conv3d(128, 256, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))
(adn): ADN(
(N): BatchNorm3d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(D): Dropout(p=0.0, inplace=False)
(A): PReLU(num_parameters=1)
)
)
(unit1): Convolution(
(conv): Conv3d(256, 256, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))
(adn): ADN(
(N): BatchNorm3d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(D): Dropout(p=0.0, inplace=False)
(A): PReLU(num_parameters=1)
)
)
)
(residual): Conv3d(128, 256, kernel_size=(1, 1, 1), stride=(1, 1, 1))
)
)
(2): Sequential(
(0): Convolution(
(conv): ConvTranspose3d(384, 64, kernel_size=(3, 3, 3), stride=(2, 2, 2), padding=(1, 1, 1), output_padding=(1, 1, 1))
(adn): ADN(
(N): BatchNorm3d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(D): Dropout(p=0.0, inplace=False)
(A): PReLU(num_parameters=1)
)
)
(1): ResidualUnit(
(conv): Sequential(
(unit0): Convolution(
(conv): Conv3d(64, 64, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))
(adn): ADN(
(N): BatchNorm3d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(D): Dropout(p=0.0, inplace=False)
(A): PReLU(num_parameters=1)
)
)
)
(residual): Identity()
)
)
)
)
(2): Sequential(
(0): Convolution(
(conv): ConvTranspose3d(128, 32, kernel_size=(3, 3, 3), stride=(2, 2, 2), padding=(1, 1, 1), output_padding=(1, 1, 1))
(adn): ADN(
(N): BatchNorm3d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(D): Dropout(p=0.0, inplace=False)
(A): PReLU(num_parameters=1)
)
)
(1): ResidualUnit(
(conv): Sequential(
(unit0): Convolution(
(conv): Conv3d(32, 32, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))
(adn): ADN(
(N): BatchNorm3d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(D): Dropout(p=0.0, inplace=False)
(A): PReLU(num_parameters=1)
)
)
)
(residual): Identity()
)
)
)
)
(2): Sequential(
(0): Convolution(
(conv): ConvTranspose3d(64, 16, kernel_size=(3, 3, 3), stride=(2, 2, 2), padding=(1, 1, 1), output_padding=(1, 1, 1))
(adn): ADN(
(N): BatchNorm3d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(D): Dropout(p=0.0, inplace=False)
(A): PReLU(num_parameters=1)
)
)
(1): ResidualUnit(
(conv): Sequential(
(unit0): Convolution(
(conv): Conv3d(16, 16, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))
(adn): ADN(
(N): BatchNorm3d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(D): Dropout(p=0.0, inplace=False)
(A): PReLU(num_parameters=1)
)
)
)
(residual): Identity()
)
)
)
)
(2): Sequential(
(0): Convolution(
(conv): ConvTranspose3d(32, 2, kernel_size=(3, 3, 3), stride=(2, 2, 2), padding=(1, 1, 1), output_padding=(1, 1, 1))
(adn): ADN(
(N): BatchNorm3d(2, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(D): Dropout(p=0.0, inplace=False)
(A): PReLU(num_parameters=1)
)
)
(1): ResidualUnit(
(conv): Sequential(
(unit0): Convolution(
(conv): Conv3d(2, 2, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))
)
)
(residual): Identity()
)
)
)
)

[yiheng-wang-nv]

Hi @frmrz , sorry for late reply.

Assume the actual Unet model you need has some layers that are different from the pretrained weights, you need to modify the state dict. For example, if you need a Unet that has 5 channels output, the code is like:

unet_model = load_from_mmar(
    item=mmar[RemoteMMARKeys.NAME],
    mmar_dir=root_dir,
    map_location=device,
    pretrained=True)

model = UNet(
    spatial_dims=3,
    in_channels=1,
    out_channels=5,
    channels=(16, 32, 64, 128, 256),
    strides=(2, 2, 2, 2),
    num_res_units=2,
    norm=Norm.BATCH,
)

model_dict = model.state_dict()

pretrain_state_dict = unet_model.state_dict()

pretrain_state_dict = {
    k: v for k, v in pretrain_state_dict.items() if (k in model_dict) and (model_dict[k].shape == pretrain_state_dict[k].shape)
}

# # uncomment the following lines to show the printed results
# for key in model_dict.keys():
#     if key not in pretrain_state_dict.keys():
#         print("the key {} will not be loaded.".format(key))

model_dict.update(pretrain_state_dict)
model.load_state_dict(model_dict)

[frmrz]

Thanks! Very clear!