Transfer Learning: How to modify the first conv2d layer of Alexnet to accomodate for 9 channel input?

aitutakiv · June 16, 2017, 9:01pm

I think the easiest way to do this is to create your model by modifying the Alexnet model from Vision so that the first conv layer has 9 channels

github.com

pytorch/vision/blob/master/torchvision/models/alexnet.py#L18


model_urls = {
'alexnet': 'https://download.pytorch.org/models/alexnet-owt-4df8aa71.pth',
}




class AlexNet(nn.Module):


def __init__(self, num_classes=1000):
    super(AlexNet, self).__init__()
    self.features = nn.Sequential(
        nn.Conv2d(3, 64, kernel_size=11, stride=4, padding=2),
        nn.ReLU(inplace=True),
        nn.MaxPool2d(kernel_size=3, stride=2),
        nn.Conv2d(64, 192, kernel_size=5, padding=2),
        nn.ReLU(inplace=True),
        nn.MaxPool2d(kernel_size=3, stride=2),
        nn.Conv2d(192, 384, kernel_size=3, padding=1),
        nn.ReLU(inplace=True),
        nn.Conv2d(384, 256, kernel_size=3, padding=1),
        nn.ReLU(inplace=True),
        nn.Conv2d(256, 256, kernel_size=3, padding=1),

Then to copy the weights over you can copy them from the original Alexnet state_dict to the state_dict of your modified model. It should be direct assignment except for the first conv layer.