Change output channels of transfer learning mmar models #611

frmrz · 2022-03-10T08:39:48Z

frmrz
Mar 10, 2022

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()
)
)
)
)

Answered by yiheng-wang-nv

Mar 16, 2022

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 pretrai…

View full answer

yiheng-wang-nv · 2022-03-16T04:18:47Z

yiheng-wang-nv
Mar 16, 2022
Collaborator

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)

1 reply

frmrz Mar 18, 2022
Author

Thanks! Very clear!

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Change output channels of transfer learning mmar models #611

{{title}}

Replies: 1 comment 1 reply

{{title}}

{{editor}}'s edit

{{editor}}'s edit

{{title}}

Select a reply

Change output channels of transfer learning mmar models #611

frmrz Mar 10, 2022

Replies: 1 comment · 1 reply

yiheng-wang-nv Mar 16, 2022 Collaborator

frmrz Mar 18, 2022 Author

frmrz
Mar 10, 2022

Replies: 1 comment 1 reply

yiheng-wang-nv
Mar 16, 2022
Collaborator

frmrz Mar 18, 2022
Author