Anomalib PatchCore 预训练权重本地化实践：避免重复下载的完整解决方案

1次阅读

没有评论

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

在工业场景中使用 Anomalib PatchCore 进行异常检测时，每次运行训练脚本都会触发预训练权重自动下载机制，这带来三个明显问题：

网络依赖风险：生产环境往往限制外网访问，导致脚本因下载失败而中止
重复存储浪费：相同版本的权重文件被多次下载到不同项目的缓存目录
时间成本激增：实测显示，ResNet50 backbone 的权重下载平均耗时 47 秒（10 次训练任务累计浪费近 8 分钟）

不同系统的默认缓存位置如下：

Linux/MacOS: ~/.cache/torch/hub/checkpoints/
Windows: C:\Users\<username>\.cache\torch\hub\checkpoints\

通过 Python 代码可动态获取路径：

from torch.hub import get_dir
cache_path = get_dir() + "/checkpoints/"

model:
  backbone: resnet50
  pretrained_weights: /mnt/shared/weights/resnet50-0676ba61.pth  # 绝对路径

export ANOMALIB_WEIGHTS_DIR=/opt/pretrained_weights

对应读取逻辑：

import os
weight_path = os.getenv(
    "ANOMALIB_WEIGHTS_DIR", 
    default="~/.cache/torch/hub/checkpoints"
)

import hashlib
def verify_weights(filepath, expected_md5):
    with open(filepath, "rb") as f:
        md5 = hashlib.md5(f.read()).hexdigest()
    if md5 != expected_md5:
        raise ValueError("Weights file corrupted")
    return True

# ResNet50 官方权重 MD5
verify_weights("resnet50.pth", "0676ba61b34242c5dfea4603d9f050d7")

# configs/model/patchcore.yaml
model:
  name: patchcore
  backbone: wide_resnet50_2
  pretrained_weights: /assets/pretrained/wide_resnet50_2.pth
  layers: ["layer2", "layer3"]

metrics:
  image: ["AUROC"]
  pixel: ["AUROC"]

try:
    model = PatchCore.load_from_checkpoint(config["pretrained_weights"])
except FileNotFoundError:
    logger.warning("Local weights not found, falling back to download")
    model = PatchCore(backbone=config["backbone"])

方案	优点	缺点
NFS	零配置接入	单点故障风险
S3	高可用	需要 AWS 认证
MinIO	自托管兼容 S3	需额外部署

# .gitlab-ci.yml
stages:
  - setup

cache_weights:
  stage: setup
  script:
    - mkdir -p /shared/weights
    - if [! -f "/shared/weights/resnet50.pth"]; then
        wget -P /shared/weights https://download.pytorch.org/models/resnet50-0676ba61.pth
      fi
  artifacts:
    paths:
      - /shared/weights

# 按需加载权重
from functools import partial

class LazyWeightsLoader:
    def __init__(self, path):
        self.path = path
        self._weights = None

    @property 
    def weights(self):
        if self._weights is None:
            self._weights = torch.load(self.path, map_location="cpu")
        return self._weights

import time

def test_speed():
    # 原始方式
    start = time.time()
    model = PatchCore(backbone="resnet50")  # 自动下载
    print(f"Download mode: {time.time() - start:.2f}s")

    # 本地模式
    start = time.time()
    model = PatchCore.load_from_checkpoint("./weights/resnet50.pth")
    print(f"Local mode: {time.time() - start:.2f}s")

典型输出结果：

Download mode: 48.31s
Local mode: 1.07s

import pytest
from unittest.mock import patch

@patch("torch.hub.download_url_to_file")
def test_offline(mock_download):
    with pytest.raises(FileNotFoundError):
        # 模拟断网环境
        mock_download.side_effect = RuntimeError("No internet")
        model = PatchCore(backbone="resnet50")

    # 应能正常加载本地权重
    assert PatchCore.load_from_checkpoint("./local_weights.pth")

问题现象	解决方案
Permission denied	执行`chmod +r /path/to/weights`
Windows 路径反斜杠转义	使用 `r"C:\path\to\file"` 原始字符串
MD5 校验失败	重新下载并检查磁盘空间
配置文件未生效	确认环境变量优先级高于配置文件