pytorch获取vgg16-feature层输出

实际应用时可能比较想获取VGG中间层的输出,那么就可以如下操作:

import numpy as np
import torch
from torchvision import models
from torch.autograd import Variable
import torchvision.transforms as transforms


class CNNShow():
    def __init__(self, model):
        self.model = model
        self.model.eval()

        self.created_image = self.image_for_pytorch(np.uint8(np.random.uniform(150, 180, (224, 224, 3))))


    def show(self):
        x = self.created_image
        for index, layer in enumerate(self.model):
            print(index,layer)
            x = layer(x)

    def image_for_pytorch(self,Data):
        transform = transforms.Compose([
            transforms.ToTensor(),  # range [0, 255] -> [0.0,1.0]
            transforms.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225))
        ]
        )
        imData = transform(Data)
        imData = Variable(torch.unsqueeze(imData, dim=0), requires_grad=True)
        return imData

if __name__ == '__main__':

    pretrained_model = models.vgg16(pretrained=True).features
    CNN = CNNShow(pretrained_model)
    CNN.show()

猜你喜欢

转载自blog.csdn.net/emiedon/article/details/79574113