Как преобразовать предварительно обученные слои FC в слои CONV в Pytorch

Question

Я хочу преобразовать предварительно обученный CNN (например, VGG-16) в полностью сверточную сеть в Pytorch. Как я могу это сделать?

Любой совет будет полезен.

Answer 1

Надеюсь, еще не поздно Я не совсем уверен, какие части вы хотите сохранить, но если вы хотите манипулировать предварительно обученной моделью или повторно использовать ее части, мой подход будет следующим:

1. Загрузите предварительно обученную модель, например,

   import torch
   import torch.nn as nn
   from torchvision import models
   
   model = models.resnet101(pretrained=True)
   

2. Извлеките интересующие вас детали и создайте новую модель из этих частей, например

   list(model.modules()) # to inspect the modules of your model
   my_model = nn.Sequential(*list(model.modules())[:-1]) # strips off last linear layer
   

Конечно, вы можете извлекать и повторно использовать любые части модели, которые вам могут понадобиться, а также добавлять новые модули, просто измените список соответствующим образом.

Answer 2

Вы можете сделать это следующим образом (см. Комментарии для описания):

import torch
import torch.nn as nn
from torchvision import models

# 1. LOAD PRE-TRAINED VGG16
model = models.vgg16(pretrained=True)

# 2. GET CONV LAYERS
features = model.features

# 3. GET FULLY CONNECTED LAYERS
fcLayers = nn.Sequential(
    # stop at last layer
    *list(model.classifier.children())[:-1]
)

# 4. CONVERT FULLY CONNECTED LAYERS TO CONVOLUTIONAL LAYERS

### convert first fc layer to conv layer with 512x7x7 kernel
fc = fcLayers[0].state_dict()
in_ch = 512
out_ch = fc["weight"].size(0)

firstConv = nn.Conv2d(in_ch, out_ch, 7, 7)

### get the weights from the fc layer
firstConv.load_state_dict({"weight":fc["weight"].view(out_ch, in_ch, 7, 7),
                           "bias":fc["bias"]})

# CREATE A LIST OF CONVS
convList = [firstConv]

# Similarly convert the remaining linear layers to conv layers 
for layer in enumerate(fcLayers[1:]):
    if isinstance(module, nn.Linear):
        # Convert the nn.Linear to nn.Conv
        fc = module.state_dict()
        in_ch = fc["weight"].size(1)
        out_ch = fc["weight"].size(0)
        conv = nn.Conv2d(in_ch, out_ch, 1, 1)

        conv.load_state_dict({"weight":fc["weight"].view(out_ch, in_ch, 1, 1),
            "bias":fc["bias"]})

        convList += [conv]
    else:
        # Append other layers such as ReLU and Dropout
        convList += [layer]

# Set the conv layers as a nn.Sequential module
convLayers = nn.Sequential(*convList)