Add official VoxelMorph CT registration pipeline
This commit is contained in:
admin
400
model_and_train.py
Normal file
400
model_and_train.py
Normal file
@@ -0,0 +1,400 @@
|
||||
"""官方 VoxelMorph 训练适配器。
|
||||
|
||||
本项目以 https://github.com/voxelmorph/voxelmorph 的 PyTorch 核心为运行核心：
|
||||
- 网络: ``voxelmorph.nn.models.VxmPairwise``
|
||||
- 空间变换: 官方 ``voxelmorph.nn.modules.SpatialTransformer``
|
||||
- 平滑正则: 官方推荐的 ``neurite.nn.modules.SpatialGradient``
|
||||
- 相似度: 默认使用带 epsilon 的本地稳定 NCC；可切换到 ``neurite.nn.modules.NCC``
|
||||
|
||||
输入张量约定：
|
||||
- NIfTI 数组轴顺序为 (X, Y, Z)。
|
||||
- 官方 PyTorch VoxelMorph 输入为 (B, C, X, Y, Z)。
|
||||
- DDF 通道顺序与空间轴一致，即 (X, Y, Z)。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
from contextlib import nullcontext
|
||||
import json
|
||||
import os
|
||||
from pathlib import Path
|
||||
from typing import Dict, Iterable, List, Sequence, Tuple
|
||||
|
||||
import numpy as np
|
||||
|
||||
from config import DEFAULT_CHECKPOINT
|
||||
|
||||
OFFICIAL_VXM_REPO = "https://github.com/voxelmorph/voxelmorph"
|
||||
OFFICIAL_VXM_COMMIT = "db73f34b910bcefcb520f7f40a1bc4a3e0b6401d"
|
||||
|
||||
try: # 让缺 torch 的环境仍能 py_compile。
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
except Exception as _torch_exc: # pragma: no cover
|
||||
torch = None # type: ignore
|
||||
nn = None # type: ignore
|
||||
F = None # type: ignore
|
||||
_TORCH_IMPORT_ERROR = _torch_exc
|
||||
else:
|
||||
_TORCH_IMPORT_ERROR = None
|
||||
|
||||
|
||||
def configure_voxelmorph_backend() -> None:
|
||||
"""显式启用官方 VoxelMorph/Neurite 的 PyTorch 后端。"""
|
||||
|
||||
os.environ.setdefault("NEURITE_BACKEND", "pytorch")
|
||||
os.environ.setdefault("VXM_BACKEND", "pytorch")
|
||||
|
||||
|
||||
def require_torch():
|
||||
if torch is None:
|
||||
raise RuntimeError(
|
||||
"缺少 PyTorch，无法训练或推理。建议使用 conda 创建 CUDA 环境："
|
||||
"conda env create -f environment.yml"
|
||||
) from _TORCH_IMPORT_ERROR
|
||||
return torch
|
||||
|
||||
|
||||
def require_official_voxelmorph():
|
||||
"""导入官方 VoxelMorph 与 Neurite，并检查关键 API。"""
|
||||
|
||||
configure_voxelmorph_backend()
|
||||
require_torch()
|
||||
try:
|
||||
import neurite as ne # type: ignore
|
||||
import voxelmorph as vxm # type: ignore
|
||||
except Exception as exc: # pragma: no cover
|
||||
raise RuntimeError(
|
||||
"缺少官方 VoxelMorph/Neurite。请运行: conda env create -f environment.yml "
|
||||
"或 pip install -r requirements.txt"
|
||||
) from exc
|
||||
|
||||
if not hasattr(vxm.nn.models, "VxmPairwise"):
|
||||
raise RuntimeError("当前 voxelmorph 版本缺少 vxm.nn.models.VxmPairwise，请安装官方 dev 版本。")
|
||||
return vxm, ne
|
||||
|
||||
|
||||
def _require_nibabel():
|
||||
try:
|
||||
import nibabel as nib # type: ignore
|
||||
except Exception as exc: # pragma: no cover
|
||||
raise RuntimeError("缺少 nibabel，请先安装项目依赖。") from exc
|
||||
return nib
|
||||
|
||||
|
||||
if torch is not None:
|
||||
|
||||
class LocalNCCLoss(nn.Module):
|
||||
"""本地 NCC 兜底。
|
||||
|
||||
官方 VoxelMorph 当前把 NCC 放在 Neurite 中；若用户安装的 Neurite
|
||||
接口短期变动，本类保证训练脚本仍可运行，但网络核心仍是官方 VoxelMorph。
|
||||
"""
|
||||
|
||||
def __init__(self, window_size: int = 9, eps: float = 1e-5):
|
||||
super().__init__()
|
||||
self.window_size = int(window_size)
|
||||
self.eps = float(eps)
|
||||
|
||||
def forward(self, fixed, warped):
|
||||
win = [self.window_size] * 3
|
||||
padding = self.window_size // 2
|
||||
filt = torch.ones((1, 1, *win), dtype=fixed.dtype, device=fixed.device)
|
||||
win_volume = float(np.prod(win))
|
||||
|
||||
fixed_sum = F.conv3d(fixed, filt, padding=padding)
|
||||
warped_sum = F.conv3d(warped, filt, padding=padding)
|
||||
fixed2_sum = F.conv3d(fixed * fixed, filt, padding=padding)
|
||||
warped2_sum = F.conv3d(warped * warped, filt, padding=padding)
|
||||
cross_sum = F.conv3d(fixed * warped, filt, padding=padding)
|
||||
|
||||
fixed_mean = fixed_sum / win_volume
|
||||
warped_mean = warped_sum / win_volume
|
||||
cross = cross_sum - warped_mean * fixed_sum - fixed_mean * warped_sum + fixed_mean * warped_mean * win_volume
|
||||
fixed_var = fixed2_sum - 2 * fixed_mean * fixed_sum + fixed_mean * fixed_mean * win_volume
|
||||
warped_var = warped2_sum - 2 * warped_mean * warped_sum + warped_mean * warped_mean * win_volume
|
||||
|
||||
cc = (cross * cross) / (fixed_var * warped_var + self.eps)
|
||||
return -torch.mean(cc)
|
||||
|
||||
|
||||
class NegativeSimilarity(nn.Module):
|
||||
"""把 Neurite NCC 分数转换为可最小化的 loss。"""
|
||||
|
||||
def __init__(self, module: nn.Module):
|
||||
super().__init__()
|
||||
self.module = module
|
||||
|
||||
def forward(self, fixed, warped):
|
||||
return -torch.mean(self.module(fixed, warped))
|
||||
|
||||
|
||||
else:
|
||||
|
||||
class LocalNCCLoss: # type: ignore[no-redef]
|
||||
def __init__(self, *args, **kwargs):
|
||||
require_torch()
|
||||
|
||||
|
||||
class NegativeSimilarity: # type: ignore[no-redef]
|
||||
def __init__(self, *args, **kwargs):
|
||||
require_torch()
|
||||
|
||||
|
||||
def load_nifti_tensor(path: str | Path, device: str | "torch.device" = "cpu") -> Tuple["torch.Tensor", np.ndarray]:
|
||||
"""读取 NIfTI 并转换为官方 VoxelMorph 输入张量。
|
||||
|
||||
NIfTI data: (X, Y, Z)
|
||||
Torch tensor: (B, C, X, Y, Z)
|
||||
"""
|
||||
|
||||
require_torch()
|
||||
nib = _require_nibabel()
|
||||
img = nib.load(str(path), mmap=True)
|
||||
data = np.asanyarray(img.dataobj, dtype=np.float32)
|
||||
if data.ndim > 3:
|
||||
data = data[..., 0]
|
||||
data = np.nan_to_num(data.astype(np.float32, copy=False), copy=False)
|
||||
tensor = torch.from_numpy(data.copy())[None, None].to(device=device, dtype=torch.float32)
|
||||
return tensor, img.affine.copy()
|
||||
|
||||
|
||||
def resolve_device(device: str):
|
||||
require_torch()
|
||||
if device == "auto":
|
||||
return torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||||
return torch.device(device)
|
||||
|
||||
|
||||
def build_vxm_model(
|
||||
nb_features: Sequence[int],
|
||||
integration_steps: int,
|
||||
device,
|
||||
):
|
||||
"""构建官方 VoxelMorph Pairwise 模型。"""
|
||||
|
||||
vxm, _ = require_official_voxelmorph()
|
||||
model = vxm.nn.models.VxmPairwise(
|
||||
ndim=3,
|
||||
source_channels=1,
|
||||
target_channels=1,
|
||||
nb_features=list(int(v) for v in nb_features),
|
||||
integration_steps=int(integration_steps),
|
||||
device=str(device),
|
||||
)
|
||||
return model.to(device)
|
||||
|
||||
|
||||
def build_similarity_loss(loss_name: str, ncc_window: int, ncc_impl: str = "local"):
|
||||
"""构建相似度 loss。
|
||||
|
||||
颈部 CT 预处理后通常有大面积零填充背景，Neurite NCC 在低方差窗口中可能
|
||||
返回非有限值。因此默认使用带 epsilon 的本地 NCC；模型核心仍为官方
|
||||
VoxelMorph。需要对照官方依赖时，可传入 ``ncc_impl="neurite"``。
|
||||
"""
|
||||
|
||||
_, ne = require_official_voxelmorph()
|
||||
loss_name = loss_name.lower()
|
||||
if loss_name == "ncc":
|
||||
ncc_impl = ncc_impl.lower()
|
||||
if ncc_impl == "neurite":
|
||||
return NegativeSimilarity(ne.nn.modules.NCC(window_size=int(ncc_window)))
|
||||
if ncc_impl == "local":
|
||||
return LocalNCCLoss(window_size=int(ncc_window))
|
||||
raise ValueError("ncc_impl 只能是 local 或 neurite。")
|
||||
if loss_name == "mse":
|
||||
return ne.nn.modules.MSE()
|
||||
raise ValueError("loss_name 只能是 ncc 或 mse。")
|
||||
|
||||
|
||||
def build_gradient_loss():
|
||||
"""官方平滑项：Neurite SpatialGradient。"""
|
||||
|
||||
_, ne = require_official_voxelmorph()
|
||||
return ne.nn.modules.SpatialGradient(penalty="l2")
|
||||
|
||||
|
||||
def train_pair(
|
||||
moving_path: str | Path,
|
||||
fixed_path: str | Path,
|
||||
checkpoint_path: str | Path = DEFAULT_CHECKPOINT,
|
||||
epochs: int = 200,
|
||||
learning_rate: float = 1e-4,
|
||||
smooth_weight: float = 0.01,
|
||||
image_loss: str = "ncc",
|
||||
ncc_window: int = 9,
|
||||
ncc_impl: str = "local",
|
||||
nb_features: Sequence[int] = (16, 16, 16, 16, 16),
|
||||
integration_steps: int = 0,
|
||||
device: str = "auto",
|
||||
use_amp: bool = True,
|
||||
save_every: int = 50,
|
||||
) -> List[Dict[str, float]]:
|
||||
"""使用一对 moving/fixed 进行无监督 VoxelMorph 训练。"""
|
||||
|
||||
require_official_voxelmorph()
|
||||
device_obj = resolve_device(device)
|
||||
moving, _ = load_nifti_tensor(moving_path, device=device_obj)
|
||||
fixed, _ = load_nifti_tensor(fixed_path, device=device_obj)
|
||||
if moving.shape != fixed.shape:
|
||||
raise ValueError(f"Moving 与 Fixed 形状不一致: {tuple(moving.shape)} vs {tuple(fixed.shape)}")
|
||||
|
||||
model = build_vxm_model(nb_features=nb_features, integration_steps=integration_steps, device=device_obj)
|
||||
similarity_loss = build_similarity_loss(image_loss, ncc_window=ncc_window, ncc_impl=ncc_impl).to(device_obj)
|
||||
grad_loss_fn = build_gradient_loss().to(device_obj)
|
||||
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
|
||||
amp_enabled = bool(use_amp and device_obj.type == "cuda")
|
||||
scaler = torch.amp.GradScaler("cuda", enabled=amp_enabled)
|
||||
|
||||
checkpoint_path = Path(checkpoint_path)
|
||||
checkpoint_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
history: List[Dict[str, float]] = []
|
||||
|
||||
for epoch in range(1, epochs + 1):
|
||||
model.train()
|
||||
optimizer.zero_grad(set_to_none=True)
|
||||
|
||||
try:
|
||||
autocast_context = torch.amp.autocast("cuda", enabled=True) if amp_enabled else nullcontext()
|
||||
with autocast_context:
|
||||
displacement, warped = model(
|
||||
moving,
|
||||
fixed,
|
||||
return_warped_source=True,
|
||||
return_field_type="displacement",
|
||||
)
|
||||
|
||||
loss_sim = similarity_loss(fixed.float(), warped.float())
|
||||
loss_smooth = grad_loss_fn(displacement.float())
|
||||
loss = loss_sim + smooth_weight * loss_smooth
|
||||
|
||||
if not torch.isfinite(loss):
|
||||
raise RuntimeError(
|
||||
"训练 loss 出现非有限值。可尝试使用 --ncc-impl local、"
|
||||
"--image-loss mse，或检查预处理后是否有大面积常值区域。"
|
||||
)
|
||||
|
||||
scaler.scale(loss).backward()
|
||||
scaler.step(optimizer)
|
||||
scaler.update()
|
||||
except RuntimeError as exc:
|
||||
if "out of memory" in str(exc).lower():
|
||||
raise RuntimeError(
|
||||
"GPU/内存不足。可尝试减小 --target-shape，或降低 --nb-features，"
|
||||
"再重新预处理与训练。"
|
||||
) from exc
|
||||
raise
|
||||
|
||||
row = {
|
||||
"epoch": float(epoch),
|
||||
"loss": float(loss.detach().cpu()),
|
||||
"image_loss": float(loss_sim.detach().cpu()),
|
||||
"smooth_loss": float(loss_smooth.detach().cpu()),
|
||||
}
|
||||
history.append(row)
|
||||
print(
|
||||
f"Epoch {epoch:04d}/{epochs} | "
|
||||
f"loss={row['loss']:.6f} image={row['image_loss']:.6f} smooth={row['smooth_loss']:.6f}"
|
||||
)
|
||||
|
||||
if epoch == epochs or (save_every > 0 and epoch % save_every == 0):
|
||||
save_checkpoint(
|
||||
checkpoint_path,
|
||||
model=model,
|
||||
epoch=epoch,
|
||||
input_shape_xyz=tuple(int(v) for v in moving.shape[2:]),
|
||||
nb_features=nb_features,
|
||||
integration_steps=integration_steps,
|
||||
config={
|
||||
"moving_path": str(moving_path),
|
||||
"fixed_path": str(fixed_path),
|
||||
"learning_rate": learning_rate,
|
||||
"smooth_weight": smooth_weight,
|
||||
"image_loss": image_loss,
|
||||
"ncc_window": ncc_window,
|
||||
"ncc_impl": ncc_impl,
|
||||
},
|
||||
history=history,
|
||||
)
|
||||
|
||||
history_path = checkpoint_path.with_suffix(".history.json")
|
||||
history_path.write_text(json.dumps(history, ensure_ascii=False, indent=2), encoding="utf-8")
|
||||
return history
|
||||
|
||||
|
||||
def save_checkpoint(
|
||||
path: str | Path,
|
||||
model,
|
||||
epoch: int,
|
||||
input_shape_xyz: Sequence[int],
|
||||
nb_features: Sequence[int],
|
||||
integration_steps: int,
|
||||
config: Dict,
|
||||
history: List[Dict[str, float]],
|
||||
) -> None:
|
||||
require_torch()
|
||||
path = Path(path)
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
torch.save(
|
||||
{
|
||||
"model_state_dict": model.state_dict(),
|
||||
"epoch": int(epoch),
|
||||
"input_shape_xyz": tuple(int(v) for v in input_shape_xyz),
|
||||
"nb_features": list(int(v) for v in nb_features),
|
||||
"integration_steps": int(integration_steps),
|
||||
"official_core": {
|
||||
"library": "voxelmorph",
|
||||
"repo": OFFICIAL_VXM_REPO,
|
||||
"commit": OFFICIAL_VXM_COMMIT,
|
||||
"model_class": "voxelmorph.nn.models.VxmPairwise",
|
||||
},
|
||||
"config": config,
|
||||
"history": history,
|
||||
},
|
||||
path,
|
||||
)
|
||||
|
||||
|
||||
def build_arg_parser() -> argparse.ArgumentParser:
|
||||
parser = argparse.ArgumentParser(description="用官方 VoxelMorph 训练 3D 颈部 CT 形变配准模型。")
|
||||
parser.add_argument("--moving", required=True, help="预处理后的 moving NIfTI。")
|
||||
parser.add_argument("--fixed", required=True, help="预处理后的 fixed NIfTI。")
|
||||
parser.add_argument("--checkpoint", default=str(DEFAULT_CHECKPOINT), help="模型权重保存路径。")
|
||||
parser.add_argument("--epochs", type=int, default=200)
|
||||
parser.add_argument("--lr", type=float, default=1e-4)
|
||||
parser.add_argument("--smooth-weight", type=float, default=0.01)
|
||||
parser.add_argument("--image-loss", choices=["ncc", "mse"], default="ncc")
|
||||
parser.add_argument("--ncc-window", type=int, default=9)
|
||||
parser.add_argument("--ncc-impl", choices=["local", "neurite"], default="local")
|
||||
parser.add_argument("--nb-features", type=int, nargs="+", default=[16, 16, 16, 16, 16])
|
||||
parser.add_argument("--integration-steps", type=int, default=0, help="0 为普通 dense flow；>0 使用 scaling-and-squaring。")
|
||||
parser.add_argument("--device", default="auto")
|
||||
parser.add_argument("--no-amp", action="store_true", help="关闭 CUDA AMP 混合精度。")
|
||||
parser.add_argument("--save-every", type=int, default=50)
|
||||
return parser
|
||||
|
||||
|
||||
def main(argv: Iterable[str] | None = None) -> None:
|
||||
args = build_arg_parser().parse_args(argv)
|
||||
train_pair(
|
||||
moving_path=args.moving,
|
||||
fixed_path=args.fixed,
|
||||
checkpoint_path=args.checkpoint,
|
||||
epochs=args.epochs,
|
||||
learning_rate=args.lr,
|
||||
smooth_weight=args.smooth_weight,
|
||||
image_loss=args.image_loss,
|
||||
ncc_window=args.ncc_window,
|
||||
ncc_impl=args.ncc_impl,
|
||||
nb_features=args.nb_features,
|
||||
integration_steps=args.integration_steps,
|
||||
device=args.device,
|
||||
use_amp=not args.no_amp,
|
||||
save_every=args.save_every,
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user