Fix KV260 PatchCore XModel compilation

.gitignore

@@ -366,3 +366,5 @@ apps/api/fastapi_app_np.py
 *.pyc
 production_package/*
 tests/__pycache__/*
+compiled_patchcore_kv260/*
+quantize_result/*

anomavision/quantize/model/backends/xmodel/patchcore.py

@@ -0,0 +1,83 @@
+"""Single-DPU PatchCore graph for AMD/Xilinx KV260.
+
+The normal AnomaVision PatchCore implementation is unchanged. This backend
+keeps the PatchCore similarity tensor in NCHW and uses a DPU-friendly channel
+maximum instead of aten::amax, which caused Vitis AI to insert a CPU transpose
+before the reduction.
+"""
+
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class PatchCoreKV260(nn.Module):
+    """DPU-friendly whole PatchCore graph for KV260."""
+
+    def __init__(self, model: nn.Module) -> None:
+        super().__init__()
+
+        self.backbone = model.embeddings_extractor.backbone
+
+        memory_bank = F.normalize(
+            model.memory_bank.float(),
+            dim=-1,
+        )
+
+        # Keep memory bank in [819, 64, 1, 1] so Vitis AI can treat
+        # each memory vector as a 1x1 convolution filter.
+        self.register_buffer(
+            "memory_bank",
+            memory_bank.reshape(819, 64, 1, 1),
+        )
+
+    def forward(self, x: torch.Tensor):
+        # ResNet layer1. Keep the feature map in native NCHW layout.
+        x = self.backbone.conv1(x)
+        x = self.backbone.bn1(x)
+        x = self.backbone.relu(x)
+        x = self.backbone.maxpool(x)
+        features = self.backbone.layer1(x)
+
+        # Compute similarity against all 819 normalized memory vectors.
+        # [B, 64, 56, 56] -> [B, 819, 56, 56]
+        similarity = F.conv2d(
+            features,
+            self.memory_bank,
+        )
+
+        # Reduce the memory-bank/channel dimension with torch.max rather than
+        # torch.amax. Vitis AI maps the explicit max reduction more reliably
+        # for the KV260 DPU and avoids the aten::amax CPU transpose path.
+        max_similarity = torch.max(
+            similarity,
+            dim=1,
+            keepdim=True,
+        ).values
+
+        # DPU-friendly squared cosine distance:
+        #   2 - 2 * max(cosine_similarity)
+        # Avoid sqrt/relu/clamp because they introduce unsupported/CPU ops.
+        distance = torch.add(
+            max_similarity * -2.0,
+            2.0,
+        )
+
+        # 56x56 -> 224x224 anomaly map.
+        score_map = F.interpolate(
+            distance,
+            size=(224, 224),
+            mode="bilinear",
+            align_corners=False,
+        ).squeeze(1)
+
+        # Image-level anomaly score.
+        image_score = F.max_pool2d(
+            distance,
+            kernel_size=(56, 56),
+            stride=1,
+        ).flatten(1)
+
+        return image_score, score_map

config.yml

@@ -1,7 +1,7 @@
 # =========================
 # Dataset / preprocessing (shared by train, detect, eval, stream)
 # =========================
-dataset_path: "D:/01-DATA"               # Root dataset folder (MVTec-style: contains train/test subfolders)
+dataset_path: "/workspace/dataset"               # Root dataset folder (MVTec-style: contains train/test subfolders)
 class_name: "bottle"                     # Class name for MVTec dataset
 resize: [224, 224]                       # Resize dimensions before processing [width, height]
 crop_size:                               # Final crop size [width, height]

howto.md

@@ -0,0 +1,150 @@
+# PatchCore → KV260 XModel
+
+Simple beginner guide for quantizing **AnomaVision PatchCore** with **Vitis AI 3.5** and generating an **XModel for KV260**.
+
+## 1. Activate Vitis AI
+
+```bash
+conda activate vitis-ai-pytorch
+```
+
+Check:
+
+```bash
+vai_c_xir -h
+```
+
+---
+
+## 2. Go to AnomaVision
+
+```bash
+cd /workspace/AnomaVision
+```
+
+---
+
+## 3. Check the model
+
+Our PatchCore model:
+
+```text
+distributions/patchcore/bottle/anomav_exp/model.pt
+```
+
+Calibration images:
+
+```text
+/workspace/dataset/bottle/train/good
+```
+
+---
+
+## 4. Create calibration data
+
+PatchCore needs **normal/good images** for calibration.
+
+Example:
+
+```bash
+ls /workspace/dataset/bottle/train/good | head
+```
+
+You should see images such as:
+
+```text
+000.png
+001.png
+002.png
+...
+```
+
+---
+
+## 5. Run INT8 calibration
+
+Run:
+
+```bash
+python quantize_patchcore_kv260.py \
+  --model distributions/patchcore/bottle/anomav_exp/model.pt \
+  --calibration-dir /workspace/dataset/bottle/train/good \
+  --output-dir compiled_patchcore_kv260 \
+  --quant_mode calib
+```
+
+Successful calibration should finish with:
+
+```text
+Calibration finished.
+Quant config exported.
+```
+
+---
+
+## 6. Generate the XModel
+
+Run:
+
+```bash
+python quantize_patchcore_kv260.py \
+  --model distributions/patchcore/bottle/anomav_exp/model.pt \
+  --calibration-dir /workspace/dataset/bottle/train/good \
+  --output-dir compiled_patchcore_kv260 \
+  --quant_mode test
+```
+
+The important output is:
+
+```text
+compiled_patchcore_kv260/PatchCoreKV260_int.xmodel
+```
+
+---
+
+## 7. Check the XModel
+
+```bash
+ls -lh compiled_patchcore_kv260/*.xmodel
+```
+
+Then:
+
+```bash
+python -c "import xir; g=xir.Graph.deserialize('compiled_patchcore_kv260/PatchCoreKV260_int.xmodel'); print('XModel OK:', g.get_name()); print('Ops:', len(g.get_ops()))"
+```
+
+Expected:
+
+```text
+XModel OK: PatchCoreKV260
+Ops: 98
+```
+
+---
+
+## 8. Compile for KV260
+
+Use the KV260 architecture file:
+
+```bash
+vai_c_xir \
+  -x compiled_patchcore_kv260/PatchCoreKV260_int.xmodel \
+  -a /opt/vitis_ai/compiler/arch/DPUCZDX8G/KV260/arch.json \
+  -o compiled_patchcore_kv260/compiled
+```
+
+---
+
+## Final result
+
+You want:
+
+```text
+compiled_patchcore_kv260/
+├── PatchCoreKV260_int.xmodel
+└── compiled/
+    └── PatchCoreKV260_int.xmodel
+```
+
+The final XModel is intended for the **AMD/Xilinx KV260 DPU**.

quantize_patchcore_kv260.py

@@ -0,0 +1,89 @@
+import argparse
+from pathlib import Path
+
+# Vitis AI 3.5.0 deploy-optimizer workaround.
+import nndct_shared.compile.deploy_optimizer as _deploy_optimizer
+import torch
+from PIL import Image
+from pytorch_nndct.apis import torch_quantizer
+
+from anomavision.quantize.model.backends.xmodel.patchcore import PatchCoreKV260
+
+for _name, _obj in vars(_deploy_optimizer).items():
+    if isinstance(_obj, type) and hasattr(_obj, "fuse_transpose_matmul"):
+        _obj.fuse_transpose_matmul = lambda self: None
+
+        print(f"[KV260] Disabled {_name}.fuse_transpose_matmul")
+
+
+def load_calibration_images(directory, size=224, limit=50):
+    paths = []
+    directory = Path(directory)
+    for ext in ("*.jpg", "*.jpeg", "*.png", "*.bmp", "*.webp"):
+        paths.extend(directory.glob(ext))
+    paths = sorted(paths)[:limit]
+    if not paths:
+        raise RuntimeError(f"No calibration images found in {directory}")
+
+    tensors = []
+    for path in paths:
+        with Image.open(path) as image:
+            image = image.convert("RGB").resize((size, size))
+            tensor = torch.from_numpy(__import__("numpy").array(image))
+            tensors.append(tensor.permute(2, 0, 1).float() / 255.0)
+    return tensors
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model", required=True)
+    parser.add_argument("--calibration-dir", required=True)
+    parser.add_argument("--output-dir", default="./quantize_result")
+    parser.add_argument("--quant_mode", choices=["calib", "test"], default="calib")
+    args = parser.parse_args()
+
+    print("Loading:", args.model)
+    model = torch.load(args.model, map_location="cpu", weights_only=False)
+    model.eval()
+    print("Original model:", type(model))
+    print("Memory bank:", tuple(model.memory_bank.shape))
+
+    wrapper = PatchCoreKV260(model)
+    wrapper.eval()
+
+    dummy = torch.randn(1, 3, 224, 224)
+    print("Testing wrapper...")
+    with torch.no_grad():
+        outputs = wrapper(dummy)
+    print("Output 0:", tuple(outputs[0].shape))
+    print("Output 1:", tuple(outputs[1].shape))
+
+    calibration_images = load_calibration_images(args.calibration_dir)
+    print("Calibration images:", len(calibration_images))
+
+    quantizer = torch_quantizer(args.quant_mode, wrapper, (dummy,))
+    quant_model = quantizer.quant_model
+    quant_model.eval()
+
+    print("Running calibration...")
+    with torch.no_grad():
+        for i, image in enumerate(calibration_images):
+            quant_model(image.unsqueeze(0))
+            if (i + 1) % 10 == 0:
+                print(f"  {i + 1}/{len(calibration_images)}")
+
+    output_dir = Path(args.output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    if args.quant_mode == "calib":
+        quantizer.export_quant_config()
+        print("\nCalibration finished.")
+        print("Quant config exported.")
+    else:
+        quantizer.export_xmodel(output_dir=str(output_dir), deploy_check=False)
+        print("\nXMODEL exported to:")
+        print(output_dir)
+
+
+if __name__ == "__main__":
+    main()