Pretrained densenet for multichannel input

[anneserious]

I noticed that it is possible to initialize and train a pretrained densenet121 with 7-channel input without error.
How is this possible? Shoudn't the pretrained weights only work for 3-channel input (since the net is pretrained on imagenet)?
What am I missing?

Kind regards
Anne

[Nic-Ma]

Hi @yiheng-wang-nv ,

Could you please help share more details to help here?

Thanks in advance.

[yiheng-wang-nv]

Hi @anneserious , as you can see from this line:
https://github.com/Project-MONAI/MONAI/blob/master/monai/networks/nets/densenet.py#L226
If you initialize a densenet with 7-channel input, its first conv layer will not load the pretrained weights due to the shapes' difference, and of course, there is no error here.

You can run the following code (the path of the pretrained weights may be different) to check if weights of the first conv layer are loaded.

from monai.networks.nets import densenet121
import torch

pre_train_file = '/root/.cache/torch/hub/checkpoints/densenet121-a639ec97.pth'
checkpoint = torch.load(pre_train_file)
print(checkpoint['features.conv0.weight'][0][0])

den121 = densenet121(pretrained=True, spatial_dims=2, in_channels=3, out_channels=2)
print(den121.features.conv0.weight[0][0])

den121_7channel = densenet121(pretrained=True, spatial_dims=2, in_channels=7, out_channels=2)
print(den121_7channel.features.conv0.weight[0][0])

[anneserious]

Hi Yiheng,

tanks for your reply, this does clear up a few things. In order to see which weights are loaded when loading the pretrained model I added a line in _load_state_dict:

    """
    This function is used to load pretrained models.
    Adapted from `PyTorch Hub 2D version
    <https://github.com/pytorch/vision/blob/master/torchvision/models/densenet.py>`_
    """
    pattern = re.compile(
        r"^(.*denselayer\d+\.(?:norm|relu|conv))\.((?:[12])\.(?:weight|bias|running_mean|running_var))$"
    )

    state_dict = load_state_dict_from_url(model_url, progress=progress)
    for key in list(state_dict.keys()):
        res = pattern.match(key)
        if res:
            new_key = res.group(1) + res.group(2)
            state_dict[new_key] = state_dict[key]
            del state_dict[key]

    model_dict = model.state_dict()
    state_dict_new = {
        k: v for k, v in state_dict.items() if ((k in model_dict) and (model_dict[k].shape == state_dict[k].shape)) 
    }
    
    
    updated_state_list=    [
        k for k, v in state_dict.items() if ((k in model_dict) )and (model_dict[k].shape == state_dict[k].shape) 
    ]
    state_dict=state_dict_new
    
    print(updated_state_list)
    model_dict.update(state_dict)
    model.load_state_dict(model_dict)

This prints out the list of loaded weights. When calling
den121 = densenet121(pretrained=True, spatial_dims=2, in_channels=3, out_channels=1000)
the resulting llist is surprisingly short:

['features.conv0.weight', 'features.norm0.weight', 'features.norm0.bias', 'features.norm0.running_mean', 'features.norm0.running_var', 'features.transition1.norm.weight', 'features.transition1.norm.bias', 'features.transition1.norm.running_mean', 'features.transition1.norm.running_var', 'features.transition1.conv.weight', 'features.transition2.norm.weight', 'features.transition2.norm.bias', 'features.transition2.norm.running_mean', 'features.transition2.norm.running_var', 'features.transition2.conv.weight', 'features.transition3.norm.weight', 'features.transition3.norm.bias', 'features.transition3.norm.running_mean', 'features.transition3.norm.running_var', 'features.transition3.conv.weight', 'features.norm5.weight', 'features.norm5.bias', 'features.norm5.running_mean', 'features.norm5.running_var']

Not sure if this is desried behavoir. There might be a different naming convention in the monai densenet and the torch densenet:
checkpoint.keys()

odict_keys(['features.conv0.weight', 'features.norm0.weight', 'features.norm0.bias', 'features.norm0.running_mean', 'features.norm0.running_var', 'features.denseblock1.denselayer1.norm.1.weight', 'features.denseblock1.denselayer1.norm.1.bias', 'features.denseblock1.denselayer1.norm.1.running_mean', 'features.denseblock1.denselayer1.norm.1.running_var', 'features.denseblock1.denselayer1.conv.1.weight', 'features.denseblock1.denselayer1.norm.2.weight', 'features.denseblock1.denselayer1.norm.2.bias', 'features.denseblock1.denselayer1.norm.2.running_mean', 'features.denseblock1.denselayer1.norm.2.running_var', 'features.denseblock1.denselayer1.conv.2.weight', 'features.denseblock1.denselayer2.norm.1.weight', 'features.denseblock1.denselayer2.norm.1.bias', 'features.denseblock1.denselayer2.norm.1.running_mean', 'features.denseblock1.denselayer2.norm.1.running_var', 'features.denseblock1.denselayer2.conv.1.weight', 'features.denseblock1.denselayer2.norm.2.weight', 'features.denseblock1.denselayer2.norm.2.bias', 'features.denseblock1.denselayer2.norm.2.running_mean', 'features.denseblock1.denselayer2.norm.2.running_var', 'features.denseblock1.denselayer2.conv.2.weight', 'features.denseblock1.denselayer3.norm.1.weight', 'features.denseblock1.denselayer3.norm.1.bias', 'features.denseblock1.denselayer3.norm.1.running_mean',...

den121.state_dict().keys()

odict_keys(['features.conv0.weight', 'features.norm0.weight', 'features.norm0.bias', 'features.norm0.running_mean', 'features.norm0.running_var', 'features.norm0.num_batches_tracked', 'features.denseblock1.denselayer1.layers.norm1.weight', 'features.denseblock1.denselayer1.layers.norm1.bias', 'features.denseblock1.denselayer1.layers.norm1.running_mean', 'features.denseblock1.denselayer1.layers.norm1.running_var', 'features.denseblock1.denselayer1.layers.norm1.num_batches_tracked', 'features.denseblock1.denselayer1.layers.conv1.weight', 'features.denseblock1.denselayer1.layers.norm2.weight', 'features.denseblock1.denselayer1.layers.norm2.bias', 'features.denseblock1.denselayer1.layers.norm2.running_mean', 'features.denseblock1.denselayer1.layers.norm2.running_var', 'features.denseblock1.denselayer1.layers.norm2.num_batches_tracked', 'features.denseblock1.denselayer1.layers.conv2.weight', 'features.denseblock1.denselayer2.layers.norm1.weight', 'features.denseblock1.denselayer2.layers.norm1.bias', 'features.denseblock1.denselayer2.layers.norm1.running_mean', 'features.denseblock1.denselayer2.layers.norm1.running_var', 'features.denseblock1.denselayer2.layers.norm1.num_batches_tracked', 'features.denseblock1.denselayer2.layers.conv1.weight', 'features.denseblock1.denselayer2.layers.norm2.weight', 'features.denseblock1.denselayer2.layers.norm2.bias', 'features.denseblock1.denselayer2.layers.norm2.running_mean', 'features.denseblock1.denselayer2.layers.norm2.running_var', 'features.denseblock1.denselayer2.layers.norm2.num_batches_tracked', 'features.denseblock1.denselayer2.layers.conv2.weight', 'features.denseblock1.denselayer3.layers.norm1.weight', 'features.denseblock1.denselayer3.layers.norm1.bias', 'features.denseblock1.denselayer3.layers.norm1.running_mean', 'features.denseblock1.denselayer3.layers.norm1.running_var',...

So most of the layers are not loaded when using the pretrained network, even for in_channel=3. It this the intended behavior, and if yes, why?

Thanks and kind regards
Anne

[yiheng-wang-nv]

Hi @anneserious , thanks for disclosing this naming issue, I checked the history and found that since this commit, the names of this network changed which lead to the issue. Let me fix it.

[yiheng-wang-nv]

Hi @anneserious , as shown in #1684 , the function has been updated to meet the latest network structure and it has been merged into the master branch. Could you please pull the lastest MONAI code, and now the pretrained weights can be loaded correctly.
To prevent the same issue, I added a new test case ( https://github.com/Project-MONAI/MONAI/blob/master/tests/test_densenet.py#L93 ) which will check the consistency between the original densenet121 (from torchvision) and our MONAI's version by input a same tensor into these two versions and compare the results.

[anneserious]

Hi @yiheng-wang-nv, awesome, thanks a lot!