pytorch可视化网络层特征

Inception V3

def preprocess_image(im, resize_im=True):
    mean = [0.485, 0.456, 0.406]
    std = [0.229, 0.224, 0.225]
    # Resize image
    if resize_im:
        cv2im = cv2.resize(im, (224, 224))
    im_as_arr = np.float32(im)
    im_as_arr = np.ascontiguousarray(im_as_arr[..., ::-1])
    im_as_arr = im_as_arr.transpose(2, 0, 1)  # Convert array to D,W,H
    # 正则化
    for channel, _ in enumerate(im_as_arr):
        im_as_arr[channel] /= 255
        im_as_arr[channel] -= mean[channel]
        im_as_arr[channel] /= std[channel]
    # 转为float
    im_as_ten = torch.from_numpy(im_as_arr).float()
    # 增加一维在起始. Tensor shape = 1,3,224,224
    im_as_ten.unsqueeze_(0)
    # 转为Variable
    im_as_var = Variable(im_as_ten, requires_grad=True)
    return im_as_var


class FeatureVisualization():
    def __init__(self,img_path,selected_layer):
        self.img_path=img_path
        self.selected_layer=selected_layer
        self.pretrained_model = models.inception_v3(pretrained=True, aux_logits=False)
        self.pretrained_model.fc = nn.Linear(2048, 2)
        self.pretrained_model.load_state_dict(torch.load('model.weight'))
        for p in self.pretrained_model.parameters():
            p.requires_grad = False
        self.pretrained_model = nn.Sequential(*list(self.pretrained_model.children())[:16])

    def process_image(self):
        img=cv2.imread(self.img_path)
        img=preprocess_image(img)
        return img

    def get_feature(self):
        # input = Variable(torch.randn(1, 3, 224, 224))
        input=self.process_image()
        print(input.shape)
        x=input
        for index,layer in enumerate(self.pretrained_model):
            x=layer(x)
            if (index == self.selected_layer):
                return x

    def get_single_feature(self):
        features=self.get_feature()
        print(features.shape)

        feature=features[:,0,:,:]
        print(feature.shape)

        feature=feature.view(feature.shape[1],feature.shape[2])
        print(feature.shape)

        return feature

    def save_feature_to_img(self):
        #to numpy
        feature=self.get_single_feature()
        feature=feature.data.numpy()

        #sigmod
        feature= 1.0/(1+np.exp(-1*feature))

        # to [0,255]
        feature=np.round(feature*255)
        print(feature[0])
        cv2.imwrite('./15_img.jpg',feature)


if __name__=='__main__':
    # get class
    myClass=FeatureVisualization('100105_30.000.jpg',15)
    print (myClass.pretrained_model)
    myClass.save_feature_to_img()