How to modify the final FC layer based on DenseNet？

[shenao1995]

I want to add an SPP(Spatial Pyramid Pooling) in front of densenet FC, but I can't start. I would like to ask how to modify it. Thank you!

[rijobro]

Is this what you had in mind?

import torch.nn as nn
from monai.networks.blocks import SimpleASPP
from monai.networks.nets import DenseNet
import types

im_size = (1,) + tuple(train_ds[0]["image"].shape)

# Create model
model = DenseNet(spatial_dims=2, in_channels=1, out_channels=6,
                 init_features=2, growth_rate=2, block_config=(2,)).to(device)


# Print summary: before
summary(model, im_size[1:])

# Figure out number of channels leading up to ASPP
rand_im = torch.randn(im_size,device=device)
out_features = model.features(rand_im)
# Create ASPP
model.aspp = SimpleASPP(
    spatial_dims=2, in_channels=out_features.shape[1], conv_out_channels=6).to(device)
# Get number of channels after ASPP
out_aspp_channels = model.aspp(out_features).shape[1]
# Modify linear layer accordingly
out_lin_channels = model.class_layers[-1].out_features
model.class_layers[-1] = nn.Linear(out_aspp_channels, out_lin_channels).to(device)

# Modify forward method
def forward_w_aspp(self, x):
    x = self.features(x)
    x = self.aspp(x)
    x = self.class_layers(x)
    return x
model.forward = types.MethodType(forward_w_aspp, model)

# Print summary: after
summary(model, im_size[1:])

[shenao1995]

Thank you very much!
I modified it according to your code. My purpose is to input images of different sizes, but an error will be reported in the dataloader. How can I solve this problem. Do I want to rewrite dataset?
RuntimeError: stack expects each tensor to be equal size, but got [1, 78, 104] at entry 0 and [1, 25, 20] at entry 1

[rijobro]

To ensure that your images are all the same size, and that they are all a size that is acceptable for the network, think about using a transform such as ResizeWithPadOrCrop. ResizeWithPadOrCrop will either crop or pad your images to a desired size. As to what size the images need to be, that's something you'll have to figure out based on your network shape.

Also thinking about it a little more, the following code snippet is slightly cleaner than the first example. Inherit from DenseNet and create a new class:

class DenseNetASPP(DenseNet):
    def __init__(
        self,
        spatial_dims,
        in_channels,
        out_channels,
        init_features = 64,
        growth_rate = 32,
        block_config = (6, 12, 24, 16),
        bn_size = 4,
        dropout_prob = 0.0,
        aspp_conv_out_channels = 6,
    ):
        # initialise normal densenet
        super().__init__(
            spatial_dims, in_channels, out_channels,
            init_features, growth_rate, block_config,
            bn_size, dropout_prob,
        )
        # create aspp module
        aspp_in_features = self.class_layers[-1].in_features
        self.aspp = SimpleASPP(
            spatial_dims=spatial_dims, in_channels=aspp_in_features, 
            conv_out_channels=aspp_conv_out_channels
        )
        # replace last linear component with updated number of input channels
        aspp_out_features = self.aspp.conv_k1.out_channels
        lin_out_channels = model.class_layers[-1].out_features
        self.class_layers[-1] = nn.Linear(aspp_out_features, lin_out_channels)
        
    def forward(self, x):
        x = self.features(x)
        x = self.aspp(x)
        x = self.class_layers(x)
        return x
        

[shenao1995]

Thank you very much for your help. I trained after adjusting the code, but the loss gradient did not drop. Then, I changed the parameters, and the loss gradient did not decrease.

[Vivekrpg]

I have a question related to the final layer in densenet model implementation in monai.
densenet model has a final stage as a classifier which is linear ( same as proposed in the paper). There are a some implementations of densenet that has a softmax as the last layer. A couple of questions around this?

• how to attach/replace linear classifier stage with softmax in monai platform ?
• are there scenarios when one method is preferred over the other?