change line ending

.gitignore

@@ -1,2 +1,3 @@
 *.pth
 .idea/
+submit/

data/dataset.py

@@ -3,6 +3,7 @@ import torch
 import lmdb
 import os
 import pickle
+from PIL import Image
 from io import BytesIO
 from torch.utils.data import Dataset
 from torchvision.datasets.folder import default_loader
@@ -60,9 +61,15 @@ class LMDBDataset(Dataset):
 
     def __getitem__(self, i):
         with self.db.begin(write=False) as txn:
-            sample = torch.load(BytesIO(txn.get("{}".format(i).encode())))
+            sample = Image.open(BytesIO(txn.get("{}".format(i).encode())))
+            if sample.mode != "RGB":
+                sample = sample.convert("RGB")
             if self.transform is not None:
-                sample = self.transform(sample)
+                try:
+                    sample = self.transform(sample)
+                except RuntimeError as re:
+                    print(sample.format, sample.size, sample.mode)
+                    raise re
             return sample
 
 

data/lmdbify.py

@@ -3,34 +3,29 @@ import pickle
 from io import BytesIO
 import argparse
 
-import torch
 import lmdb
 from data.dataset import CARS, ImprovedImageFolder
-import torchvision
 from tqdm import tqdm
 
 
+def content_loader(path):
+    with open(path, "rb") as f:
+        return f.read()
+
+
 def dataset_to_lmdb(dataset, lmdb_path):
-    env = lmdb.open(lmdb_path, map_size=1099511627776, subdir=os.path.isdir(lmdb_path))
+    env = lmdb.open(lmdb_path, map_size=1099511627776*2, subdir=os.path.isdir(lmdb_path))
     with env.begin(write=True) as txn:
         for i in tqdm(range(len(dataset)), ncols=50):
-            buffer = BytesIO()
-            torch.save(dataset[i], buffer)
-            txn.put("{}".format(i).encode(), buffer.getvalue())
+            txn.put("{}".format(i).encode(), bytearray(dataset[i]))
         txn.put(b"classes_list", pickle.dumps(dataset.classes_list))
         txn.put(b"__len__", pickle.dumps(len(dataset)))
 
 
 def transform(save_path, dataset_path):
     print(save_path, dataset_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 = CARS("/data/few-shot/STANFORD-CARS/", transform=dt)
-    origin_dataset = ImprovedImageFolder(dataset_path, transform=dt)
+    # origin_dataset = CARS("/data/few-shot/STANFORD-CARS/", loader=content_loader)
+    origin_dataset = ImprovedImageFolder(dataset_path, loader=content_loader)
     dataset_to_lmdb(origin_dataset, save_path)
 
 

test.py

@@ -1,6 +1,6 @@
 import torch
 from torch.utils.data import DataLoader
-from torchvision import transforms
+import torchvision
 from data import dataset
 
 import argparse
@@ -48,7 +48,13 @@ def evaluate(query, target, support):
 
 
 def test(lmdb_path, import_path):
-    origin_dataset = dataset.LMDBDataset(lmdb_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)
     N = 5
     K = 5
     episodic_dataset = dataset.EpisodicDataset(
@@ -59,7 +65,7 @@ def test(lmdb_path, import_path):
     )
     print(episodic_dataset)
 
-    data_loader = DataLoader(episodic_dataset, batch_size=16, pin_memory=False)
+    data_loader = DataLoader(episodic_dataset, batch_size=20, pin_memory=False)
     device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 
     submit = import_module(f"submit.{import_path}")
@@ -69,11 +75,8 @@ def test(lmdb_path, import_path):
 
     accs = []
 
-    load_st = time.time()
     with torch.no_grad():
-        for item in data_loader:
-            st = time.time()
-            print("load", time.time() - load_st)
+        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
@@ -81,21 +84,17 @@ def test(lmdb_path, import_path):
             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)
-            print("compute", time.time() - st)
-            load_st = time.time()
-
-        accs.append(evaluate(query_batch, item["target"], support_batch))
+            accs.append(evaluate(query_batch, item["target"], support_batch))
     print(torch.tensor(accs).mean().item())
-    print("time: ", time.time() - st)
 
 
 if __name__ == '__main__':
     setup_seed(100)
-    defined_path = ["/data/few-shot/lmdb/mini-imagenet/val.lmdb",
-                    "/data/few-shot/lmdb/CUB_200_2011/data.lmdb",
-                    "/data/few-shot/lmdb/STANFORD-CARS/train.lmdb",
-                    # "/data/few-shot/lmdb/Plantae/data.lmdb",
-                    # "/data/few-shot/lmdb/Places365/val.lmdb"
+    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",
                     ]
     parser = argparse.ArgumentParser(description="test")
     parser.add_argument('-i', "--import_path", required=True)