resnet18 用imagenet预训练好的模型 提取fc之前的特征 仅输出数据

# -*- coding: utf-8 -*-
''' resnet18 用imagenet预训练好的模型 提取fc之前的特征 仅输出数据'''
import os.path
import torch
import torch.nn as nn
from torchvision import models, transforms
from torch.autograd import Variable
import numpy as np
from PIL import Image
import warnings
warnings.filterwarnings("ignore")

features_dir = 'D:/'

img_path = "./A.jpg"
file_name = img_path.split('/')[-1]
feature_path = os.path.join(features_dir, file_name + '.txt')

print(feature_path)

transform1 = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor()]
)

img = Image.open(img_path)
img1 = transform1(img)

# print(img1)

resnet18_feature_extractor = models.resnet18(pretrained=False)
resnet18_feature_extractor.load_state_dict(torch.load('./resnet18-f37072fd.pth'))
resnet18_feature_extractor.fc = nn.Linear(512, 512)
torch.nn.init.eye(resnet18_feature_