How to modify densenet121 model to take input with 1 channel?

1

I wanted to use densenet121 with grayscale images, so looking to change it to accept 1 channel input. I tried the following :

model = torchvision.models.densenet121(pretrained=pretrained)
        model.features[0] = nn.Conv2d(1, 64, kernel_size=7, stride=2, padding=3, bias=False)
        num_ftrs = model.classifier.in_features 
        model.classifier = nn.Linear(num_ftrs, n_classes)

But getting error: “RuntimeError: stack expects each tensor to be equal size, but got [1, 224, 224] at entry 0 and [4, 224, 224] at entry 17”

Is there any other way I can modify densenet121 to accept 1 channel input?

2

Your code works fine for me:

model = torchvision.models.densenet121(pretrained=True)
model.features[0] = nn.Conv2d(1, 64, kernel_size=7, stride=2, padding=3, bias=False)

x = torch.randn(2, 1, 224,224)
out = model(x)
print(out.shape)
# torch.Size([2, 1000])
1 Like