使用MONAI transforms中SimpleITK时Python无法调用GPU的问题

核心问题原因 #

1. 自定义的N4ITKTransform继承自基础Transform类，而非MONAI专为字典格式数据设计的MapTransform，导致与MONAI的字典数据处理流程不兼容。
2. SimpleITK的处理完全在CPU上进行，生成的numpy数组转成张量后默认留在CPU，若未显式转移到GPU，会导致GPU上的模型无法调用GPU计算。

具体修复步骤 #

1. 修正自定义Transform的继承类 #

将N4ITKTransform改为继承MapTransform，确保符合MONAI字典数据的处理规范：

import numpy as np
import SimpleITK as sitk
from monai.transforms import MapTransform  # 替换原Transform导入

class N4ITKTransform(MapTransform):
    def __init__(self, keys):
        super().__init__(keys)
    
    def __call__(self, data):
        # 复制输入字典避免修改原数据
        processed_data = dict(data)
        # 仅对指定key（这里是image）做N4偏置校正
        for key in self.keys:
            if key != "image":
                continue
            filtered_channels = []
            for channel in processed_data[key]:
                # numpy转SimpleITK图像
                sitk_img = sitk.GetImageFromArray(channel)
                sitk_img = sitk.Cast(sitk_img, sitk.sitkFloat32)
                # 执行N4校正
                corrector = sitk.N4BiasFieldCorrectionImageFilter()
                corrected_sitk_img = corrector.Execute(sitk_img)
                # SimpleITK图像转回numpy
                filtered_channels.append(sitk.GetArrayFromImage(corrected_sitk_img))
            processed_data[key] = np.stack(filtered_channels)
        return processed_data

使用时需指定作用的keys：

train_transforms = Compose(monai_load + [N4ITKTransform(keys=["image"])] + monai_transforms)

2. 确保张量正确转移到GPU #

有两种方式实现：

方式一：在训练循环中显式转移

import torch

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)

# 训练循环
for batch in train_loader:
    # 将batch中所有张量转移到GPU
    images = batch["image"].to(device, non_blocking=True)
    segmentations = batch["segmentation"].to(device, non_blocking=True)
    # 执行模型前向传播等操作
    outputs = model(images)
    # ... 后续损失计算、反向传播步骤

方式二：在变换流程中直接转GPU

修改变换链，用EnsureTyped替代ToTensord并指定设备，让预处理后直接生成GPU张量：

# 调整monai_load，移除EnsureTyped(data_type="numpy")
monai_load = [
    LoadImaged(keys=["image","segmentation"], image_only=False, reader=PILReader()),
    AddChanneld(keys=["segmentation","image"]),   
    RepeatChanneld(keys=["image"], repeats=3),
    AsChannelFirstd(keys=["image"], channel_dim = 0), 
]

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# 调整monai_transforms
monai_transforms =[
    AsDiscreted(keys=["segmentation"], threshold=0.5),
    # EnsureTyped自动将numpy转成张量并转移到指定设备
    EnsureTyped(keys=["image", "segmentation"], device=device),
]

train_transforms = Compose(monai_load + [N4ITKTransform(keys=["image"])] + monai_transforms)

3. 额外检查项 #

• 确认PILReader已正确导入（需补充from monai.data import PILReader）。
• 保持num_workers=0，多进程数据加载时SimpleITK可能出现兼容性问题。

内容的提问来源于stack exchange，提问作者Federico Digiacomo