This commit is contained in:
2020-07-23 22:32:28 +08:00
parent 3a72dcb5f0
commit ead93c1b0e
5 changed files with 78 additions and 34 deletions

33
test.py
View File

@@ -49,12 +49,11 @@ def evaluate(query, target, support):
def test(lmdb_path, import_path):
dt = torchvision.transforms.Compose([
torchvision.transforms.Resize((256, 256)),
torchvision.transforms.CenterCrop(224),
torchvision.transforms.ToTensor(),
torchvision.transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])
origin_dataset = dataset.LMDBDataset(lmdb_path, transform=dt)
origin_dataset = dataset.LMDBDataset(lmdb_path, transform=None)
N = 5
K = 5
episodic_dataset = dataset.EpisodicDataset(
@@ -65,8 +64,8 @@ def test(lmdb_path, import_path):
)
print(episodic_dataset)
data_loader = DataLoader(episodic_dataset, batch_size=20, pin_memory=False)
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
data_loader = DataLoader(episodic_dataset, batch_size=8, pin_memory=False)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
submit = import_module(f"submit.{import_path}")
@@ -78,12 +77,19 @@ def test(lmdb_path, import_path):
with torch.no_grad():
for item in tqdm(data_loader):
item = convert_tensor(item, device, non_blocking=True)
# item["query"]: B x NK x 3 x W x H
# item["support"]: B x NK x 3 x W x H
# item["query"]: B x ANK x 3 x W x H
# item["support"]: B x ANK x 3 x W x H
# item["target"]: B x NK
batch_size = item["target"].size(0)
query_batch = extractor(item["query"].view([-1, *item["query"].shape[-3:]])).view(batch_size, N * K, -1)
support_batch = extractor(item["support"].view([-1, *item["query"].shape[-3:]])).view(batch_size, N, K, -1)
image_size = item["query"].shape[-3:]
A = int(item["query"].size(1) / (N * K))
query_batch = extractor(item["query"].view([-1, *image_size])).view(batch_size, N * K, A, -1)
support_batch = extractor(item["support"].view([-1, *image_size])).view(batch_size, N, K, A, -1)
query_batch = torch.mean(query_batch, -2)
support_batch = torch.mean(support_batch, -2)
accs.append(evaluate(query_batch, item["target"], support_batch))
print(torch.tensor(accs).mean().item())
@@ -91,11 +97,12 @@ def test(lmdb_path, import_path):
if __name__ == '__main__':
setup_seed(100)
defined_path = [
"/data/few-shot/lmdb/dogs/data.lmdb",
"/data/few-shot/lmdb/flowers/data.lmdb",
"/data/few-shot/lmdb/256-object/data.lmdb",
"/data/few-shot/lmdb/dtd/data.lmdb",
]
"/data/few-shot/lmdb256/dogs.lmdb",
"/data/few-shot/lmdb256/flowers.lmdb",
"/data/few-shot/lmdb256/256-object.lmdb",
"/data/few-shot/lmdb256/dtd.lmdb",
"/data/few-shot/lmdb256/mini-imagenet-test.lmdb"
]
parser = argparse.ArgumentParser(description="test")
parser.add_argument('-i', "--import_path", required=True)
args = parser.parse_args()