3

我想在 128x128 大小的图像上训练 VGG。我不想将它们重新缩放到 224x224 以节省 GPU 内存和训练时间。这样做的正确方法是什么?

4

1 回答 1

3

最好的方法是保持卷积部分不变,替换全连接层。这样,甚至可以为网络的卷积部分采用预训练的权重。全连接层必须随机初始化。通过这种方式,可以微调具有较小输入大小的网络。

这里有一些pytorch代码

import torch
from torch.autograd import Variable
import torchvision
import torch.nn as nn

from torchvision.models.vgg import model_urls

VGG_TYPES = {'vgg11' : torchvision.models.vgg11, 
             'vgg11_bn' : torchvision.models.vgg11_bn, 
             'vgg13' : torchvision.models.vgg13, 
             'vgg13_bn' : torchvision.models.vgg13_bn, 
             'vgg16' : torchvision.models.vgg16, 
             'vgg16_bn' : torchvision.models.vgg16_bn,
             'vgg19_bn' : torchvision.models.vgg19_bn, 
             'vgg19' : torchvision.models.vgg19}


class Custom_VGG(nn.Module):

    def __init__(self,
                 ipt_size=(128, 128), 
                 pretrained=True, 
                 vgg_type='vgg19_bn', 
                 num_classes=1000):
        super(Custom_VGG, self).__init__()

        # load convolutional part of vgg
        assert vgg_type in VGG_TYPES, "Unknown vgg_type '{}'".format(vgg_type)
        vgg_loader = VGG_TYPES[vgg_type]
        vgg = vgg_loader(pretrained=pretrained)
        self.features = vgg.features

        # init fully connected part of vgg
        test_ipt = Variable(torch.zeros(1,3,ipt_size[0],ipt_size[1]))
        test_out = vgg.features(test_ipt)
        self.n_features = test_out.size(1) * test_out.size(2) * test_out.size(3)
        self.classifier = nn.Sequential(nn.Linear(self.n_features, 4096),
                                        nn.ReLU(True),
                                        nn.Dropout(),
                                        nn.Linear(4096, 4096),
                                        nn.ReLU(True),
                                        nn.Dropout(),
                                        nn.Linear(4096, num_classes)
                                       )
        self._init_classifier_weights()

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

    def _init_classifier_weights(self):
        for m in self.classifier:
            if isinstance(m, nn.Linear):
                m.weight.data.normal_(0, 0.01)
                m.bias.data.zero_()

要创建一个 vgg,只需调用它:

vgg = Custom_VGG(ipt_size=(128, 128), pretrained=True)
于 2018-01-29T04:24:30.520 回答