前言

因为最近可用的开源代码库里找到了一个网络，因为要提高网络训练速度，网络的输入就改成了一组组特征，所以就需要提前利用特征提取模型进行特征提取，并将特征存为数组。

代码

import os
import numpy as np
from tqdm import tqdm
from PIL import Image
import torch
import torch.utils.data
def main(device, in_jit_model, out_feat_dir):



    model = torch.jit.load(in_jit_model, 'cpu')
    model.eval()
    model = model.to(device)
    f = open('.../xxx.txt', 'r')

    files = f.readlines()  # 读取整个文件所有行，保存在 list 列表中

    feats = []
    tmp = '0'
    n = 0
    for filename in files:
            filename = filename[:-1]
            image = Image.open('.../' + filename)
            image = image.convert("RGB")
            image = np.array(image)
            # image = np.expand_dims(image, axis=0)



            image = torch.tensor(image)
            image = image.float()
            image /= 255
            image = image.permute(2,1,0)
            # image = data_transform(image)
            image = torch.tensor(np.expand_dims(image, axis=0))
            B, N = [1,1]
            image = image.contiguous().to(device)
            t_y = model(image)

            assert t_y.ndim == 2

            t_y = t_y.reshape(B, N, t_y.shape[-1])

            if filename[:12]==tmp :
                feats.extend(t_y.cpu().detach().numpy())
                n = n+1
                print(filename+' get {0} feat'.format(n))

            elif tmp =='0':
                feats.extend(t_y.cpu().detach().numpy())
                n = n + 1
                print(filename+' get {0} feat'.format(n))
            else:
                files2 = open('.../xxx.txt',
                    'r')
                files2 = files2.readlines()  # 读取整个文件所有行，保存在 list 列表中
                for k in files2:
                    k = k[:-1]
                    if tmp == k[:12]:
                        feats = np.stack(feats, 0)
                        out_feat_file = f'{out_feat_dir}/{k}.svs.npy'
                        os.makedirs(os.path.dirname(out_feat_file), exist_ok=True)
                        np.save(out_feat_file, feats, allow_pickle=False)
                        n = 0
                        print(k+"----------save 1 feats")
                        feats = []
                        feats.extend(t_y.cpu().detach().numpy())
                        break



            tmp = filename[:12]



if __name__ == '__main__':

    device = 'cuda:0'
    in_jit_model = '.../xxx.pt'
    out_feat_dir = '..../...'
    main(
        device=device,
        in_jit_model=in_jit_model,
        out_feat_dir=out_feat_dir
    )