fixes

inference.py

@@ -82,8 +82,11 @@ class InferenceService:
     def embed_images(self, images: List[Image.Image]) -> List[np.ndarray]:
         if len(images) == 0:
             return []
-        batch = torch.stack([self.transform(img) for img in images]).to(self.device)
-        emb = self.resnet(batch)
+        batch = torch.stack([self.transform(img) for img in images])
+        batch = batch.to(self.device, memory_format=torch.channels_last)
+        use_amp = (self.device == "cuda")
+        with torch.autocast(device_type=("cuda" if use_amp else "cpu"), enabled=use_amp):
+            emb = self.resnet(batch)
         emb = nn.functional.normalize(emb, dim=-1)
         return [e.detach().cpu().numpy().astype(np.float32) for e in emb]
 

train_resnet.py

@@ -27,10 +27,12 @@ def parse_args() -> argparse.Namespace:
 def main() -> None:
     args = parse_args()
     device = "cuda" if torch.cuda.is_available() else ("mps" if torch.backends.mps.is_available() else "cpu")
+    if device == "cuda":
+        torch.backends.cudnn.benchmark = True
 
     dataset = PolyvoreTripletDataset(args.data_root, split="train")
 
-    loader = DataLoader(dataset, batch_size=args.batch_size, shuffle=True, num_workers=4, pin_memory=True)
+    loader = DataLoader(dataset, batch_size=args.batch_size, shuffle=True, num_workers=4, pin_memory=(device=="cuda"))
     model = ResNetItemEmbedder(embedding_dim=args.embedding_dim).to(device)
     optimizer = optim.AdamW(model.parameters(), lr=args.lr, weight_decay=1e-4)
     criterion = nn.TripletMarginLoss(margin=0.2, p=2)
@@ -45,12 +47,13 @@ def main() -> None:
         for batch in loader:
             # Expect batch as (anchor, positive, negative)
             anchor, positive, negative = batch
-            anchor = anchor.to(device)
-            positive = positive.to(device)
-            negative = negative.to(device)
-            emb_a = model(anchor)
-            emb_p = model(positive)
-            emb_n = model(negative)
+            anchor = anchor.to(device, memory_format=torch.channels_last, non_blocking=True)
+            positive = positive.to(device, memory_format=torch.channels_last, non_blocking=True)
+            negative = negative.to(device, memory_format=torch.channels_last, non_blocking=True)
+            with torch.autocast(device_type=("cuda" if device=="cuda" else "cpu"), enabled=(device=="cuda")):
+                emb_a = model(anchor)
+                emb_p = model(positive)
+                emb_n = model(negative)
             loss = criterion(emb_a, emb_p, emb_n)
             optimizer.zero_grad(set_to_none=True)
             loss.backward()

train_vit_triplet.py

@@ -42,13 +42,15 @@ def embed_outfit(imgs: List[torch.Tensor], embedder: ResNetItemEmbedder, device:
 def main() -> None:
     args = parse_args()
     device = "cuda" if torch.cuda.is_available() else ("mps" if torch.backends.mps.is_available() else "cpu")
+    if device == "cuda":
+        torch.backends.cudnn.benchmark = True
 
     dataset = PolyvoreOutfitTripletDataset(args.data_root, split="train")
 
     def collate(batch):
         return batch  # variable length handled inside training loop
 
-    loader = DataLoader(dataset, batch_size=args.batch_size, shuffle=True, num_workers=4, pin_memory=True, collate_fn=collate)
+    loader = DataLoader(dataset, batch_size=args.batch_size, shuffle=True, num_workers=4, pin_memory=(device=="cuda"), collate_fn=collate)
 
     model = OutfitCompatibilityModel(embedding_dim=args.embedding_dim).to(device)
     embedder = ResNetItemEmbedder(embedding_dim=args.embedding_dim).to(device).eval()
@@ -80,10 +82,11 @@ def main() -> None:
             N = torch.cat(negative_tokens, dim=0)
 
             # get outfit-level embeddings via ViT encoder pooled output
-            ea = model.encoder(A).mean(dim=1)
-            ep = model.encoder(P).mean(dim=1)
-            en = model.encoder(N).mean(dim=1)
-            loss = triplet(ea, ep, en)
+            with torch.autocast(device_type=("cuda" if device=="cuda" else "cpu"), enabled=(device=="cuda")):
+                ea = model.encoder(A).mean(dim=1)
+                ep = model.encoder(P).mean(dim=1)
+                en = model.encoder(N).mean(dim=1)
+                loss = triplet(ea, ep, en)
             optimizer.zero_grad(set_to_none=True)
             loss.backward()
             optimizer.step()

utils/data_fetch.py

@@ -43,7 +43,35 @@ def ensure_dataset_ready() -> Optional[str]:
 
     # Download the HF dataset snapshot into root
     try:
-        snapshot_download("Stylique/Polyvore", repo_type="dataset", local_dir=root, local_dir_use_symlinks=False)
+        # Only fetch what's needed to run and prepare splits
+        allow = [
+            "images.zip",
+            "images/*.jpg",
+            "images/*.jpeg",
+            "images/*.png",
+            "train.json",
+            "valid.json",
+            "test.json",
+            "fill_in_blank_*.json",
+            "compatibility_*.txt",
+            "polyvore_item_metadata.json",
+            "polyvore_outfit_titles.json",
+            "categories.csv",
+        ]
+        ignore = [
+            "**/*hglmm*",
+            "disjoint/*",
+            "nondisjoint/*",
+            "*/large/*",
+        ]
+        snapshot_download(
+            "Stylique/Polyvore",
+            repo_type="dataset",
+            local_dir=root,
+            local_dir_use_symlinks=False,
+            allow_patterns=allow,
+            ignore_patterns=ignore,
+        )
     except Exception as e:  # pragma: no cover
         print(f"Failed to download Stylique/Polyvore dataset: {e}")
         return None