Pytorch сохранение и загрузка VGG16 с передачей знаний

1

Я сохраняю VGG16 с передачей знаний, используя следующую инструкцию:

torch.save(model.state_dict(), 'checkpoint.pth')

и перезагрузку, используя следующий оператор:

state_dict = torch.load('checkpoint.pth') model.load_state_dict (state_dict)

Это работает до тех пор, пока я перезагружаю модель VGG16 и даю ей те же настройки, что и раньше, со следующим кодом:

model = models.vgg16(pretrained=True)
model.cuda()
for param in model.parameters(): param.requires_grad = False

class Network(nn.Module):
    def __init__(self, input_size, output_size, hidden_layers, drop_p=0.5):

#             input_size: integer, size of the input
#             output_size: integer, size of the output layer
#             hidden_layers: list of integers, the sizes of the hidden layers
#             drop_p: float between 0 and 1, dropout probability

        super().__init__()
        # Add the first layer, input to a hidden layer
        self.hidden_layers = nn.ModuleList([nn.Linear(input_size, hidden_layers[0])])

        # Add a variable number of more hidden layers
        layer_sizes = zip(hidden_layers[:-1], hidden_layers[1:])
        self.hidden_layers.extend([nn.Linear(h1, h2) for h1, h2 in layer_sizes])
        self.output = nn.Linear(hidden_layers[-1], output_size)
        self.dropout = nn.Dropout(p=drop_p)

    def forward(self, x):
        ''' Forward pass through the network, returns the output logits '''

        # Forward through each layer in 'hidden_layers', with ReLU activation and dropout
        for linear in self.hidden_layers:
            x = F.relu(linear(x))
            x = self.dropout(x)

        x = self.output(x)
        return F.log_softmax(x, dim=1)

classifier = Network(25088, 102, [4096], drop_p=0.5)
model.classifier = classifier

Как этого избежать? Как я могу перезагрузить модель без перезагрузки VGG16 и переопределить классификатор?

Теги:
pytorch
vgg-net

1 ответ

1
Лучший ответ

Почему бы не переопределить модель VGG16 напрямую? просмотреть vgg.py для подробностей

class VGG_New(nn.Module):
    def __init__(self, features, num_classes=1000, init_weights=True):
        super(VGG, self).__init__()
        self.features = features
        # change here with you code
        self.classifier = nn.Sequential(
            nn.Linear(512 * 7 * 7, 4096),
            nn.ReLU(True),
            nn.Dropout(),
            nn.Linear(4096, 4096),
            nn.ReLU(True),
            nn.Dropout(),
            nn.Linear(4096, num_classes),
        )
        if init_weights:
            self._initialize_weights()

    def forward(self, x):
        x = self.features(x)
        x = x.view(x.size(0), -1)
        x = self.classifier(x)
        return x

    def _initialize_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
                if m.bias is not None:
                    nn.init.constant_(m.bias, 0)
            elif isinstance(m, nn.BatchNorm2d):
                nn.init.constant_(m.weight, 1)
                nn.init.constant_(m.bias, 0)
            elif isinstance(m, nn.Linear):
                nn.init.normal_(m.weight, 0, 0.01)
                nn.init.constant_(m.bias, 0)

тогда вес груза для функций только

pretrained_dict=torch.load(vgg_weight)
model_dict=vgg_new.state_dict()
# 1. filter out unnecessary keys
pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict}
# or filter with key value
# pretrained_dict = {k: v for k, v in pretrained_dict.items() if k.find('classifier')==-1}
# 2. overwrite entries in the existing state dict
model_dict.update(pretrained_dict)
vgg_new.load_state_dict(model_dict)
  • 0
    также может показать код для сохранения веса?

Ещё вопросы

Сообщество Overcoder
Наверх
Меню