137 lines
4.1 KiB
Python
Executable File
137 lines
4.1 KiB
Python
Executable File
import torch
|
|
from torch.utils.data import DataLoader
|
|
import torchvision
|
|
from data import dataset
|
|
import torch.nn as nn
|
|
|
|
|
|
def setup_seed(seed):
|
|
torch.manual_seed(seed)
|
|
torch.cuda.manual_seed_all(seed)
|
|
torch.backends.cudnn.deterministic = True
|
|
|
|
|
|
def euclidean_dist(x, y):
|
|
"""
|
|
Compute euclidean distance between two tensors
|
|
"""
|
|
# x: N x D
|
|
# y: M x D
|
|
n = x.size(0)
|
|
m = y.size(0)
|
|
d = x.size(1)
|
|
if d != y.size(1):
|
|
raise Exception
|
|
|
|
x = x.unsqueeze(1).expand(n, m, d)
|
|
y = y.unsqueeze(0).expand(n, m, d)
|
|
|
|
return torch.pow(x - y, 2).sum(2)
|
|
|
|
|
|
def evaluate(query, target, support):
|
|
"""
|
|
:param query: NK x D vector
|
|
:param target: NK x 1 vector
|
|
:param support: N x K x D vector
|
|
:return:
|
|
"""
|
|
prototypes = support.mean(1)
|
|
distance = euclidean_dist(query, prototypes)
|
|
indices = distance.argmin(1)
|
|
return torch.eq(target, indices).float().mean()
|
|
|
|
|
|
class Flatten(nn.Module):
|
|
def __init__(self):
|
|
super(Flatten, self).__init__()
|
|
|
|
def forward(self, x):
|
|
return x.view(x.size(0), -1)
|
|
|
|
# def make_extractor():
|
|
# resnet50 = torchvision.models.resnet50(pretrained=True)
|
|
# resnet50.to(torch.device("cuda"))
|
|
# resnet50.fc = torch.nn.Identity()
|
|
# resnet50.eval()
|
|
#
|
|
# def extract(images):
|
|
# with torch.no_grad():
|
|
# return resnet50(images)
|
|
# return extract
|
|
|
|
|
|
def make_extractor():
|
|
model = resnet18()
|
|
model.to(torch.device("cuda"))
|
|
model.eval()
|
|
|
|
def extract(images):
|
|
with torch.no_grad():
|
|
return model(images)
|
|
return extract
|
|
|
|
|
|
def resnet18(model_path="ResNet18Official.pth"):
|
|
"""Constructs a ResNet-18 model.
|
|
Args:
|
|
pretrained (bool): If True, returns a model pre-trained on ImageNet
|
|
"""
|
|
model_w_fc = torchvision.models.resnet18(pretrained=False)
|
|
seq = list(model_w_fc.children())[:-1]
|
|
seq.append(Flatten())
|
|
model = torch.nn.Sequential(*seq)
|
|
# model.load_state_dict(torch.load(model_path), strict=False)
|
|
model.load_state_dict(torch.load(model_path, map_location ='cpu'), strict=False)
|
|
# model.load_state_dict(torch.load(model_path))
|
|
model.eval()
|
|
|
|
return model
|
|
|
|
|
|
def test():
|
|
data_transform = torchvision.transforms.Compose([
|
|
torchvision.transforms.Resize([int(224*1.15), int(224*1.15)]),
|
|
torchvision.transforms.CenterCrop(224),
|
|
torchvision.transforms.ToTensor(),
|
|
torchvision.transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
|
|
])
|
|
origin_dataset = dataset.CARS("/data/few-shot/STANFORD-CARS/", transform=data_transform)
|
|
|
|
episodic_dataset = dataset.EpisodicDataset(
|
|
origin_dataset, # 抽取数据集
|
|
torch.randint(5, 10, (1,)).tolist()[0], # N
|
|
torch.randint(1, 10, (1,)).tolist()[0], # K
|
|
5 # 任务数目
|
|
)
|
|
print(episodic_dataset)
|
|
|
|
data_loader = DataLoader(episodic_dataset)
|
|
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
|
|
extractor = make_extractor()
|
|
accs = []
|
|
for item in data_loader:
|
|
support_list = []
|
|
query_list = []
|
|
class_id_list = []
|
|
for class_id in item:
|
|
for i in range(len(item[class_id])):
|
|
item[class_id][i] = item[class_id][i].to(device)
|
|
num_support_set = len(item[class_id]) // 2
|
|
num_query_set = len(item[class_id]) - num_support_set
|
|
support_list.append(torch.stack([extractor(pair) for pair in item[class_id][:num_support_set]]))
|
|
query_list.append(torch.stack([extractor(pair) for pair in item[class_id][num_support_set:]]))
|
|
class_id_list.extend([class_id]*num_query_set)
|
|
|
|
query = torch.squeeze(torch.cat(query_list)).to(device)
|
|
support = torch.squeeze(torch.stack(support_list)).to(device)
|
|
target = torch.squeeze(torch.tensor(class_id_list)).to(device)
|
|
|
|
accs.append(evaluate(query, target, support))
|
|
|
|
print(accs)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
setup_seed(10)
|
|
test() |