Pytorch提取模型特征向量
# -*- coding: utf-8 -*-
"""
dj
"""
import torch
import torch.nn as nn
import os
from torchvision import models, transforms
from torch.autograd import Variable
import numpy as np
from PIL import Image
import torchvision.models as models
import pretrainedmodels
import pandas as pd
class FCViewer(nn.Module):
def forward(self, x):
return x.view(x.size(0), -1)
class M(nn.Module):
def __init__(self, backbone1, drop, pretrained=True):
super(M,self).__init__()
if pretrained:
img_model = pretrainedmodels.__dict__[backbone1](num_classes=1000, pretrained='imagenet')
else:
img_model = pretrainedmodels.__dict__[backbone1](num_classes=1000, pretrained=None)
self.img_encoder = list(img_model.children())[:-2]
self.img_encoder.append(nn.AdaptiveAvgPool2d(1))
self.img_encoder = nn.Sequential(*self.img_encoder)
if drop > 0:
self.img_fc = nn.Sequential(FCViewer())
else:
self.img_fc = nn.Sequential(
FCViewer())
def forward(self, x_img):
x_img = self.img_encoder(x_img)
x_img = self.img_fc(x_img)
return x_img
model1=M('resnet18',0,pretr