test

.gitignore

@@ -0,0 +1,2 @@
+*.pth
+.idea/

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 "<EpisodicDataset N={} K={} NUM={}>".format(self.num_class, self.num_set, self.num_episodes)

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()