Pytorch，预训练模型：如何同时使用特征和分类器

Question

我正在使用vgg16提取图像特征向量。我想从倒数第二层得到1 1 4096个向量。

我的代码：

def get_model():
    model = models.vgg16(pretrained=True)#.features[:].classifier[:4]
    model = model.eval()
    # model.cuda()  # send the model to GPU, DO NOT include this line if you haven't a GPU
    return model

但是我只能从最后一层得到1 1 1000个向量。

我知道如何使用feathers和classifier，但我不知道如何同时使用它们。

仅使用分类器：

仅使用羽毛：

同时使用它们：

日志：

Traceback (most recent call last):
  File "/mnt/c/Users/sunji/PycharmProjects/image_cluster_pytorch/main.py", line 7, in <module>
    model = calc.get_model()
  File "/mnt/c/Users/sunji/PycharmProjects/image_cluster_pytorch/imagecluster/calc.py", line 17, in get_model
    model = models.vgg16(pretrained=True).features[:].classifier[:4]
  File "/usr/local/lib/python3.8/dist-packages/torch/nn/modules/module.py", line 771, in __getattr__
    raise ModuleAttributeError("'{}' object has no attribute '{}'".format(
torch.nn.modules.module.ModuleAttributeError: 'Sequential' object has no attribute 'classifier'

Answer 1

我自己找到了解决方案。

很抱歉麻烦堆栈溢出。

代码如下：

def get_model():
    model = models.vgg16(pretrained=True)
    model.features = model.features[:]
    model.classifier = model.classifier[:4]

    model = model.eval()
    # model.cuda()  # send the model to GPU, DO NOT include this line if you haven't a GPU
    return model

结果：

我认为这是正确的答案。