如何高效下载与部署anomalib预训练模型：避坑指南与最佳实践

1次阅读

共计 2746 个字符，预计需要花费 7 分钟才能阅读完成。

在工业缺陷检测等场景中，anomalib 提供的预训练模型（Pretrained Model）能大幅减少开发时间。但在实际下载和部署过程中，开发者常遇到以下问题：

国外源下载速度慢：直接从官方源下载模型权重时，国内用户常遇到几 KB/ s 的龟速
版本兼容性问题：PyTorch 与 CUDA 版本不匹配导致ImportError，尤其是新旧版本 anomalib 对 PyTorch 的要求差异较大
模型校验失败：网络波动导致下载文件损坏，但缺乏自动重试机制

通过替换 pip 和 conda 源为国内镜像，可提升依赖安装速度。以下是清华大学源的配置方法：

# 永久配置 pip 镜像源
pip config set global.index-url https://pypi.tuna.tsinghua.edu.cn/simple

# conda 换源（Linux/macOS）conda config --add channels https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/main/
conda config --add channels https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/free/
conda config --set show_channel_urls yes

推荐使用 conda 创建独立环境，并锁定关键库版本：

# 创建指定 Python 版本的环境
conda create -n anomalib_env python=3.8
conda activate anomalib_env

# 安装匹配的 PyTorch（以 CUDA 11.3 为例）conda install pytorch==1.12.1 torchvision==0.13.1 torchaudio==0.12.1 cudatoolkit=11.3 -c pytorch

以下是带自动重试和校验的模型加载代码：

import hashlib
import os
from anomalib.models import Padim
from pathlib import Path
import requests

MODEL_URL = "https://github.com/openvinotoolkit/anomalib/releases/download/model/padim_model.pth"
EXPECTED_MD5 = "a1b2c3d4e5f6..."  # 替换为实际 MD5 值

# 带校验的下载函数
def download_with_retry(url, save_path, max_retry=3):
    for attempt in range(max_retry):
        try:
            response = requests.get(url, stream=True)
            response.raise_for_status()

            # 计算下载文件的 MD5
            md5 = hashlib.md5()
            with open(save_path, "wb") as f:
                for chunk in response.iter_content(chunk_size=8192):
                    f.write(chunk)
                    md5.update(chunk)

            if md5.hexdigest() == EXPECTED_MD5:
                return True
            os.remove(save_path)  # 校验失败删除文件
        except Exception as e:
            print(f"Attempt {attempt + 1} failed: {str(e)}")
    return False

# 模型加载
def load_padim_model():
    model_path = Path("./models/padim_model.pth")
    if not model_path.exists():
        model_path.parent.mkdir(exist_ok=True)
        if not download_with_retry(MODEL_URL, model_path):
            raise RuntimeError("Failed to download model after retries")

    model = Padim.load_from_checkpoint(model_path)
    return model

anomalib 版本	PyTorch 版本	CUDA 版本
0.3.x	1.11.0	11.3
0.4.x	1.12.1	11.6

对于大模型可采用分块加载：

from torch import nn

class ChunkedModel(nn.Module):
    def __init__(self, model_path):
        super().__init__()
        self.model_path = model_path
        self.model = None

    def load_chunk(self):
        if self.model is None:
            self.model = Padim.load_from_checkpoint(self.model_path)

import cv2
import numpy as np
from torchvision.transforms import ToTensor

model = load_padim_model()
model.eval()

# 测试图像预处理
test_img = cv2.imread("test.jpg")
img_tensor = ToTensor()(test_img).unsqueeze(0)

# 推理
with torch.no_grad():
    prediction = model(img_tensor)
    anomaly_map = prediction["anomaly_map"].squeeze().numpy()

# 可视化结果
heatmap = (anomaly_map * 255).astype(np.uint8)
heatmap = cv2.applyColorMap(heatmap, cv2.COLORMAP_JET)
blended = cv2.addWeighted(test_img, 0.7, heatmap, 0.3, 0)