From bf201c506dba623ac2b6120e8f0e6ac149448c60 Mon Sep 17 00:00:00 2001 From: Ray Wong Date: Sun, 5 Jul 2020 23:20:45 +0800 Subject: [PATCH] test --- .gitignore | 2 + data/__init__.py | 0 data/dataset.py | 68 +++++++++++++++++++++++ test.py | 137 +++++++++++++++++++++++++++++++++++++++++++++++ 4 files changed, 207 insertions(+) create mode 100755 .gitignore create mode 100755 data/__init__.py create mode 100755 data/dataset.py create mode 100755 test.py diff --git a/.gitignore b/.gitignore new file mode 100755 index 0000000..e2994bb --- /dev/null +++ b/.gitignore @@ -0,0 +1,2 @@ +*.pth +.idea/ \ No newline at end of file diff --git a/data/__init__.py b/data/__init__.py new file mode 100755 index 0000000..e69de29 diff --git a/data/dataset.py b/data/dataset.py new file mode 100755 index 0000000..0c72fb2 --- /dev/null +++ b/data/dataset.py @@ -0,0 +1,68 @@ +from scipy.io import loadmat +import torch +from torch.utils.data import Dataset +from torchvision.datasets.folder import default_loader +from torchvision.datasets import ImageFolder +from pathlib import Path +from collections import defaultdict + + +class CARS(Dataset): + def __init__(self, root, loader=default_loader, transform=None): + self.root = Path(root) + self.transform = transform + self.loader = loader + self.annotations = loadmat(self.root / "devkit/cars_train_annos.mat")["annotations"][0] + self.annotations = {d[-1].item(): d[-2].item() - 1 for d in self.annotations} + self.classes_list = defaultdict(list) + for i in range(len(self.annotations)): + self.classes_list[self.annotations["{:05d}.jpg".format(i + 1)]].append(i) + + def __len__(self): + return len(self.annotations) + + def __getitem__(self, item): + file_name = "{:05d}.jpg".format(item+1) + target = self.annotations[file_name] + sample = self.loader(self.root / "cars_train" / file_name) + if self.transform is not None: + sample = self.transform(sample) + return sample + + +class ImprovedImageFolder(ImageFolder): + def __init__(self, root, loader=default_loader, transform=None): + super().__init__(root, transform, loader=loader) + self.classes_list = defaultdict(list) + for i in range(len(self)): + self.classes_list[self.samples[i][-1]].append(i) + + def __getitem__(self, item): + return super().__getitem__(item)[0] + + +class EpisodicDataset(Dataset): + def __init__(self, origin_dataset, num_class, num_set, num_episodes): + self.origin = origin_dataset + self.num_class = num_class + assert self.num_class < len(self.origin.classes_list) + self.num_set = num_set*2 # 2*K + self.num_episodes = num_episodes + + def __len__(self): + return self.num_episodes + + def __getitem__(self, _): + random_classes = torch.randint(high=len(self.origin.classes_list), size=(self.num_class,)).tolist() + item = {} + for i in random_classes: + image_list = self.origin.classes_list[i] + if len(image_list) > self.num_set: + idx_list = torch.randperm(len(image_list))[:self.num_set].tolist() + else: + idx_list = torch.randint(high=len(image_list), size=(self.num_set,)).tolist() + item[i] = [self.origin[idx] for idx in idx_list] + return item + + def __repr__(self): + return "".format(self.num_class, self.num_set, self.num_episodes) diff --git a/test.py b/test.py new file mode 100755 index 0000000..698815e --- /dev/null +++ b/test.py @@ -0,0 +1,137 @@ +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() \ No newline at end of file