Cutomize VGG11 to train

vision
1

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

2

Hi,

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.

bests

3

Thanks its usefull for me