401 lines
14 KiB
Python
401 lines
14 KiB
Python
"""官方 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()
|