Cutomize VGG11 to train

I would like to customize the VGG 11 model so that it has an input channel of 1, that is, the images to train are grayscale, please help me


I think easiest way is too define a super class and change only first layer, something like this:

import torch
import torch.nn as nn
import torchvision

class MyVGG(nn.Module):
    def __init__(self, in_channel=1):
        super(MyVGG, self).__init__()
        # use PyTorch implementation
        self.model = torchvision.models.vgg11()
        # replace first layer of model with your own desired
        self.model.features[0] = nn.Conv2d(1, 64, 3, 1, 1)

    def forward(self, x):
        return self.model(x)

myvgg = MyVGG()

Another approach is to use the source code of models here and change this line to in_channels = 1.

This thread may also help you.


Thanks its usefull for me