pytorch提取softmax前的特征并保存为txt文件

最近这段时间，一直在找一些有用的传统特征想用在我的下一个工作里面，但是发现很难找到一个合适的高维特征来解决手头的问题。然后就想能不能使用深度网络xceptionet，resnet以及densenet等来训练，将全连接层前的输入tensor（一般是1024D或者是2048D，当然多少维的都有）当做深度特征，然后再把这些特征用在后续的one-class SVM以及聚类中。
首先想到，能不能在深度学习的代码里面直接用接后续的操作（one-class SVM以及聚类），使用python实现后续的这些操作，但是这个显然有点难。因为平时习惯了（依赖） 使用MATLAB，直接调用一些工具，可以轻松实现这些目的。所以这时候，另外一条路就是，把想要的tensor保存到txt中输出来，再用MATLAB读取这些txt文件，后续操作就比较容易处理了。
查了一些资料，并进行了一些修改，形成了亲测可用的代码，话不多说先上码，代码后有比较详细的解释，希望能够帮到大家。

# -*- coding: utf-8 -*-
import os, torch, glob
import numpy as np
from torch.autograd import Variable
from PIL import Image  
from torchvision import models, transforms
import torch.nn as nn
import shutil
from densenet import *
features_dir = '/home/yaohaichao/densenet/features/'
 
def extractor(img_path, saved_path, net, use_gpu):
    transform = transforms.Compose([
            transforms.CenterCrop(224),
            transforms.ToTensor()]
    )
    
    img = Image.open(img_path)
    img = transform(img)
    
    x = Variable(torch.unsqueeze(img, dim=0).float(), requires_grad=False)
    if use_gpu:
        x = x.cuda()
        net = net.cuda()
    y = net(x).cpu()
    y = np.array(y)
    #print(y)
    #y = y.data.numpy()
    y = np.around(y, decimals=2)

    np.savetxt(saved_path, y, delimiter='\n')
    
if __name__ == '__main__':
    files_list = []
    files_list.extend(glob.glob(r"/home/yaohaichao/Dataset/TEST_qp10_obb/*.png"))
    
        
    use_gpu = torch.cuda.is_available()
    densenet_feature_extractor = densenet121(num_classes = 1024) #要提前知道你使用的网络在全连接层的输入维度即特征维度
    if use_gpu:
        densenet_feature_extractor = densenet_feature_extractor.cuda()
        densenet_feature_extractor = nn.DataParallel(densenet_feature_extractor)
   
    pretrained_dict = torch.load('/home/yaohaichao/densenet/output/epoch_29.pth')
    model_dict = densenet_feature_extractor.state_dict()
    pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict}
    del pretrained_dict["module.classifier.weight"]
    del pretrained_dict["module.classifier.bias"]
    model_dict.update(pretrained_dict)
    densenet_feature_extractor.load_state_dict(model_dict)
    densenet_feature_extractor.classifier = nn.Linear(1024, 1024)
    torch.nn.init.eye_(densenet_feature_extractor.classifier.weight)
    
    densenet_feature_extractor.eval()
    for param in densenet_feature_extractor.parameters():
        param.requires_grad = False   

    for x_path in files_list:
        print(x_path)
        fx_path = os.path.join(features_dir, x_path[39:-4] + '.txt')
        extractor(x_path, fx_path, densenet_feature_extractor, use_gpu)

依次将”/home/yaohaichao/Dataset/TEST_qp10_obb/*.png”文件夹下的png图像读取出文件名形成一个list，加载已经训练好的pth模型，这里要注意，我们认为提特征这个过程是在测试阶段的，是在模型训练好了之后，在进行特征的提取。

    pretrained_dict = torch.load('/home/yaohaichao/densenet/output/epoch_29.pth')
    model_dict = densenet_feature_extractor.state_dict()
    pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict}
    del pretrained_dict["module.classifier.weight"]
    del pretrained_dict["module.classifier.bias"]
    model_dict.update(pretrained_dict)
    densenet_feature_extractor.load_state_dict(model_dict)
    densenet_feature_extractor.classifier = nn.Linear(1024, 1024)

这段代码，是将已有模型加载进来，module.classifier.weight和module.classifier.bias在不同的模型里面也是不一样的，有的是module.fc.bias和module.fc.weight,这跟模型的定义有关系。
比较关键的是下面这两句：

densenet_feature_extractor.classifier = nn.Linear(1024, 1024)
torch.nn.init.eye_(densenet_feature_extractor.classifier.weight)

我们如何将1024的特征向量输出？这里的技巧是，将全连接层设置为1024×1024，并且将全连接层的weight设置成单位矩阵，这就使得这个网络的输出就是我们想要的特征。
如此，直接保存网络的输出就可以了。
注意，由于我要将保存的txt改名与相对应的图片一致，所以这里用了

fx_path = os.path.join(features_dir, x_path[39:-4] + '.txt')

x_path[39:-4]是在整个文件路径的字符串中提取了文件名。