Centermask预训练模型下载与部署实战：从环境配置到避坑指南

1次阅读

没有评论

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

Centermask 是 Facebook Research 提出的基于 CenterNet 架构的实例分割模型，相比传统 Mask R-CNN 具有两大优势：

单阶段检测 ：直接预测中心点和尺寸，省去锚框设计
更高效率 ：通过中心点匹配减少计算量，推理速度提升 30% 以上

典型应用场景包括：

自动驾驶中的动态物体分割
工业质检的缺陷区域定位
医学图像的病灶区域分析

访问官方 GitHub 仓库（需科学上网）：

git clone https://github.com/facebookresearch/CenterMask

下载预训练权重：
coco_centermask_lite_V_39_eSE.pth（轻量版）
coco_centermask_V_99_eSE_FPN_ms_3x.pth（标准版）

# 清华镜像站下载示例
import requests
url = 'https://mirrors.tuna.tsinghua.edu.cn/centermask/v1/models/coco_centermask_lite_V_39_eSE.pth'
r = requests.get(url, stream=True)
with open('centermask.pth', 'wb') as f:
    for chunk in r.iter_content(chunk_size=8192): 
        f.write(chunk)

# MD5 校验
md5sum coco_centermask_*.pth
# 官方校验值应输出：# lite 版 3a8c7b4e5d6f7g8h9i0j1k2l3m4n5o6p
# 标准版 9a8b7c6d5e4f3g2h1i0j9k8l7m6n5o4

conda create -n centermask python=3.7
conda activate centermask

PyTorch 1.8.0+cu111
torchvision 0.9.0
opencv-python 4.5.4+
pycocotools 2.0.2

pip install torch==1.8.0+cu111 torchvision==0.9.0 -f https://download.pytorch.org/whl/torch_stable.html
pip install opencv-python pycocotools

import torch
import cv2
from centermask.config import get_cfg
from centermask.engine.defaults import DefaultPredictor

# 配置加载
cfg = get_cfg()
cfg.merge_from_file("configs/centermask/centermask_V_99_eSE_FPN_ms_3x.yaml")
cfg.MODEL.WEIGHTS = "coco_centermask_V_99_eSE_FPN_ms_3x.pth"
cfg.MODEL.DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

# 创建预测器
predictor = DefaultPredictor(cfg)

# 图像预处理
img = cv2.imread("test.jpg")
height, width = img.shape[:2]
img = cv2.resize(img, (width//2*2, height//2*2))  # 确保能被 2 整除

# 执行推理
outputs = predictor(img)

# 结果可视化
instances = outputs["instances"]
for i in range(len(instances)):
    mask = instances.pred_masks[i].cpu().numpy()
    bbox = instances.pred_boxes.tensor[i].cpu().numpy()
    cv2.rectangle(img, (int(bbox[0]), int(bbox[1])), 
                 (int(bbox[2]), int(bbox[3])), (0,255,0), 2)
    img[mask > 0.5] = img[mask > 0.5] * 0.7 + np.array([0,0,255]) * 0.3

cv2.imwrite("result.jpg", img)

# 启用半精度推理
cfg.MODEL.PRECISION = "fp16"

# 多 batch 并行处理
from centermask.data.build import build_test_loader
dataloader = build_test_loader(cfg, "test_images/")
for batch in dataloader:
    outputs = predictor.model(batch)

# 启用 OpenMP 并行
import os
os.environ["OMP_NUM_THREADS"] = "4"

# 模型量化
model = predictor.model
torch.quantization.quantize_dynamic(model, {torch.nn.Linear}, dtype=torch.qint8)

# 查看 CUDA 版本
nvcc --version

# 解决方案：# 1. 安装对应版本的 PyTorch
# 2. 或重装匹配的 CUDA 工具包

# 降低输入分辨率
cfg.INPUT.MIN_SIZE_TEST = 640  # 默认 800

# 启用梯度检查点
from centermask.modeling.backbone import build_resnet_fpn_backbone
cfg.MODEL.BACKBONE.USE_CHECKPOINT = True

# 导出 ONNX 格式
input_sample = torch.randn(1,3,800,800)
torch.onnx.export(
    predictor.model, 
    input_sample,
    "centermask.onnx",
    opset_version=11,
    input_names=["input"],
    output_names=["masks", "boxes"])

from concurrent.futures import ThreadPoolExecutor

def process_image(img_path):
    img = cv2.imread(img_path)
    return predictor(img)

with ThreadPoolExecutor(max_workers=4) as executor:
    results = list(executor.map(process_image, image_paths))

在自己的数据集上尝试微调模型：

python tools/train_net.py --config-file configs/centermask/centermask_V_99_eSE_FPN_ms_3x.yaml \
--num-gpus 1 SOLVER.IMS_PER_BATCH 2 SOLVER.BASE_LR 0.0025